Update app.py

app.py

@@ -6,15 +6,6 @@ from diffusers import StableDiffusionPipeline
 from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 
 
-# model_id = "CompVis/stable-diffusion-v1-4"
-model_id = "./stable-diffusion-v1-4"
-pip_1_4 = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pip_1_4 = pip_1_4.to("cuda")
-
-# model_id = "runwayml/stable-diffusion-v1-5"
-model_id = "./stable-diffusion-v1-5"
-pip_1_5 = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pip_1_5 = pip_1_5.to("cuda")
 
 # model_id = "stabilityai/stable-diffusion-2-1"
 model_id = "./stable-diffusion-2-1"
@@ -32,12 +23,14 @@ def infer(prompt, sd_options, seed, b1, b2, s1, s2):
     global seed_prev
     global sd_image_prev
 
-    if sd_options == 'SD1.5':
-        pip = pip_1_5
-    elif sd_options == 'SD2.1':
-        pip = pip_2_1
-    else:
-        pip = pip_1_4
+    # if sd_options == 'SD1.5':
+    #     pip = pip_1_5
+    # elif sd_options == 'SD2.1':
+    #     pip = pip_2_1
+    # else:
+    #     pip = pip_1_4
+
+    pip = pip_2_1
 
     run_baseline = False
     if prompt != prompt_prev or sd_options != sd_options_prev or seed != seed_prev:
@@ -78,6 +71,9 @@ examples = [
     [
         "half human half cat, a human cat hybrid",
     ],
+    [
+        "a drone flying over a snowy forest."
+    ],
 ]
     
     
@@ -94,7 +90,7 @@ h1 {
 
 block = gr.Blocks(css='style.css')
 
-options = ['SD1.4', 'SD1.5', 'SD2.1']
+options = ['SD2.1']
 
 with block:
     gr.Markdown("SD vs. FreeU.")