sampler

modules/sd_samplers_common.py

@@ -5,7 +5,7 @@ import torch
 from PIL import Image
 from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared, sd_models
 from modules.shared import opts, state
-from ldm_patched.modules import model_management
+from modules_forge.forge_sampler import sampling_prepare, sampling_cleanup
 import k_diffusion.sampling
 
 
@@ -177,20 +177,16 @@ def apply_refiner(cfg_denoiser, x):
     cfg_denoiser.p.extra_generation_params['Refiner'] = refiner_checkpoint_info.short_title
     cfg_denoiser.p.extra_generation_params['Refiner switch at'] = refiner_switch_at
 
+    sampling_cleanup(sd_models.model_data.get_sd_model().forge_objects.unet)
+
     with sd_models.SkipWritingToConfig():
         sd_models.reload_model_weights(info=refiner_checkpoint_info)
 
-    refiner = sd_models.model_data.get_sd_model()
-
     devices.torch_gc()
     cfg_denoiser.p.setup_conds()
     cfg_denoiser.update_inner_model()
 
-    inference_memory = refiner.current_controlnet_required_memory
-    unet_patcher = refiner.forge_objects.unet
-    model_management.load_models_gpu(
-        [unet_patcher],
-        unet_patcher.memory_required([x.shape[0]] * 2 + list(x.shape[1:])) + inference_memory)
+    sampling_prepare(sd_models.model_data.get_sd_model().forge_objects.unet, x=x)
     return True
 
 

modules/sd_samplers_kdiffusion.py

@@ -7,7 +7,7 @@ from modules.script_callbacks import ExtraNoiseParams, extra_noise_callback
 
 from modules.shared import opts
 import modules.shared as shared
-import ldm_patched.modules.model_management
+from modules_forge.forge_sampler import sampling_prepare, sampling_cleanup
 
 
 samplers_k_diffusion = [
@@ -141,11 +141,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
         return sigmas
 
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        inference_memory = self.model_wrap.inner_model.current_controlnet_required_memory
         unet_patcher = self.model_wrap.inner_model.forge_objects.unet
-        ldm_patched.modules.model_management.load_models_gpu(
-            [unet_patcher],
-            unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
+        sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x)
 
         self.model_wrap.log_sigmas = self.model_wrap.log_sigmas.to(unet_patcher.current_device)
         self.model_wrap.sigmas = self.model_wrap.sigmas.to(unet_patcher.current_device)
@@ -201,14 +198,13 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         self.add_infotext(p)
 
+        sampling_cleanup(unet_patcher)
+
         return samples
 
     def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        inference_memory = self.model_wrap.inner_model.current_controlnet_required_memory
         unet_patcher = self.model_wrap.inner_model.forge_objects.unet
-        ldm_patched.modules.model_management.load_models_gpu(
-            [unet_patcher],
-            unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
+        sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x)
 
         self.model_wrap.log_sigmas = self.model_wrap.log_sigmas.to(unet_patcher.current_device)
         self.model_wrap.sigmas = self.model_wrap.sigmas.to(unet_patcher.current_device)
@@ -256,6 +252,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         self.add_infotext(p)
 
+        sampling_cleanup(unet_patcher)
+
         return samples
 
 

modules/sd_samplers_timesteps.py

@@ -7,7 +7,7 @@ from modules.script_callbacks import ExtraNoiseParams, extra_noise_callback
 
 from modules.shared import opts
 import modules.shared as shared
-import ldm_patched.modules.model_management
+from modules_forge.forge_sampler import sampling_prepare, sampling_cleanup
 
 
 samplers_timesteps = [
@@ -98,11 +98,8 @@ class CompVisSampler(sd_samplers_common.Sampler):
         return timesteps
 
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        inference_memory = self.model_wrap.inner_model.current_controlnet_required_memory
         unet_patcher = self.model_wrap.inner_model.forge_objects.unet
-        ldm_patched.modules.model_management.load_models_gpu(
-            [unet_patcher],
-            unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
+        sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x)
 
         self.model_wrap.inner_model.alphas_cumprod = self.model_wrap.inner_model.alphas_cumprod.to(unet_patcher.current_device)
 
@@ -146,14 +143,13 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         self.add_infotext(p)
 
+        sampling_cleanup(unet_patcher)
+
         return samples
 
     def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        inference_memory = self.model_wrap.inner_model.current_controlnet_required_memory
         unet_patcher = self.model_wrap.inner_model.forge_objects.unet
-        ldm_patched.modules.model_management.load_models_gpu(
-            [unet_patcher],
-            unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
+        sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x)
 
         self.model_wrap.inner_model.alphas_cumprod = self.model_wrap.inner_model.alphas_cumprod.to(unet_patcher.current_device)
 
@@ -178,6 +174,8 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         self.add_infotext(p)
 
+        sampling_cleanup(unet_patcher)
+
         return samples