Update app.py

app.py

@@ -2,14 +2,40 @@ from diffusers import StableDiffusionPipeline
 import gradio as gr
 import torch
 
-device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
+models = [
+  "nitrosocke/Arcane-Diffusion",
+  "nitrosocke/archer-diffusion",
+  "nitrosocke/elden-ring-diffusion",
+  "nitrosocke/spider-verse-diffusion"
+]
+
+prompt_prefixes = {
+  models[0]: "arcane style ",
+  models[1]: "archer style ",
+  models[2]: "elden ring style ",
+  models[3]: "spiderverse style "
+}
 
-pipe = StableDiffusionPipeline.from_pretrained("nitrosocke/Arcane-Diffusion", torch_dtype=torch.float16)
+current_model = models[0]
+pipe = StableDiffusionPipeline.from_pretrained(current_model, torch_dtype=torch.float16)
 if torch.cuda.is_available():
   pipe = pipe.to("cuda")
 
+device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
+
+def on_model_change(model):
+
+    global current_model
+    global pipe
+    if model != current_model:
+        current_model = model
+        pipe = StableDiffusionPipeline.from_pretrained(current_model, torch_dtype=torch.float16)
+        if torch.cuda.is_available():
+            pipe = pipe.to("cuda")
+
 def inference(prompt, guidance, steps):
-    prompt = prompt + ", arcane style"
+
+    prompt = prompt_prefixes[current_model] + prompt
     image = pipe(prompt, num_inference_steps=int(steps), guidance_scale=guidance, width=512, height=512).images[0]
     return image
 
@@ -26,11 +52,11 @@ with gr.Blocks() as demo:
                 "
               >
                 <h1 style="font-weight: 900; margin-bottom: 7px;">
-                  Arcane Diffusion
+                  Finetuned Diffusion
                 </h1>
               </div>
               <p style="margin-bottom: 10px; font-size: 94%">
-               Demo for a fine-tuned Stable Diffusion model trained on images from the TV Show Arcane.
+               Demo for multiple fine-tuned Stable Diffusion models, trained on different styles: Arcane, Archer, Elden Ring, Spiderverse.
               </p>
             </div>
         """
@@ -38,20 +64,22 @@ with gr.Blocks() as demo:
     with gr.Row():
         
         with gr.Column():
-            prompt = gr.Textbox(label="prompt", placeholder="' , arcane style' is appended automatically")
-            guidance = gr.Slider(label="guidance scale", value=7.5, maximum=15)
-            steps = gr.Slider(label="steps", value=50, maximum=100, minimum=2)
+            model = gr.Dropdown(label="Model", choices=models, value=models[0])
+            prompt = gr.Textbox(label="Prompt", placeholder="{} is added automatically".format(prompt_prefixes[current_model]))
+            guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
+            steps = gr.Slider(label="Steps", value=50, maximum=100, minimum=2)
             run = gr.Button(value="Run")
             gr.Markdown(f"Running on: {device}")
         with gr.Column():
             image_out = gr.Image(height=512)
 
+    model.change(on_model_change, inputs=model, outputs=[])
     run.click(inference, inputs=[prompt, guidance, steps], outputs=image_out)
     gr.Examples([
         ["jason bateman disassembling the demon core", 7.5, 50],
         ["portrait of dwayne johnson", 7.0, 75],
-        ["portrait of a beautiful alyx vance half life, volume lighting, concept art, by greg rutkowski!!, colorful, xray melting colors!!", 7, 50],
-        ["Aloy from Horizon: Zero Dawn, half body portrait, videogame cover art, highly detailed, digital painting, artstation, concept art, smooth, detailed armor, sharp focus, beautiful face, illustration, art by Artgerm and greg rutkowski and alphonse mucha", 7, 50],
+        ["portrait of a beautiful alyx vance half life", 7, 50],
+        ["Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7, 50],
         ["fantasy portrait painting, digital art", 4, 30],
     ], [prompt, guidance, steps], image_out, inference, cache_examples=torch.cuda.is_available())
     gr.HTML('''