Spaces:
Running
on
Zero
Running
on
Zero
| 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") | |
| 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() |