Upload 2 files

app.py

@@ -48,7 +48,7 @@ pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_ret
 # https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
 # https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
 @spaces.GPU()
-def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm=True)): # , progress=gr.Progress(track_tqdm=True) # gradio.exceptions.Error: 'Model load Error: too many values to unpack (expected 2)'
+def change_base_model(repo_id: str, cn_on: bool): # , progress=gr.Progress(track_tqdm=True) # gradio.exceptions.Error: 'Model load Error: too many values to unpack (expected 2)'
     global pipe
     global controlnet_union
     global controlnet
@@ -59,7 +59,7 @@ def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm
         if cn_on:
             #progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
-            #clear_cache()
+            clear_cache()
             controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype)#.to(device)
             controlnet = FluxMultiControlNetModel([controlnet_union])#.to(device)
             pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype)#.to(device)
@@ -71,7 +71,7 @@ def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm
         else:
             #progress(0, desc=f"Loading model: {repo_id}")
             print(f"Loading model: {repo_id}")
-            #clear_cache()
+            clear_cache()
             pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype)#, vae=taef1 .to(device)
             pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
             last_model = repo_id

mod.py

@@ -74,10 +74,13 @@ def is_repo_exists(repo_id):
 
 
 def clear_cache():
-    torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_peak_memory_stats()
-    gc.collect()
+    try:
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()
+        torch.cuda.reset_peak_memory_stats()
+        gc.collect()
+    except Exception as e:
+        print(e)
 
 
 def deselect_lora():
@@ -348,3 +351,4 @@ load_prompt_enhancer.zerogpu = True
 fuse_loras.zerogpu = True
 preprocess_image.zerogpu = True
 get_control_params.zerogpu = True
+clear_cache.zerogpu = True