|
|
@ -39,6 +39,7 @@ class ControlBase: |
|
|
|
self.global_average_pooling = False |
|
|
|
self.global_average_pooling = False |
|
|
|
self.timestep_range = None |
|
|
|
self.timestep_range = None |
|
|
|
self.compression_ratio = 8 |
|
|
|
self.compression_ratio = 8 |
|
|
|
|
|
|
|
self.upscale_algorithm = 'nearest-exact' |
|
|
|
|
|
|
|
|
|
|
|
if device is None: |
|
|
|
if device is None: |
|
|
|
device = comfy.model_management.get_torch_device() |
|
|
|
device = comfy.model_management.get_torch_device() |
|
|
@ -80,6 +81,7 @@ class ControlBase: |
|
|
|
c.timestep_percent_range = self.timestep_percent_range |
|
|
|
c.timestep_percent_range = self.timestep_percent_range |
|
|
|
c.global_average_pooling = self.global_average_pooling |
|
|
|
c.global_average_pooling = self.global_average_pooling |
|
|
|
c.compression_ratio = self.compression_ratio |
|
|
|
c.compression_ratio = self.compression_ratio |
|
|
|
|
|
|
|
c.upscale_algorithm = self.upscale_algorithm |
|
|
|
|
|
|
|
|
|
|
|
def inference_memory_requirements(self, dtype): |
|
|
|
def inference_memory_requirements(self, dtype): |
|
|
|
if self.previous_controlnet is not None: |
|
|
|
if self.previous_controlnet is not None: |
|
|
@ -165,7 +167,7 @@ class ControlNet(ControlBase): |
|
|
|
if self.cond_hint is not None: |
|
|
|
if self.cond_hint is not None: |
|
|
|
del self.cond_hint |
|
|
|
del self.cond_hint |
|
|
|
self.cond_hint = None |
|
|
|
self.cond_hint = None |
|
|
|
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio, 'nearest-exact', "center").to(dtype).to(self.device) |
|
|
|
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio, self.upscale_algorithm, "center").to(dtype).to(self.device) |
|
|
|
if x_noisy.shape[0] != self.cond_hint.shape[0]: |
|
|
|
if x_noisy.shape[0] != self.cond_hint.shape[0]: |
|
|
|
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) |
|
|
|
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) |
|
|
|
|
|
|
|
|
|
|
@ -435,12 +437,13 @@ def load_controlnet(ckpt_path, model=None): |
|
|
|
return control |
|
|
|
return control |
|
|
|
|
|
|
|
|
|
|
|
class T2IAdapter(ControlBase): |
|
|
|
class T2IAdapter(ControlBase): |
|
|
|
def __init__(self, t2i_model, channels_in, compression_ratio, device=None): |
|
|
|
def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None): |
|
|
|
super().__init__(device) |
|
|
|
super().__init__(device) |
|
|
|
self.t2i_model = t2i_model |
|
|
|
self.t2i_model = t2i_model |
|
|
|
self.channels_in = channels_in |
|
|
|
self.channels_in = channels_in |
|
|
|
self.control_input = None |
|
|
|
self.control_input = None |
|
|
|
self.compression_ratio = compression_ratio |
|
|
|
self.compression_ratio = compression_ratio |
|
|
|
|
|
|
|
self.upscale_algorithm = upscale_algorithm |
|
|
|
|
|
|
|
|
|
|
|
def scale_image_to(self, width, height): |
|
|
|
def scale_image_to(self, width, height): |
|
|
|
unshuffle_amount = self.t2i_model.unshuffle_amount |
|
|
|
unshuffle_amount = self.t2i_model.unshuffle_amount |
|
|
@ -466,7 +469,7 @@ class T2IAdapter(ControlBase): |
|
|
|
self.control_input = None |
|
|
|
self.control_input = None |
|
|
|
self.cond_hint = None |
|
|
|
self.cond_hint = None |
|
|
|
width, height = self.scale_image_to(x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio) |
|
|
|
width, height = self.scale_image_to(x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio) |
|
|
|
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, width, height, 'nearest-exact', "center").float().to(self.device) |
|
|
|
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, width, height, self.upscale_algorithm, "center").float().to(self.device) |
|
|
|
if self.channels_in == 1 and self.cond_hint.shape[1] > 1: |
|
|
|
if self.channels_in == 1 and self.cond_hint.shape[1] > 1: |
|
|
|
self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True) |
|
|
|
self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True) |
|
|
|
if x_noisy.shape[0] != self.cond_hint.shape[0]: |
|
|
|
if x_noisy.shape[0] != self.cond_hint.shape[0]: |
|
|
@ -485,12 +488,13 @@ class T2IAdapter(ControlBase): |
|
|
|
return self.control_merge(control_input, mid, control_prev, x_noisy.dtype) |
|
|
|
return self.control_merge(control_input, mid, control_prev, x_noisy.dtype) |
|
|
|
|
|
|
|
|
|
|
|
def copy(self): |
|
|
|
def copy(self): |
|
|
|
c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio) |
|
|
|
c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio, self.upscale_algorithm) |
|
|
|
self.copy_to(c) |
|
|
|
self.copy_to(c) |
|
|
|
return c |
|
|
|
return c |
|
|
|
|
|
|
|
|
|
|
|
def load_t2i_adapter(t2i_data): |
|
|
|
def load_t2i_adapter(t2i_data): |
|
|
|
compression_ratio = 8 |
|
|
|
compression_ratio = 8 |
|
|
|
|
|
|
|
upscale_algorithm = 'nearest-exact' |
|
|
|
|
|
|
|
|
|
|
|
if 'adapter' in t2i_data: |
|
|
|
if 'adapter' in t2i_data: |
|
|
|
t2i_data = t2i_data['adapter'] |
|
|
|
t2i_data = t2i_data['adapter'] |
|
|
@ -522,6 +526,7 @@ def load_t2i_adapter(t2i_data): |
|
|
|
elif "backbone.0.0.weight" in keys: |
|
|
|
elif "backbone.0.0.weight" in keys: |
|
|
|
model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.0.weight'].shape[1], proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63]) |
|
|
|
model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.0.weight'].shape[1], proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63]) |
|
|
|
compression_ratio = 32 |
|
|
|
compression_ratio = 32 |
|
|
|
|
|
|
|
upscale_algorithm = 'bilinear' |
|
|
|
else: |
|
|
|
else: |
|
|
|
return None |
|
|
|
return None |
|
|
|
|
|
|
|
|
|
|
@ -532,4 +537,4 @@ def load_t2i_adapter(t2i_data): |
|
|
|
if len(unexpected) > 0: |
|
|
|
if len(unexpected) > 0: |
|
|
|
print("t2i unexpected", unexpected) |
|
|
|
print("t2i unexpected", unexpected) |
|
|
|
|
|
|
|
|
|
|
|
return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio) |
|
|
|
return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio, upscale_algorithm) |
|
|
|