Update app.py

app.py

@@ -9,6 +9,10 @@ vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype
 scheduler = DDIMScheduler.from_pretrained(model, subfolder="scheduler")
 pipe = DiffusionPipeline.from_pretrained(model, vae=vae, scheduler=scheduler, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
 
+model_15 = "runwayml/stable-diffusion-v1-5"
+scheduler_15 = DDIMScheduler.from_pretrained(model_15, subfolder="scheduler")
+pipe_15 = DiffusionPipeline.from_pretrained(model_15, vae=vae, scheduler=scheduler_15, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
+
 #pipe.enable_model_cpu_offload()
 pipe.enable_vae_tiling()
 
@@ -17,14 +21,28 @@ apply_hidiffusion(pipe)
 
 @spaces.GPU
 def run_hidiffusion(prompt, negative_prompt, progress=gr.Progress(track_tqdm=True)):
-    return pipe(prompt, guidance_scale=7.5, height=4096, width=4096, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25).images[0]
-        
+    return pipe(prompt, guidance_scale=7.5, height=2048, width=2048, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25).images[0]
+
+@spaces.GPU
+def run_hidiffusion_15(prompt, negative_prompt, progress=gr.Progress(track_tqdm=True)):
+    return pipe_15(prompt, guidance_scale=7.5, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25).images[0]
+
+
 with gr.Blocks() as demo:
-    prompt = gr.Textbox()
-    negative_prompt = gr.Textbox()
-    btn = gr.Button("Run")
+    gr.Markdown("# HiDiffusion Demo")
+    gr.Markdown("Make Stable Diffusion generated higher resolution images than what it was trained for")
+    with gr.Tab("SDXL in 2048x2048"):
+        with gr.Row():
+            prompt = gr.Textbox(label="Prompt")
+            negative_prompt = gr.Textbox(label="Negative Prompt")
+        btn = gr.Button("Run")
+    with gr.Tab("SD1.5 in 1024x1024"):
+        with gr.Row():
+            prompt_15 = gr.Textbox(label="Prompt")
+            negative_prompt_15 = gr.Textbox(label="Negative Prompt")
+        btn_15 = gr.Button("Run")
     output = gr.Image()
 
     btn.click(fn=run_hidiffusion, inputs=[prompt, negative_prompt], outputs=[output])
-    
+    btn_15.click(fn=run_hidiffusion, inputs=[prompt_15, negative_prompt_15], outputs=[output])
 demo.launch()