Update app.py

app.py

@@ -2,27 +2,50 @@ import gradio as gr
 import numpy as np
 import random
 
-# import spaces #[uncomment to use ZeroGPU]
+import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+MODEL_OPTIONS = [
+    ("stabilityai/sdxl-turbo", "SDXL Turbo (Быстро)"),
+    ("hakurei/waifu-diffusion", "Что-то альтернативное"),
+]
+DEFAULT_MODEL_ID = "stabilityai/sdxl-turbo"
+# model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+
+PIPELINES = {}
+
+def load_pipelines():
+    # SDXL Turbo
+    mid = "stabilityai/sdxl-turbo"
+    pipe = DiffusionPipeline.from_pretrained(mid, torch_dtype=torch_dtype)
+    pipe = pipe.to(device)
+    PIPELINES[mid] = pipe
+
+    # SD v1-4
+    mid = "hakurei/waifu-diffusion"
+    pipe = DiffusionPipeline.from_pretrained(mid, torch_dtype=torch_dtype)
+    pipe = pipe.to(device)
+    PIPELINES[mid] = pipe
+
+# load_pipelines()
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+# if torch.cuda.is_available():
+#     torch_dtype = torch.float16
+# else:
+#     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# pipe = pipe.to(device)
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+# MAX_SEED = np.iinfo(np.int32).max
+# MAX_IMAGE_SIZE = 1024
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
+@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
+    model_id, 
     prompt,
     negative_prompt,
     seed,
@@ -37,6 +60,8 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
+    
+    pipe = PIPELINES[model_id] 
 
     image = pipe(
         prompt=prompt,
@@ -105,7 +130,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=1024,
                 )
 
                 height = gr.Slider(
@@ -113,7 +138,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=1024,
                 )
 
             with gr.Row():
@@ -122,7 +147,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=0.0,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -130,7 +155,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=2,
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])