import gradio as gr
from transformers import AutoModel, AutoProcessor
from PIL import Image
import torch
import requests
from io import BytesIO
import trimesh
import plotly.graph_objects as go



# Load model and processor from Hugging Face
def load_model_and_processor():
    try:
        model = AutoModel.from_pretrained("zxhezexin/openlrm-mix-large-1.1")
        processor = AutoProcessor.from_pretrained("zxhezexin/openlrm-mix-large-1.1")
        return model, processor
    except Exception as e:
        print(f"Error loading model or processor: {e}")
        return None, None

model, processor = load_model_and_processor()

# Example image URL (replace this with a suitable example)
example_image_url = "https://huggingface.co/datasets/nateraw/image-folder/resolve/main/example_1.png"

# Function to load example image from URL
def load_example_image():
    try:
        response = requests.get(example_image_url)
        image = Image.open(BytesIO(response.content))
        return image
    except Exception as e:
        print(f"Error loading example image: {e}")
        return None

# Define function to generate 3D output from 2D image
def image_to_3d(image):
    if processor is None or model is None:
        return "Model or processor not loaded."
    
    try:
        # Preprocess input image
        inputs = processor(images=image, return_tensors="pt")
        
        # Run inference
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Convert outputs to a 3D mesh (replace with actual logic based on model output)
        # Assuming 'vertices' and 'faces' are returned by the model (adjust as needed)
        vertices = outputs['vertices'].numpy()  # Placeholder for vertex output
        faces = outputs['faces'].numpy()  # Placeholder for face output
        
        # Create a mesh using trimesh
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        
        # Visualize the mesh using Plotly
        fig = go.Figure(data=[go.Mesh3d(x=vertices[:,0], y=vertices[:,1], z=vertices[:,2], 
                                         i=faces[:,0], j=faces[:,1], k=faces[:,2])])
        
        return fig  # return the figure for display
    except Exception as e:
        return f"Error during inference: {str(e)}"

# Load the example image for the Gradio interface
example_image = load_example_image()

# Gradio interface setup
interface = gr.Interface(
    fn=image_to_3d,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Plot(label="3D Model"),
    title="OpenLRM Mix-Large 1.1 - Image to 3D",
    description="Upload an image to generate a 3D model using OpenLRM Mix-Large 1.1.",
    examples=[[example_image]] if example_image else None,
    theme="compact"
)

# Display a suggestion below the upload widget
interface.launch()