Update app.py

app.py

@@ -13,10 +13,20 @@ import torch
 from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 
 
-def infer(prompt, pip_sd, pip_freeu):
-    
+model_id = "CompVis/stable-diffusion-v1-4"
+        
+pip_sd = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pip_sd = pip_sd.to("cuda")
 
+pip_freeu = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pip_freeu = pip_freeu.to("cuda")
+# -------- freeu block registration
+register_free_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
+register_free_crossattn_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
+# -------- freeu block registration
 
+def infer(prompt):
+    
     print("Generating SD:")
     sd_image = pip_sd(prompt).images[0]  
 
@@ -105,26 +115,25 @@ with block:
                 )
             btn = gr.Button("Generate image", scale=0)
 
-        # sd_options = gr.Dropdown(options, label="SD options")
-        sd_options = gr.Dropdown(options, value='SD1.4', label="SD options")
-        model_id = "CompVis/stable-diffusion-v1-4"
+        # sd_options = gr.Dropdown(options, value='SD1.4', label="SD options")
+        # model_id = "CompVis/stable-diffusion-v1-4"
         
-        if sd_options == 'SD1.4':
-            model_id = "CompVis/stable-diffusion-v1-4"
-        elif sd_options == 'SD1.5':
-            model_id = "runwayml/stable-diffusion-v1-5"
-        elif sd_options == 'SD2.1':
-            model_id = "stabilityai/stable-diffusion-2-1"
+        # if sd_options == 'SD1.4':
+        #     model_id = "CompVis/stable-diffusion-v1-4"
+        # elif sd_options == 'SD1.5':
+        #     model_id = "runwayml/stable-diffusion-v1-5"
+        # elif sd_options == 'SD2.1':
+        #     model_id = "stabilityai/stable-diffusion-2-1"
             
-        pip_sd = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-        pip_sd = pip_sd.to("cuda")
-
-        pip_freeu = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-        pip_freeu = pip_freeu.to("cuda")
-        # -------- freeu block registration
-        register_free_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
-        register_free_crossattn_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
-        # -------- freeu block registration
+        # pip_sd = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+        # pip_sd = pip_sd.to("cuda")
+
+        # pip_freeu = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+        # pip_freeu = pip_freeu.to("cuda")
+        # # -------- freeu block registration
+        # register_free_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
+        # register_free_crossattn_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
+        # # -------- freeu block registration
         
         with gr.Accordion('FreeU Parameters', open=False):
             
@@ -158,11 +167,11 @@ with block:
                 image_2_label = gr.Markdown("FreeU")
             
 
-    ex = gr.Examples(examples=examples, fn=infer, inputs=[text, pip_sd, pip_freeu], outputs=[image_1, image_2], cache_examples=False)
+    ex = gr.Examples(examples=examples, fn=infer, inputs=[text], outputs=[image_1, image_2], cache_examples=False)
     ex.dataset.headers = [""]
 
-    text.submit(infer, inputs=[text, pip_sd, pip_freeu], outputs=[image_1, image_2])
-    btn.click(infer, inputs=[text, pip_sd, pip_freeu], outputs=[image_1, image_2])
+    text.submit(infer, inputs=[text], outputs=[image_1, image_2])
+    btn.click(infer, inputs=[text], outputs=[image_1, image_2])
 
 block.launch()
 # block.queue(default_enabled=False).launch(share=False)