import torch
from diffusers import DiffusionPipeline
import gradio as gr

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image",
    torch_dtype=torch.bfloat16
)

pipe.load_lora_weights(
    "ProGamerGov/qwen-360-diffusion",
    weight_name="qwen-360-diffusion-2512-int8-bf16-v2.safetensors"
)

pipe.to(DEVICE)


def generate(prompt):
    image = pipe(
        prompt,
        width=2048,
        height=1024,
        num_inference_steps=25,
        true_cfg_scale=4.0
    ).images[0]

    return image


demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(
        label="Prompt",
        value="360 degree panorama with equirectangular projection, snowy mountain valley"
    ),
    outputs=gr.Image()
)

demo.launch()