Update app.py

app.py

@@ -11,7 +11,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 MODEL_OPTIONS = [
     ("stabilityai/sdxl-turbo", "SDXL Turbo (Быстро)"),
-    ("CompVis/stable-diffusion-v1-4", "Stable Diffusion v1-4 (Классика)"),
+    # ("CompVis/stable-diffusion-v1-4", "Stable Diffusion v1-4 (Классика)"),
+    ("hakurei/waifu-diffusion", "Что-то альтернативное"),
     # ("Qwen/Qwen-Image", "Топ модель, но долго"),
 ]
 DEFAULT_MODEL_ID = "stabilityai/sdxl-turbo"
@@ -30,13 +31,19 @@ def load_pipelines():
     # SDXL Turbo
     mid = "stabilityai/sdxl-turbo"
     pipe = DiffusionPipeline.from_pretrained(mid, torch_dtype=torch_dtype)
-    # pipe = pipe.to(device)
+    pipe = pipe.to(device)
     PIPELINES[mid] = pipe
 
+    # # SD v1-4
+    # mid = "CompVis/stable-diffusion-v1-4"
+    # pipe = DiffusionPipeline.from_pretrained(mid, torch_dtype=torch_dtype)
+    # pipe = pipe.to(device)
+    # PIPELINES[mid] = pipe
+
     # SD v1-4
-    mid = "CompVis/stable-diffusion-v1-4"
+    mid = "hakurei/waifu-diffusion"
     pipe = DiffusionPipeline.from_pretrained(mid, torch_dtype=torch_dtype)
-    # pipe = pipe.to(device)
+    pipe = pipe.to(device)
     PIPELINES[mid] = pipe
 
     # # Qwen-Image 
@@ -69,7 +76,7 @@ def infer(
     generator = torch.Generator().manual_seed(seed)
     
     pipe = PIPELINES[model_id] 
-    pipe = pipe.to(device)
+    # pipe = pipe.to(device)
 
     image = pipe(
         prompt=prompt,