app.py

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(duration=200)
def text_to_pano(prompt, upscale):
    input_data = {'prompt': prompt, 'upscale': upscale}
    output = txt2panoimg(input_data)
    return 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 160 seconds 6144x3072 resolution)")
            t2p_generate = gr.Button("Generate Panorama")
        with gr.Column():
            t2p_output = Pannellum(label="Generated 360° Panorama")
    
    t2p_generate.click(
        text_to_pano,
        inputs=[t2p_input, t2p_upscale],
        outputs=t2p_output
    )

demo.launch()