2.0.80 (#520)

* Rework many patches and some UI details.
* Speed up processing.
* Move Colab to independent branch.
* Implemented CFG Scale and TSNR correction when CFG is bigger than 10.
* Implemented Developer Mode with more options to debug.

modules/default_pipeline.py

@@ -2,23 +2,49 @@ import modules.core as core
 import os
 import torch
 import modules.path
-import modules.virtual_memory as virtual_memory
 import comfy.model_management
 
+from comfy.model_patcher import ModelPatcher
 from comfy.model_base import SDXL, SDXLRefiner
-from modules.patch import cfg_patched, patched_model_function
 from modules.expansion import FooocusExpansion
 
 
 xl_base: core.StableDiffusionModel = None
 xl_base_hash = ''
 
-xl_refiner: core.StableDiffusionModel = None
-xl_refiner_hash = ''
-
 xl_base_patched: core.StableDiffusionModel = None
 xl_base_patched_hash = ''
 
+xl_refiner: ModelPatcher = None
+xl_refiner_hash = ''
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def assert_model_integrity():
+    error_message = None
+
+    if xl_base is None:
+        error_message = 'You have not selected SDXL base model.'
+
+    if xl_base_patched is None:
+        error_message = 'You have not selected SDXL base model.'
+
+    if not isinstance(xl_base.unet.model, SDXL):
+        error_message = 'You have selected base model other than SDXL. This is not supported yet.'
+
+    if not isinstance(xl_base_patched.unet.model, SDXL):
+        error_message = 'You have selected base model other than SDXL. This is not supported yet.'
+
+    if xl_refiner is not None:
+        if not isinstance(xl_refiner.model, SDXLRefiner):
+            error_message = 'You have selected refiner model other than SDXL refiner. This is not supported yet.'
+
+    if error_message is not None:
+        raise NotImplementedError(error_message)
+
+    return True
+
 
 @torch.no_grad()
 @torch.inference_mode()
@@ -31,24 +57,13 @@ def refresh_base_model(name):
     if xl_base_hash == model_hash:
         return
 
-    if xl_base is not None:
-        xl_base.to_meta()
-        xl_base = None
+    xl_base = None
+    xl_base_hash = ''
+    xl_base_patched = None
+    xl_base_patched_hash = ''
 
     xl_base = core.load_model(filename)
-    if not isinstance(xl_base.unet.model, SDXL):
-        print('Model not supported. Fooocus only support SDXL model as the base model.')
-        xl_base = None
-        xl_base_hash = ''
-        refresh_base_model(modules.path.default_base_model_name)
-        xl_base_hash = model_hash
-        xl_base_patched = xl_base
-        xl_base_patched_hash = ''
-        return
-
     xl_base_hash = model_hash
-    xl_base_patched = xl_base
-    xl_base_patched_hash = ''
     print(f'Base model loaded: {model_hash}')
     return
 
@@ -64,28 +79,18 @@ def refresh_refiner_model(name):
     if xl_refiner_hash == model_hash:
         return
 
+    xl_refiner = None
+    xl_refiner_hash = ''
+
     if name == 'None':
-        xl_refiner = None
-        xl_refiner_hash = ''
-        print(f'Refiner unloaded.')
-        return
-
-    if xl_refiner is not None:
-        xl_refiner.to_meta()
-        xl_refiner = None
-
-    xl_refiner = core.load_model(filename)
-    if not isinstance(xl_refiner.unet.model, SDXLRefiner):
-        print('Model not supported. Fooocus only support SDXL refiner as the refiner.')
-        xl_refiner = None
-        xl_refiner_hash = ''
         print(f'Refiner unloaded.')
         return
 
+    xl_refiner = core.load_unet_only(filename)
     xl_refiner_hash = model_hash
     print(f'Refiner model loaded: {model_hash}')
 
-    xl_refiner.vae.first_stage_model.to('meta')
+    # Remove VAE
     xl_refiner.vae = None
     return
 
@@ -133,6 +138,15 @@ def clip_encode_single(clip, text, verbose=False):
     return result
 
 
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate(cond):
+    c, p = cond[0]
+    c = c[..., -1280:].clone()
+    p = p["pooled_output"].clone()
+    return [[c, {"pooled_output": p}]]
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def clip_encode(sd, texts, pool_top_k=1):
@@ -158,36 +172,20 @@ def clip_encode(sd, texts, pool_top_k=1):
     return [[torch.cat(cond_list, dim=1), {"pooled_output": pooled_acc}]]
 
 
-@torch.no_grad()
-@torch.inference_mode()
-def clear_sd_cond_cache(sd):
-    if sd is None:
-        return None
-    if sd.clip is None:
-        return None
-    sd.clip.fcs_cond_cache = {}
-    return
-
-
 @torch.no_grad()
 @torch.inference_mode()
 def clear_all_caches():
-    clear_sd_cond_cache(xl_base_patched)
-    clear_sd_cond_cache(xl_refiner)
+    xl_base.clip.fcs_cond_cache = {}
+    xl_base_patched.clip.fcs_cond_cache = {}
 
 
 @torch.no_grad()
 @torch.inference_mode()
 def refresh_everything(refiner_model_name, base_model_name, loras):
     refresh_refiner_model(refiner_model_name)
-    if xl_refiner is not None:
-        virtual_memory.try_move_to_virtual_memory(xl_refiner.unet.model)
-        virtual_memory.try_move_to_virtual_memory(xl_refiner.clip.cond_stage_model)
-
     refresh_base_model(base_model_name)
-    virtual_memory.load_from_virtual_memory(xl_base.unet.model)
-
     refresh_loras(loras)
+    assert_model_integrity()
     clear_all_caches()
     return
 
@@ -203,32 +201,7 @@ expansion = FooocusExpansion()
 
 @torch.no_grad()
 @torch.inference_mode()
-def patch_all_models():
-    assert xl_base is not None
-    assert xl_base_patched is not None
-
-    xl_base.unet.model_options['sampler_cfg_function'] = cfg_patched
-    xl_base.unet.model_options['model_function_wrapper'] = patched_model_function
-
-    xl_base_patched.unet.model_options['sampler_cfg_function'] = cfg_patched
-    xl_base_patched.unet.model_options['model_function_wrapper'] = patched_model_function
-
-    if xl_refiner is not None:
-        xl_refiner.unet.model_options['sampler_cfg_function'] = cfg_patched
-        xl_refiner.unet.model_options['model_function_wrapper'] = patched_model_function
-
-    return
-
-
-@torch.no_grad()
-@torch.inference_mode()
-def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, latent=None, denoise=1.0, tiled=False):
-    patch_all_models()
-
-    if xl_refiner is not None:
-        virtual_memory.try_move_to_virtual_memory(xl_refiner.unet.model)
-    virtual_memory.load_from_virtual_memory(xl_base.unet.model)
-
+def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0):
     if latent is None:
         empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
     else:
@@ -239,7 +212,7 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             model=xl_base_patched.unet,
             positive=positive_cond[0],
             negative=negative_cond[0],
-            refiner=xl_refiner.unet,
+            refiner=xl_refiner,
             refiner_positive=positive_cond[1],
             refiner_negative=negative_cond[1],
             refiner_switch_step=switch,
@@ -247,7 +220,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
             denoise=denoise,
-            callback_function=callback
+            callback_function=callback,
+            cfg=cfg_scale,
+            sampler_name=sampler_name
         )
     else:
         sampled_latent = core.ksampler(
@@ -258,7 +233,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
             denoise=denoise,
-            callback_function=callback
+            callback_function=callback,
+            cfg=cfg_scale,
+            sampler_name=sampler_name
         )
 
     decoded_latent = core.decode_vae(vae=xl_base_patched.vae, latent_image=sampled_latent, tiled=tiled)