import gradio as gr
import torch
from PIL import Image
from model import CRM
from inference import generate3d
import numpy as np

# Load model
crm_path = "CRM.pth"  # Make sure the model is uploaded to the Space
model = CRM(torch.load(crm_path, map_location="cpu"))
model = model.to("cuda:0" if torch.cuda.is_available() else "cpu")

def generate_3d(image_path, seed=1234, scale=5.5, step=30):
    image = Image.open(image_path).convert("RGB")
    np_img = np.array(image)
    glb_path = generate3d(model, np_img, np_img, "cuda:0" if torch.cuda.is_available() else "cpu")
    return glb_path

iface = gr.Interface(
    fn=generate_3d,
    inputs=gr.Image(type="filepath"),
    outputs=gr.Model3D(),
    title="Convolutional Reconstruction Model (CRM)",
    description="Upload an image to generate a 3D model."
)

iface.launch()