import spaces
import os
from PIL import Image
import gradio as gr
from pathlib import Path
import uuid
import torch
import shutil

OUTPUT_DIR = "output"

@spaces.GPU
@torch.inference_mode()
def generate_and_process_3d(image: Image.Image) -> tuple[str | None, Image.Image | None]:
    try:
        # Export to GLB
        unique_id = str(uuid.uuid4())
        filename = f'model_{unique_id}.glb'
        output_path = os.path.join(OUTPUT_DIR, filename)
        public_url = f"https://john6666-image-to-3d-test2.hf.space/file={output_path}"
        image_path = "image.png"
        image.save(image_path)
        shutil.copy(image_path, output_path)
 
        return output_path, public_url
        
    except Exception as e:
        print(f"Error during generation: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Create Gradio app using Blocks
with gr.Blocks() as demo:
    gr.Markdown("This space is based on [Stable Point-Aware 3D](https://huggingface.co/spaces/stabilityai/stable-point-aware-3d) by Stability AI, [Text to 3D](https://huggingface.co/spaces/jbilcke-hf/text-to-3d) by jbilcke-hf.")
    
    with gr.Row():
        input_img = gr.Image(
            type="pil", label="Input Image", sources="upload", image_mode="RGBA"
        )

    with gr.Row():
        model_output = gr.Model3D(
            label="Generated .GLB model",
            clear_color=[0.0, 0.0, 0.0, 0.0],
            visible=False
        )
        output_url = gr.Textbox(label="Output URL", value="", lines=1, interactive=False)

    # Event handler
    input_img.upload(
        fn=generate_and_process_3d,
        inputs=[input_img],
        outputs=[model_output, output_url],
        api_name="generate"
    )
    
if __name__ == "__main__":
    demo.queue().launch(ssr_mode=False, allowed_paths=[Path(OUTPUT_DIR).resolve()])