Fooocus GitHub Bot Commit

This commit is generated by a GitHub bot of Fooocus

modules/sample_hijack.py

@@ -3,7 +3,7 @@ import fcbh.samplers
 import fcbh.model_management
 
 from fcbh.model_base import SDXLRefiner, SDXL
-from fcbh.sample import get_additional_models, get_models_from_cond
+from fcbh.sample import get_additional_models, get_models_from_cond, cleanup_additional_models
 from fcbh.samplers import resolve_areas_and_cond_masks, wrap_model, calculate_start_end_timesteps, \
     create_cond_with_same_area_if_none, pre_run_control, apply_empty_x_to_equal_area, encode_adm, \
     blank_inpaint_image_like
@@ -49,22 +49,6 @@ def clip_separate(cond, target_model=None, target_clip=None):
     return [[c, p]]
 
 
-@torch.no_grad()
-@torch.inference_mode()
-def force_unload_all_control(positive, negative):
-    control_nets = set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control"))
-
-    cleaned_any_model = False
-    for m in control_nets:
-        if hasattr(m, 'cleanup'):
-            m.cleanup()
-            cleaned_any_model = True
-
-    if cleaned_any_model:
-        fcbh.model_management.soft_empty_cache()
-    return
-
-
 @torch.no_grad()
 @torch.inference_mode()
 def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
@@ -129,7 +113,7 @@ def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas
         extra_args["cond_concat"] = cond_concat
 
     def refiner_switch():
-        force_unload_all_control(positive, negative)
+        cleanup_additional_models(set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")))
 
         extra_args["cond"] = positive_refiner
         extra_args["uncond"] = negative_refiner