Update app.py

app.py

@@ -6,19 +6,40 @@ import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
+# List of models
+models = {
+    "sdxl-turbo": "stabilityai/sdxl-turbo",
+    "MistoLine": "TheMistoAI/MistoLine"
+}
+
+# Cache to store loaded pipelines
+pipelines = {}
+
+# Function to load a model
+def load_model(model_name):
+    if model_name in pipelines:
+        return pipelines[model_name]
+
+    if model_name not in models:
+        raise ValueError(f"Model {model_name} is not available.")
+
+    model_path = models[model_name]
+    if torch.cuda.is_available():
+        torch.cuda.max_memory_allocated(device=device)
+        pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+        pipe.enable_xformers_memory_efficient_attention()
+    else:
+        pipe = DiffusionPipeline.from_pretrained(model_path, use_safetensors=True)
     pipe = pipe.to(device)
+    
+    pipelines[model_name] = pipe
+    return pipe
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def infer(model_name, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    pipe = load_model(model_name)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -64,6 +85,12 @@ with gr.Blocks(css=css) as demo:
         """)
         
         with gr.Row():
+            model_name = gr.Dropdown(
+                label="Select Model",
+                choices=list(models.keys()),
+                value="sdxl-turbo",
+                show_label=True
+            )
             
             prompt = gr.Text(
                 label="Prompt",
@@ -139,8 +166,8 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [model_name, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs = [result]
     )
 
-demo.queue().launch()
+demo.queue().launch()