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"
os.makedirs(OUTPUT_DIR, exist_ok=True)

@spaces.GPU
@torch.inference_mode()
def generate_and_process_3d(image: Image.Image) -> tuple[str | None, str | 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/gradio_api/file={Path(output_path).resolve()}"
        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)}")
        return None, None

css = """
.image { margin: 0px auto; object-fit: contain; !important; }
.info { align-items: center; text-align: center; }
"""

# Create Gradio app using Blocks
with gr.Blocks(theme=gr.themes.Soft(), fill_width=True, fill_height=True, elem_id="container", css=css) 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.", elem_classes="info")
    
    input_img = gr.Image(
        type="pil", label="Input Image", sources="upload", image_mode="RGBA", width=40, elem_classes="image",
    )

    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, show_copy_button=True)

    # 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()])