fix colab

launch.py

@@ -20,7 +20,7 @@ def prepare_environment():
     xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.20')
 
     comfy_repo = os.environ.get('COMFY_REPO', "https://github.com/comfyanonymous/ComfyUI")
-    comfy_commit_hash = os.environ.get('COMFY_COMMIT_HASH', "c9b562aed153cb35d4ce4126caf86995b0c63b12")
+    comfy_commit_hash = os.environ.get('COMFY_COMMIT_HASH', "2bc12d3d22efb5c63ae3a7fc342bb2dd16b31735")
 
     print(f"Python {sys.version}")
     print(f"Fooocus version: {fooocus_version.version}")

modules/core.py

@@ -9,7 +9,7 @@ import comfy.utils
 
 from comfy.sd import load_checkpoint_guess_config
 from nodes import VAEDecode, EmptyLatentImage, CLIPTextEncode
-from comfy.sample import prepare_mask, broadcast_cond, get_additional_models, cleanup_additional_models
+from comfy.sample import prepare_mask, broadcast_cond, load_additional_models, cleanup_additional_models
 from modules.samplers_advanced import KSampler, KSamplerWithRefiner
 from modules.patch import patch_all
 
@@ -133,6 +133,8 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
     if noise_mask is not None:
         noise_mask = prepare_mask(noise_mask, noise.shape, device)
 
+    comfy.model_management.load_model_gpu(model)
+
     real_model = model.model
 
     noise = noise.to(device)
@@ -141,9 +143,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
     positive_copy = broadcast_cond(positive, noise.shape[0], device)
     negative_copy = broadcast_cond(negative, noise.shape[0], device)
 
-    models = get_additional_models(positive, negative)
-    comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(
-        noise.shape[0] * noise.shape[2] * noise.shape[3]))
+    models = load_additional_models(positive, negative, model.model_dtype())
 
     sampler = KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler,
                        denoise=denoise, model_options=model.model_options)
@@ -206,6 +206,8 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
     if noise_mask is not None:
         noise_mask = prepare_mask(noise_mask, noise.shape, device)
 
+    comfy.model_management.load_model_gpu(model)
+
     noise = noise.to(device)
     latent_image = latent_image.to(device)
 
@@ -215,9 +217,7 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
     refiner_positive_copy = broadcast_cond(refiner_positive, noise.shape[0], device)
     refiner_negative_copy = broadcast_cond(refiner_negative, noise.shape[0], device)
 
-    models = get_additional_models(positive, negative)
-    comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(
-        noise.shape[0] * noise.shape[2] * noise.shape[3]))
+    models = load_additional_models(positive, negative, model.model_dtype())
 
     sampler = KSamplerWithRefiner(model=model, refiner_model=refiner, steps=steps, device=device,
                                   sampler=sampler_name, scheduler=scheduler,

modules/samplers_advanced.py

@@ -152,8 +152,7 @@ class KSamplerWithRefiner:
                                           noise.shape[3], noise.shape[2], self.device, "negative")
 
         def refiner_switch():
-            comfy.model_management.load_models_gpu([self.refiner_model_patcher], comfy.model_management.batch_area_memory(
-                noise.shape[0] * noise.shape[2] * noise.shape[3]))
+            comfy.model_management.load_model_gpu(self.refiner_model_patcher)
             self.model_denoise.inner_model = self.refiner_model_denoise.inner_model
             for i in range(len(positive)):
                 positive[i] = refiner_positive[i]