import gradio as gr
import replicate
import os
import requests
from PIL import Image
from io import BytesIO

# Get the Replicate API token from the environment variable
replicate_api_token = os.environ.get('REPLICATE_API_TOKEN')
if not replicate_api_token:
    raise ValueError("REPLICATE_API_TOKEN environment variable is not set")

# Set the Replicate API token
os.environ['REPLICATE_API_TOKEN'] = replicate_api_token

def process_image_url(image_url, remove_bg, seed, sample_steps):
    # Prepare the input for the Replicate API
    input_data = {
        "seed": seed,
        "image_path": image_url,
        "export_video": True,
        "sample_steps": sample_steps,
        "export_texmap": False,
        "remove_background": remove_bg
    }
    
    # Run the model
    output = replicate.run(
        "aryamansital/instant_mesh:e353a25cc764e0edb0aa9033df0bf4b82318dcda6d0a0cd9f2aace90566068ac",
        input=input_data
    )
    
    # Process the output
    multi_view_image = Image.open(requests.get(output[0], stream=True).raw)
    video_url = output[1]
    glb_url = output[2]
    
    return multi_view_image, video_url, glb_url

def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# MIX-AR")
        gr.Markdown("### Important Notes:")
        gr.Markdown("- Our demo exports a .glb mesh and generates a video of the 3D model.")
        gr.Markdown("- The 3D mesh generation results highly depend on the quality of generated multi-view images. Please try a different seed value if the result is unsatisfying (Default: 42).")
        
        with gr.Row():
            image_url = gr.Textbox(label="Image URL", placeholder="Enter the URL of the image")
            multi_view_image = gr.Image(label="Generated Multi-views")
        
        with gr.Row():
            with gr.Column():
                remove_bg = gr.Checkbox(label="Remove Background", value=True)
                seed = gr.Number(label="Seed Value", value=42)
                sample_steps = gr.Slider(label="Sample Steps", minimum=1, maximum=100, value=75, step=1)
                generate_btn = gr.Button("Generate")
            
            with gr.Column():
                video_output = gr.Video(label="Generated 3D Model Video")
                glb_output = gr.File(label="Download GLB Model")
        
        gr.Markdown("Note: The generated .glb model can be viewed in 3D modeling software or web-based 3D viewers.")
        
        # Connect the components
        generate_btn.click(
            process_image_url,
            inputs=[image_url, remove_bg, seed, sample_steps],
            outputs=[multi_view_image, video_output, glb_output]
        )
        
        # Add example image URLs
        example_urls = [
            "https://replicate.delivery/pbxt/Kuf9PGp3vopyQmscqIszxJWBr1OYuzr1VWb67ijDJVBs6N7D/image.png",
            # Add more example image URLs as needed
        ]
        gr.Examples(examples=example_urls, inputs=image_url)
    
    return demo

# Launch the interface
if __name__ == "__main__":
    interface = create_interface()
    interface.launch()