import spaces
import gradio as gr
from gradio_pannellum import Pannellum
import torch
from huggingface_hub import snapshot_download
from txt2panoimg import Text2360PanoramaImagePipeline
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)

@spaces.GPU
def text_to_pano(prompt, upscale):
    input_data = {'prompt': prompt, 'upscale': upscale, 'refinement': False}
    output = txt2panoimg(input_data)
    return output, output  

title = """<h1 align="center">SD-T2I-360PanoImage</h1>
<p align="center">360° Panorama Image Generation</p>
<p><center>
<a href="https://github.com/ArcherFMY/SD-T2I-360PanoImage/" target="_blank">[Github]</a>
<a href="https://huggingface.co/archerfmy0831/sd-t2i-360panoimage" target="_blank">[Models]</a>
</center></p>
"""

with gr.Blocks(theme='bethecloud/storj_theme') as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            t2p_input = gr.Textbox(label="Enter your prompt", lines=3)
            t2p_upscale = gr.Checkbox(label="Upscale (takes about 60 seconds 6144x3072 resolution)")
            t2p_generate = gr.Button("Generate Panorama")
        with gr.Column(variant="default"):
            t2p_output = Pannellum(show_label=False, interactive=True)
    
    with gr.Row():
        t2p_image_output = gr.Image(label="Generated Image")

    # Add a hidden component to store a random value
    update_trigger = gr.State(value=0)

    def generate_with_update(prompt, upscale, trigger):
        output, image = text_to_pano(prompt, upscale)
        return output, image, trigger + 1

    t2p_generate.click(
        generate_with_update,
        inputs=[t2p_input, t2p_upscale, update_trigger],
        outputs=[t2p_output, t2p_image_output, update_trigger]
    )

demo.launch()