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

with gr.Blocks(theme='bethecloud/storj_theme') as demo:
    gr.Markdown("# 360° Panorama Image Generation")

    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,
        inputs=[t2p_input, t2p_upscale],
        outputs=t2p_output
    )

demo.launch()