Update app.py

app.py

@@ -2,7 +2,6 @@
 import os
 import random
 import uuid
-import json
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -49,7 +48,7 @@ BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-def load_model(model_id):
+def load_and_prepare_model(model_id):
     pipe = StableDiffusionXLPipeline.from_pretrained(
         model_id,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -66,8 +65,8 @@ def load_model(model_id):
     
     return pipe
 
-current_model_id = MODEL_OPTIONS["RealVisXL_V4.0_Lightning"]
-pipe = load_model(current_model_id)
+# Preload and compile both models
+models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -97,9 +96,8 @@ def generate(
     num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True),
 ):
-    global pipe
-    if model_choice != current_model_id:
-        pipe = load_model(MODEL_OPTIONS[model_choice])
+    global models
+    pipe = models[model_choice]
     
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
@@ -131,14 +129,7 @@ def generate(
 
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown(DESCRIPTIONx)
-
-    with gr.Group():
         with gr.Row():
-            model_choice = gr.Dropdown(
-                label="Model",
-                choices=list(MODEL_OPTIONS.keys()),
-                value="RealVisXL_V4.0_Lightning"
-            )
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
@@ -149,6 +140,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             run_button = gr.Button("Run", scale=0)
         result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
+        with gr.Row():
+            model_choice = gr.Dropdown(
+                label="Model",
+                choices=list(MODEL_OPTIONS.keys()),
+                value="RealVisXL_V4.0_Lightning"
+            )
+
     with gr.Accordion("Advanced options", open=False, visible=False):
         num_images = gr.Slider(
             label="Number of Images",
@@ -250,4 +248,4 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown("⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.")
 
 if __name__ == "__main__":
-    demo.queue(max_size=40).launch()
+    demo.queue(max_size=40).launch()