Upload 2 files

app.py

@@ -8,15 +8,13 @@ from diffusers import DiffusionPipeline
 import copy
 import random
 import time
+from mod import models, clear_cache, get_repo_safetensors, change_base_model
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 # Initialize the base model
-models = ["camenduru/FLUX.1-dev-diffusers", "black-forest-labs/FLUX.1-schnell",
-           "sayakpaul/FLUX.1-merged", "John6666/blue-pencil-flux1-v001-fp8-flux",
-           "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux", "John6666/nepotism-fuxdevschnell-v3aio-flux"]
 base_model = models[0]
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
@@ -79,7 +77,8 @@ def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height,
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
               lora_scale, lora_repo, lora_weights, lora_trigger, progress=gr.Progress(track_tqdm=True)):
-    #if selected_index is None and not lora_repo:
+    if selected_index is None and not lora_repo:
+        gr.Info("LoRA isn't selected.")
     #    raise gr.Error("You must select a LoRA before proceeding.")
 
     if selected_index is not None and not lora_repo:
@@ -110,33 +109,12 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
     pipe.to("cpu")
     if selected_index is not None or lora_repo: pipe.unload_lora_weights()
+    clear_cache()
     return image, seed  
 
 run_lora.zerogpu = True
 
-def get_repo_safetensors(repo_id: str):
-    from huggingface_hub import HfApi
-    api = HfApi()
-    try:
-        if " " in repo_id or not api.repo_exists(repo_id): return gr.update(value="", choices=[])
-        files = api.list_repo_files(repo_id=repo_id)
-    except Exception as e:
-        print(f"Error: Failed to get {repo_id}'s info. ")
-        print(e)
-        return gr.update(choices=[])
-    files = [f for f in files if f.endswith(".safetensors")]
-    if len(files) == 0: return gr.update(value="", choices=[])
-    else: return gr.update(value=files[0], choices=files)
-
-def change_base_model(repo_id: str):
-    from huggingface_hub import HfApi
-    global pipe
-    api = HfApi()
-    try:
-        if " " in repo_id or not api.repo_exists(repo_id): return
-        pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-    except Exception as e:
-        print(e)
+
 
 css = '''
 #gen_btn{height: 100%}
@@ -147,7 +125,7 @@ css = '''
 '''
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
     title = gr.HTML(
-        """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> FLUX LoRA the Explorer</h1>""",
+        """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> FLUX LoRA the Explorer Mod</h1>""",
         elem_id="title",
     )
     selected_index = gr.State(None)

mod.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import torch
+import spaces
+from diffusers import DiffusionPipeline
+import gc
+import subprocess
+
+
+subprocess.run('pip cache purge', shell=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch.set_grad_enabled(False)
+
+
+models = ["camenduru/FLUX.1-dev-diffusers",
+          "black-forest-labs/FLUX.1-schnell",
+          "sayakpaul/FLUX.1-merged",
+          "John6666/blue-pencil-flux1-v001-fp8-flux",
+          "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux",
+          "John6666/nepotism-fuxdevschnell-v3aio-flux"
+]
+
+
+def clear_cache():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def get_repo_safetensors(repo_id: str):
+    from huggingface_hub import HfApi
+    api = HfApi()
+    try:
+        if " " in repo_id or not api.repo_exists(repo_id): return gr.update(value="", choices=[])
+        files = api.list_repo_files(repo_id=repo_id)
+    except Exception as e:
+        print(f"Error: Failed to get {repo_id}'s info. ")
+        print(e)
+        return gr.update(choices=[])
+    files = [f for f in files if f.endswith(".safetensors")]
+    if len(files) == 0: return gr.update(value="", choices=[])
+    else: return gr.update(value=files[0], choices=files)
+
+
+def change_base_model(repo_id: str):
+    from huggingface_hub import HfApi
+    global pipe
+    api = HfApi()
+    try:
+        if " " in repo_id or not api.repo_exists(repo_id): return
+        clear_cache()
+        pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+    except Exception as e:
+        print(e)