[BUG]: Sampler failing on latest commit

[Chubbly]

I'm not sure if this is caused by the smart memory management changes Comfy introduced or a strange interaction with previews (turning previews off doesn't fix it) but can't run any workflows.

Error occurred when executing KSamplerAdvanced:

module 'comfy.model_management' has no attribute 'load_if_low_vram'

File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\execution.py", line 151, in recursive_execute
output_data, output_ui = get_output_data(obj, input_data_all)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\execution.py", line 81, in get_output_data
return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\execution.py", line 74, in map_node_over_list
results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\nodes.py", line 1240, in sample
return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\custom_nodes\AIT\AITemplate\AITemplate.py", line 175, in common_ksampler
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-Impact-Pack\modules\impact\hacky.py", line 9, in informative_sample
return original_sample(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\custom_nodes\AIT\AITemplate\AITemplate.py", line 308, in sample
samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 733, in sample
samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\sampling.py", line 701, in sample_dpmpp_2m_sde_gpu
return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\sampling.py", line 613, in sample_dpmpp_2m_sde
denoised = model(x, sigmas[i] * s_in, **extra_args)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 323, in forward
out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, cond_concat=cond_concat, model_options=model_options, seed=seed)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\external.py", line 125, in forward
eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\external.py", line 151, in get_eps
return self.inner_model.apply_model(*args, **kwargs)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 311, in apply_model
out = sampling_function(self.inner_model.apply_model, x, timestep, uncond, cond, cond_scale, cond_concat, model_options=model_options, seed=seed)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 289, in sampling_function
cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options)
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 241, in calc_cond_uncond_batch
c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
File "C:\Users\hoffl\Stable Diffusion\ComfyUI_windows_portable\ComfyUI\custom_nodes\AIT\AITemplate\AITemplate.py", line 409, in get_control
self.control_model = comfy.model_management.load_if_low_vram(self.control_model)

[hlky]

caused by the smart memory management changes

Yes it would appear so. Those changes are not important, it's just a refactor, so just rollback to before 89a0767 for now.

[asagi4]

I fixed this for myself with

From 79a1181cc091f9f3707d7b431f5994d3e19260a9 Mon Sep 17 00:00:00 2001
From: asagi4 <130366179+asagi4@users.noreply.github.com>
Date: Thu, 17 Aug 2023 20:15:36 +0300
Subject: [PATCH 1/2] Fix controlnets

---
 AITemplate/AITemplate.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/AITemplate/AITemplate.py b/AITemplate/AITemplate.py
index 9f0c178..bf5c21f 100644
--- a/AITemplate/AITemplate.py
+++ b/AITemplate/AITemplate.py
@@ -342,7 +342,8 @@ class ControlNet(ControlBase):
             self.aitemplate = True
         else:
             self.aitemplate = None
-        self.control_model = control_model.to("cuda")
+        self.control_model = control_model
+        self.control_model_wrapped = comfy.sd.ModelPatcher(self.control_model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
         self.cond_hint_original = None
         self.cond_hint = None
         self.strength = 1.0
@@ -406,11 +407,11 @@ class ControlNet(ControlBase):
                 precision_scope = contextlib.nullcontext
 
             with precision_scope(comfy.model_management.get_autocast_device(self.device)):
-                self.control_model = comfy.model_management.load_if_low_vram(self.control_model)
+                comfy.model_management.load_models_gpu([self.control_model_wrapped])
                 context = torch.cat(cond['c_crossattn'], 1)
                 y = cond.get('c_adm', None)
                 control = self.control_model(x=x_noisy, hint=self.cond_hint, timesteps=t, context=context, y=y)
-                self.control_model = comfy.model_management.unload_if_low_vram(self.control_model)
+                comfy.model_management.unload_model_clones(self.control_model_wrapped)
         else:
             # AITemplate inference, returns the same as regular
             control = self.aitemplate_controlnet(x_noisy, t, cond, self.cond_hint)
-- 
2.39.3

but I don't know if it's a complete solution or just a paperbag patch.

[hlky]

@asagi4 It's about what I would have expected, if it works for you it should be good to go. Do you want to make a PR? If not I can do it and add you as co-author.