import torch
import spaces
import gradio as gr
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline

prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.bfloat16)

prior_pipeline.enable_model_cpu_offload()
decoder_pipeline.enable_model_cpu_offload()

prior = prior_pipeline.to("cuda")
decoder = decoder_pipeline.to("cuda")

@spaces.GPU
def generate(prompt, negative_prompt, steps):
    prior_output = prior(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=1024,
        height=1024,
        guidance_scale=4.0,
        num_images_per_prompt=1,
        num_inference_steps=steps
    )
    
    decoder_output = decoder(
        image_embeddings=prior_output.image_embeddings.to(torch.bfloat16),
        prompt=prompt,
        guidance_scale=0.0,
        output_type="pil",
        num_inference_steps=10,
        negative_prompt=negative_prompt
    ).images[0]
    return decoder_output

with gr.Blocks() as demo:
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value="A perfectly red apple, 32K HDR, studio lighting")
        generate_btn = gr.Button("Generate")

    with gr.Row():
        output = gr.Image(label="Output")

    with gr.Accordion("Advanced", open=False):
        negative_prompt = gr.Textbox(label="Negative Prompt", value="ugly, low quality")
        steps = gr.Slider(minimum=4, maximum=50, step=1, value=20, label="Steps")

    generate_btn.click(generate, inputs=[prompt, negative_prompt, steps], outputs=output)

demo.launch()