import gradio as gr
from gradio_pannellum import Pannellum
import torch
from huggingface_hub import snapshot_download
from txt2panoimg import Text2360PanoramaImagePipeline
from img2panoimg import Image2360PanoramaImagePipeline
from PIL import Image

# Download the model
model_path = snapshot_download("archerfmy0831/sd-t2i-360panoimage")

# Initialize pipelines
txt2panoimg = Text2360PanoramaImagePipeline(model_path, torch_dtype=torch.float16)
img2panoimg = Image2360PanoramaImagePipeline(model_path, torch_dtype=torch.float16)

# Load the default mask image
default_mask = Image.open("i2p-mask.jpg").convert("RGB")

def text_to_pano(prompt, upscale):
    input_data = {'prompt': prompt, 'upscale': upscale}
    output = txt2panoimg(input_data)
    return output

def image_to_pano(image, mask, prompt, upscale):
    image = image.resize((512, 512))
    if mask is None:
        mask = default_mask.resize((512, 512))
    else:
        mask = mask.resize((512, 512))
    input_data = {
        'prompt': prompt,
        'image': image,
        'mask': mask,
        'upscale': upscale
    }
    output = img2panoimg(input_data)
    return output

def text_to_pano_interface(prompt, upscale):
    output = text_to_pano(prompt, upscale)
    return Pannellum().update(panorama=output)

def image_to_pano_interface(image, mask, prompt, upscale):
    output = image_to_pano(image, mask, prompt, upscale)
    return Pannellum().update(panorama=output)

with gr.Blocks() as demo:
    gr.Markdown("# 360° Panorama Image Generation")

    with gr.Tab("Text to 360° Panorama"):
        with gr.Row():
            with gr.Column():
                t2p_input = gr.Textbox(label="Enter your prompt", lines=3)
                t2p_upscale = gr.Checkbox(label="Upscale (requires >16GB GPU)")
                t2p_generate = gr.Button("Generate Panorama")
            with gr.Column():
                t2p_output = Pannellum(label="Generated 360° Panorama")
        
        t2p_generate.click(
            text_to_pano_interface,
            inputs=[t2p_input, t2p_upscale],
            outputs=t2p_output
        )

    with gr.Tab("Image to 360° Panorama"):
        with gr.Row():
            with gr.Column():
                i2p_image = gr.Image(label="Upload Input Image", type="pil")
                i2p_mask = gr.Image(label="Upload Mask Image (Optional)", type="pil")
                i2p_prompt = gr.Textbox(label="Enter your prompt", lines=3)
                i2p_upscale = gr.Checkbox(label="Upscale (requires >16GB GPU)")
                i2p_generate = gr.Button("Generate Panorama")
            with gr.Column():
                i2p_output = Pannellum(label="Generated 360° Panorama")
        
        i2p_generate.click(
            image_to_pano_interface,
            inputs=[i2p_image, i2p_mask, i2p_prompt, i2p_upscale],
            outputs=i2p_output
        )

demo.launch()