import gradio as gr
import torch
from PIL import Image
import numpy as np
import os
from pathlib import Path
import tempfile
from huggingface_hub import hf_hub_download

# Set environment variable to use CPU
os.environ["SF3D_USE_CPU"] = "1"

# Import the main pipeline
from stable_fast_3d import StableFast3D

# Initialize the model
model = StableFast3D()

def process_image(image, prompt):
    # Convert image to PIL if it's not already
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Create temporary directory for output
    with tempfile.TemporaryDirectory() as tmpdir:
        output_path = Path(tmpdir) / "output.glb"
        
        # Process the image
        model.process_image(
            image=image,
            prompt=prompt,
            output_path=str(output_path)
        )
        
        # Return the GLB file
        return str(output_path)

def convert_2d_to_3d(image, prompt=None):
    """
    Convert a 2D image to a 3D model using Stable Diffusion and advanced 3D reconstruction.
    
    Args:
        image (PIL.Image): Input 2D image
        prompt (str, optional): Text prompt to guide the 3D generation
        
    Returns:
        str: Path to the generated GLB file
    """
    # TODO: Implement the actual 2D to 3D conversion logic
    # For now, return a placeholder message
    return "3D conversion will be implemented soon!"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Stable Fast 3D - Convert 2D Images to 3D Models")
    gr.Markdown("Upload a 2D image and get a 3D model in return. Optionally provide a text prompt to guide the generation.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            text_prompt = gr.Textbox(label="Text Prompt (optional)", placeholder="Enter a description to guide the 3D generation...")
            convert_btn = gr.Button("Convert to 3D")
        
        with gr.Column():
            output = gr.Text(label="Output")
    
    convert_btn.click(
        fn=convert_2d_to_3d,
        inputs=[input_image, text_prompt],
        outputs=output
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()