Update app.py

app.py

@@ -75,7 +75,7 @@ import diffusers
 import warnings
 from stablepy import logger
 from diffusers import FluxPipeline
-# import urllib.parse
+import urllib.parse
 import subprocess
 
 os.system("pip list")
@@ -246,13 +246,13 @@ class GuiSD:
                     self.active_downloads.add(lock_key)
                     break
 
-            yield f"Waiting for existing download to finish: {model_name}..."
+            yield f"Waiting for existing download to finish: {escape_html(model_name)}..."
             time.sleep(1)
 
         try:
             # download link model > model_name
             if model_name.startswith("http"):
-                yield f"Downloading model: {model_name}"
+                yield f"Downloading model: {escape_html(model_name)}"
                 model_name = download_things(DIRECTORY_MODELS, model_name, HF_TOKEN, CIVITAI_API_KEY)
                 if not model_name:
                     raise ValueError("Error retrieving model information from URL")
@@ -294,7 +294,7 @@ class GuiSD:
 
         self.status_loading = True
 
-        yield f"Loading model: {model_name}"
+        yield f"Loading model: {escape_html(model_name)}"
 
         if vae_model == "BakedVAE":
             vae_model = model_name
@@ -351,7 +351,7 @@ class GuiSD:
         self.last_load = datetime.now()
         self.status_loading = False
 
-        yield f"Model loaded: {model_name}"
+        yield f"Model loaded: {escape_html(model_name)}"
 
     # @spaces.GPU(duration=59)
     @torch.inference_mode()
@@ -481,7 +481,7 @@ class GuiSD:
 
         vae_model = vae_model if vae_model != "None" else None
         loras_list = [lora1, lora2, lora3, lora4, lora5, lora6, lora7]
-        vae_msg = f"VAE: {vae_model}" if vae_model else ""
+        vae_msg = f"VAE: {escape_html(vae_model)}" if vae_model else ""
         msg_lora = ""
 
         logger.debug(f"Config model: {model_name}, {vae_model}, {loras_list}")
@@ -675,9 +675,9 @@ class GuiSD:
 
                 for status, lora in zip(self.model.lora_status, self.model.lora_memory):
                     if status:
-                        msg_lora += f"<br>Loaded: {lora}"
+                        msg_lora += f"<br>Loaded: {escape_html(lora)}"
                     elif status is not None:
-                        msg_lora += f"<br>Error with: {lora}"
+                        msg_lora += f"<br>Error with: {escape_html(lora)}"
 
                 if msg_lora:
                     info_images += msg_lora
@@ -686,7 +686,7 @@ class GuiSD:
 
                 download_links = "<br>".join(
                     [
-                        f'<a href="{path.replace("/images/", f"/gradio_api/file={allowed_path}/")}" download="{os.path.basename(path)}">Download Image {i + 1}</a>'
+                        f'<a href="{path.replace("/images/", f"/gradio_api/file={allowed_path}/")}" download="{escape_html(os.path.basename(path))}">Download Image {i + 1}</a>'
                         for i, path in enumerate(image_path)
                     ]
                 )
@@ -770,7 +770,7 @@ def sd_gen_generate_pipeline(*args):
             )
             gr.Info(f"LoRAs in cache: {lora_cache_msg}")
 
-    msg_request = f"Requesting {gpu_duration_arg}s. of GPU time.\nModel: {sd_gen.model.base_model_id}"
+    msg_request = f"Requesting {gpu_duration_arg}s. of GPU time.\nModel: {escape_html(sd_gen.model.base_model_id)}"
     if verbose_arg:
         gr.Info(msg_request)
         print(msg_request)
@@ -935,6 +935,7 @@ with gr.Blocks(fill_width=True, fill_height=False) as app:
                         "Image resolution": gr.update(value=1024),
                         "Strength": gr.update(),
                         "Prompt emphasis": gr.update(),
+                        "VAE": gr.update(),
                     }
 
                     # Generate up to 7 LoRAs
@@ -1027,6 +1028,15 @@ with gr.Blocks(fill_width=True, fill_height=False) as app:
                                         val = filtered_models[0]
                                     else:
                                         val = name_model
+                                if key == "VAE":
+                                    val_unquote = urllib.parse.unquote(str(val))
+                                    filtered_vaes = [v for v in vae_model_list if str(val) in v or val_unquote in v]
+                                    if filtered_vaes:
+                                        val = filtered_vaes[0]
+                                    elif val in ["None", "BakedVAE"]:
+                                        val = val
+                                    else:
+                                        continue
                                 if key == "Hires upscaler":
                                     if val not in UPSCALER_KEYS:
                                         continue
@@ -1187,7 +1197,7 @@ with gr.Blocks(fill_width=True, fill_height=False) as app:
                         label="Style Prompt",
                         interactive=True,
                     )
-                    style_json_gui = gr.File(label="Style JSON File")
+                    style_json_gui = gr.File(label="Style JSON File", file_count="single", file_types=[".json"])
                     style_button = gr.Button("Load styles")
 
                     def load_json_style_file(json):
@@ -1290,6 +1300,7 @@ with gr.Blocks(fill_width=True, fill_height=False) as app:
                         image_resolution_gui,
                         strength_gui,
                         prompt_syntax_gui,
+                        vae_model_gui,
                         lora1_gui,
                         lora_scale_1_gui,
                         lora2_gui,