ControlNets unload every step (very slow performance) about ait

[FizzleDorf]

Ah just seeing it break it everywhere now. I'll merge #41 so this won't be an issue. Just a little worried it feels unfinished. The TODO here is getting the controlnets to not unload every step

[FizzleDorf]

I renamed this issue to be in line with the current bug.

[asagi4]

@FizzleDorf I have this patch locally to just completely skip unloading and things work at sane speeds:

diff --git a/AITemplate/AITemplate.py b/AITemplate/AITemplate.py
index 74f4664..b2a76b4 100644
--- a/AITemplate/AITemplate.py
+++ b/AITemplate/AITemplate.py
@@ -411,7 +411,6 @@ class ControlNet(ControlBase):
                 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)
-                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)

But I'm pretty sure that's just a memory leak that papers over the performance problem, so I haven't made a PR for it.

[TheMindExpansionNetwork]

Same issue let me know you all find a work around

[TheMindExpansionNetwork]

@FizzleDorf I have this patch locally to just completely skip unloading and things work at sane speeds:

diff --git a/AITemplate/AITemplate.py b/AITemplate/AITemplate.py
index 74f4664..b2a76b4 100644
--- a/AITemplate/AITemplate.py
+++ b/AITemplate/AITemplate.py
@@ -411,7 +411,6 @@ class ControlNet(ControlBase):
                 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)
-                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)

But I'm pretty sure that's just a memory leak that papers over the performance problem, so I haven't made a PR for it.

prob dumb question How do I patch it for my comfy

just want to start using this flow but the only issue

Error occurred when executing KSampler:

Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\execution.py", line 151, in recursive_execute
output_data, output_ui = get_output_data(obj, input_data_all)
File "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\nodes.py", line 1211, in sample
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
File "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-Impact-Pack\modules\impact\hacky.py", line 22, in informative_sample
raise e
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-Impact-Pack\modules\impact\hacky.py", line 9, in informative_sample
return original_sample(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 676, in sample
samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 880, in sample_unipc
x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 730, in sample
model_prev_list = [self.model_fn(x, vec_t)]
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 421, in model_fn
return self.data_prediction_fn(x, t)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 403, in data_prediction_fn
noise = self.noise_prediction_fn(x, t)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 397, in noise_prediction_fn
return self.model(x, t)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 329, in model_fn
return noise_pred_fn(x, t_continuous)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\extra_samplers\uni_pc.py", line 297, in noise_pred_fn
output = model(x, t_input, **model_kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\external.py", line 98, in predict_eps_discrete_timestep
return (input - self(input, sigma, **kwargs)) / utils.append_dims(sigma, input.ndim)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in call_impl
return forward_call(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\external.py", line 151, in get_eps
return self.inner_model.apply_model(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\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 "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\controlnet.py", line 162, in get_control
control = self.control_model(x=x_noisy.to(self.control_model.dtype), hint=self.cond_hint, timesteps=t, context=context.to(self.control_model.dtype), y=y)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\cldm\cldm.py", line 283, in forward
emb = self.time_embed(t_emb)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\container.py", line 217, in forward
input = module(input)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\python_embeded\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "K:\AI ART UNIVERSE\COMFY-UI\ComfyUI_windows_portable\ComfyUI\comfy\ops.py", line 18, in forward
return torch.nn.functional.linear(input, self.weight, self.bias)

[asagi4]

@TheMindExpansionNetwork just updating AIT should get rid of that exception. My patch just deals with the slowness.

If you want to use it, you should be able to copy the patch as-is to a text file and use "git apply" in the AIT repository to apply it, but you should know that doing so means you'll have to deal with your local version of AIT being different from what's on github, and the patch may interfere with updating the nodes if a conflicting change is introduced.