67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
import torch
|
|
from torch.nn import Linear
|
|
from types import MethodType
|
|
import comfy.model_management
|
|
import comfy.samplers
|
|
from comfy.cldm.cldm import ControlNet
|
|
from comfy.controlnet import ControlLora
|
|
|
|
def patch_controlnet(model, control_net):
|
|
import comfy.controlnet
|
|
if isinstance(control_net, ControlLora):
|
|
del_keys = []
|
|
for k in control_net.control_weights:
|
|
if k.startswith("label_emb.0.0."):
|
|
del_keys.append(k)
|
|
|
|
for k in del_keys:
|
|
control_net.control_weights.pop(k)
|
|
|
|
super_pre_run = ControlLora.pre_run
|
|
super_copy = ControlLora.copy
|
|
|
|
super_forward = ControlNet.forward
|
|
|
|
def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
|
|
with torch.cuda.amp.autocast(enabled=True):
|
|
context = model.model.diffusion_model.encoder_hid_proj(context)
|
|
return super_forward(self, x, hint, timesteps, context, **kwargs)
|
|
|
|
def KolorsControlLora_pre_run(self, *args, **kwargs):
|
|
result = super_pre_run(self, *args, **kwargs)
|
|
|
|
if hasattr(self, "control_model"):
|
|
self.control_model.forward = MethodType(
|
|
KolorsControlNet_forward, self.control_model)
|
|
return result
|
|
|
|
control_net.pre_run = MethodType(
|
|
KolorsControlLora_pre_run, control_net)
|
|
|
|
def KolorsControlLora_copy(self, *args, **kwargs):
|
|
c = super_copy(self, *args, **kwargs)
|
|
c.pre_run = MethodType(
|
|
KolorsControlLora_pre_run, c)
|
|
return c
|
|
|
|
control_net.copy = MethodType(KolorsControlLora_copy, control_net)
|
|
|
|
elif isinstance(control_net, comfy.controlnet.ControlNet):
|
|
model_label_emb = model.model.diffusion_model.label_emb
|
|
control_net.control_model.label_emb = model_label_emb
|
|
control_net.control_model_wrapped.model.label_emb = model_label_emb
|
|
super_forward = ControlNet.forward
|
|
|
|
def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
|
|
with torch.cuda.amp.autocast(enabled=True):
|
|
context = model.model.diffusion_model.encoder_hid_proj(context)
|
|
return super_forward(self, x, hint, timesteps, context, **kwargs)
|
|
|
|
control_net.control_model.forward = MethodType(
|
|
KolorsControlNet_forward, control_net.control_model)
|
|
|
|
else:
|
|
raise NotImplementedError(f"Type {control_net} not supported for KolorsControlNetPatch")
|
|
|
|
return control_net
|