From 8befcdd9d621c0919dcbb15abb847b5fcf0a5729 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Wed, 28 Feb 2024 19:05:08 +0100 Subject: [PATCH] PixArt controlnet test --- PixArt/conf.py | 13 ++ PixArt/loader.py | 34 +++- PixArt/models/PixArt.py | 1 + PixArt/models/PixArtMS.py | 26 ++- PixArt/models/pixart_controlnet.py | 312 +++++++++++++++++++++++++++++ PixArt/nodes.py | 53 ++++- 6 files changed, 426 insertions(+), 13 deletions(-) create mode 100644 PixArt/models/pixart_controlnet.py diff --git a/PixArt/conf.py b/PixArt/conf.py index 9919c5e..5ba98b0 100644 --- a/PixArt/conf.py +++ b/PixArt/conf.py @@ -38,6 +38,19 @@ pixart_conf = { }, } +pixart_conf.update({ # controlnet models + "ControlPixArtHalf": { + "target": "ControlPixArtHalf", + "unet_config": pixart_conf["PixArt_XL_2"]["unet_config"], + "sampling_settings": pixart_conf["PixArt_XL_2"]["sampling_settings"], + }, + "ControlPixArtMSHalf": { + "target": "ControlPixArtMSHalf", + "unet_config": pixart_conf["PixArtMS_XL_2"]["unet_config"], + "sampling_settings": pixart_conf["PixArtMS_XL_2"]["sampling_settings"], + } +}) + pixart_res = { "PixArtMS_XL_2": { # models/PixArtMS 1024x1024 '0.25': [512, 2048], '0.26': [512, 1984], '0.27': [512, 1920], '0.28': [512, 1856], diff --git a/PixArt/loader.py b/PixArt/loader.py index 1d66fa3..67edee2 100644 --- a/PixArt/loader.py +++ b/PixArt/loader.py @@ -3,6 +3,7 @@ import comfy.latent_formats import comfy.model_patcher import comfy.model_base import comfy.utils +import comfy.conds import torch from comfy import model_management from .diffusers_convert import convert_state_dict @@ -23,6 +24,27 @@ class EXM_PixArt(comfy.supported_models_base.BASE): def model_type(self, state_dict, prefix=""): return comfy.model_base.ModelType.EPS +class EXM_PixArt_Model(comfy.model_base.BaseModel): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def extra_conds(self, **kwargs): + out = super().extra_conds(**kwargs) + + img_hw = kwargs.get("img_hw", None) + if img_hw is not None: + out["img_hw"] = comfy.conds.CONDRegular(torch.tensor(img_hw)) + + aspect_ratio = kwargs.get("aspect_ratio", None) + if aspect_ratio is not None: + out["aspect_ratio"] = comfy.conds.CONDRegular(torch.tensor(aspect_ratio)) + + cn_hint = kwargs.get("cn_hint", None) + if cn_hint is not None: + out["cn_hint"] = comfy.conds.CONDRegular(cn_hint) + + return out + def load_pixart(model_path, model_conf): state_dict = comfy.utils.load_torch_file(model_path) state_dict = state_dict.get("model", state_dict) @@ -48,7 +70,7 @@ def load_pixart(model_path, model_conf): unet_dtype = manual_cast_dtype model_conf = EXM_PixArt(model_conf) # convert to object - model = comfy.model_base.BaseModel( + model = EXM_PixArt_Model( # same as comfy.model_base.BaseModel model_conf, model_type=comfy.model_base.ModelType.EPS, device=model_management.get_torch_device() @@ -60,6 +82,16 @@ def load_pixart(model_path, model_conf): elif model_conf.model_target == "PixArt": from .models.PixArt import PixArt model.diffusion_model = PixArt(**model_conf.unet_config) + elif model_conf.model_target == "ControlPixArtMSHalf": + from .models.PixArtMS import PixArtMS + from .models.pixart_controlnet import ControlPixArtMSHalf + model.diffusion_model = PixArtMS(**model_conf.unet_config) + model.diffusion_model = ControlPixArtMSHalf(model.diffusion_model) + elif model_conf.model_target == "ControlPixArtHalf": + from .models.PixArt import PixArt + from .models.pixart_controlnet import ControlPixArtHalf + model.diffusion_model = PixArt(**model_conf.unet_config) + model.diffusion_model = ControlPixArtHalf(model.diffusion_model) else: raise NotImplementedError(f"Unknown model target '{model_conf.model_target}'") diff --git a/PixArt/models/PixArt.py b/PixArt/models/PixArt.py index 4428ab7..13b6727 100644 --- a/PixArt/models/PixArt.py +++ b/PixArt/models/PixArt.py @@ -28,6 +28,7 @@ class PixArtBlock(nn.Module): def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., window_size=0, input_size=None, use_rel_pos=False, **block_kwargs): super().__init__() + self.hidden_size = hidden_size self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.attn = WindowAttention(hidden_size, num_heads=num_heads, qkv_bias=True, input_size=input_size if window_size == 0 else (window_size, window_size), diff --git a/PixArt/models/PixArtMS.py b/PixArt/models/PixArtMS.py index a9ef67b..66b3cc3 100644 --- a/PixArt/models/PixArtMS.py +++ b/PixArt/models/PixArtMS.py @@ -182,30 +182,34 @@ class PixArtMS(PixArt): x = self.unpatchify(x) # (N, out_channels, H, W) return x - def forward(self, x, timesteps, context, y=None, **kwargs): + def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, **kwargs): """ Forward pass that adapts comfy input to original forward function x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) timesteps: (N,) tensor of diffusion timesteps context: (N, 1, 120, C) conditioning - y: extra conditioning. + img_hw: height|width conditioning + aspect_ratio: aspect ratio conditioning """ - ## aspect ratio based on the latent image shape. - # Ideally, these should only be used as a fallback with the real ones - # passed in `y` to allow different values to be used for cont/uncond. + ## size/ar from cond with fallback based on the latent image shape. bs = x.shape[0] - data_info = { - "img_hw" : torch.tensor( + data_info = {} + if img_hw is None: + data_info["img_hw"] = torch.tensor( [[x.shape[2]*8, x.shape[3]*8]], dtype=self.dtype, device=x.device - ).repeat(bs, 1), - "aspect_ratio" : torch.tensor( + ).repeat(bs, 1) + else: + data_info["img_hw"] = img_hw.to(x.dtype).to(x.device) + if aspect_ratio is None or True: + data_info["aspect_ratio"] = torch.tensor( [[x.shape[2]/x.shape[3]]], dtype=self.dtype, device=x.device - ).repeat(bs, 1), - } + ).repeat(bs, 1) + else: + data_info["aspect_ratio"] = aspect_ratio.to(x.dtype).to(x.device) ## Still accepts the input w/o that dim but returns garbage if len(context.shape) == 3: diff --git a/PixArt/models/pixart_controlnet.py b/PixArt/models/pixart_controlnet.py new file mode 100644 index 0000000..37fa4c1 --- /dev/null +++ b/PixArt/models/pixart_controlnet.py @@ -0,0 +1,312 @@ +import re +import torch +import torch.nn as nn + +from copy import deepcopy +from torch import Tensor +from torch.nn import Module, Linear, init +from typing import Any, Mapping + +from .PixArt import PixArt, get_2d_sincos_pos_embed +from .PixArtMS import PixArtMSBlock, PixArtMS +from .utils import auto_grad_checkpoint + +# The implementation of ControlNet-Half architrecture +# https://github.com/lllyasviel/ControlNet/discussions/188 +class ControlT2IDitBlockHalf(Module): + def __init__(self, base_block: PixArtMSBlock, block_index: 0) -> None: + super().__init__() + self.copied_block = deepcopy(base_block) + self.block_index = block_index + + for p in self.copied_block.parameters(): + p.requires_grad_(True) + + self.copied_block.load_state_dict(base_block.state_dict()) + self.copied_block.train() + + self.hidden_size = hidden_size = base_block.hidden_size + if self.block_index == 0: + self.before_proj = Linear(hidden_size, hidden_size) + init.zeros_(self.before_proj.weight) + init.zeros_(self.before_proj.bias) + self.after_proj = Linear(hidden_size, hidden_size) + init.zeros_(self.after_proj.weight) + init.zeros_(self.after_proj.bias) + + def forward(self, x, y, t, mask=None, c=None): + + if self.block_index == 0: + # the first block + c = self.before_proj(c) + c = self.copied_block(x + c, y, t, mask) + c_skip = self.after_proj(c) + else: + # load from previous c and produce the c for skip connection + c = self.copied_block(c, y, t, mask) + c_skip = self.after_proj(c) + + return c, c_skip + + +# The implementation of ControlPixArtHalf net +class ControlPixArtHalf(Module): + # only support single res model + def __init__(self, base_model: PixArt, copy_blocks_num: int = 13) -> None: + super().__init__() + self.dtype = torch.get_default_dtype() + self.base_model = base_model.eval() + self.controlnet = [] + self.copy_blocks_num = copy_blocks_num + self.total_blocks_num = len(base_model.blocks) + for p in self.base_model.parameters(): + p.requires_grad_(False) + + # Copy first copy_blocks_num block + for i in range(copy_blocks_num): + self.controlnet.append(ControlT2IDitBlockHalf(base_model.blocks[i], i)) + self.controlnet = nn.ModuleList(self.controlnet) + + def __getattr__(self, name: str) -> Tensor or Module: + if name in ['forward', 'forward_with_dpmsolver', 'forward_with_cfg', 'forward_c', 'load_state_dict']: + return self.__dict__[name] + elif name in ['base_model', 'controlnet']: + return super().__getattr__(name) + else: + return getattr(self.base_model, name) + + def forward_c(self, c): + self.h, self.w = c.shape[-2]//self.patch_size, c.shape[-1]//self.patch_size + pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(self.pos_embed.shape[-1], (self.h, self.w), lewei_scale=self.lewei_scale, base_size=self.base_size)).unsqueeze(0).to(c.device).to(self.dtype) + return self.x_embedder(c) + pos_embed if c is not None else c + + # def forward(self, x, t, c, **kwargs): + # return self.base_model(x, t, c=self.forward_c(c), **kwargs) + def forward_raw(self, x, timestep, y, mask=None, data_info=None, c=None, **kwargs): + # modify the original PixArtMS forward function + if c is not None: + c = c.to(self.dtype) + c = self.forward_c(c) + """ + Forward pass of PixArt. + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t: (N,) tensor of diffusion timesteps + y: (N, 1, 120, C) tensor of class labels + """ + x = x.to(self.dtype) + timestep = timestep.to(self.dtype) + y = y.to(self.dtype) + pos_embed = self.pos_embed.to(self.dtype) + self.h, self.w = x.shape[-2]//self.patch_size, x.shape[-1]//self.patch_size + x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(timestep.to(x.dtype)) # (N, D) + t0 = self.t_block(t) + y = self.y_embedder(y, self.training) # (N, 1, L, D) + if mask is not None: + if mask.shape[0] != y.shape[0]: + mask = mask.repeat(y.shape[0] // mask.shape[0], 1) + mask = mask.squeeze(1).squeeze(1) + y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1]) + y_lens = mask.sum(dim=1).tolist() + else: + y_lens = [y.shape[2]] * y.shape[0] + y = y.squeeze(1).view(1, -1, x.shape[-1]) + + # define the first layer + x = auto_grad_checkpoint(self.base_model.blocks[0], x, y, t0, y_lens, **kwargs) # (N, T, D) #support grad checkpoint + + if c is not None: + # update c + for index in range(1, self.copy_blocks_num + 1): + c, c_skip = auto_grad_checkpoint(self.controlnet[index - 1], x, y, t0, y_lens, c, **kwargs) + x = auto_grad_checkpoint(self.base_model.blocks[index], x + c_skip, y, t0, y_lens, **kwargs) + + # update x + for index in range(self.copy_blocks_num + 1, self.total_blocks_num): + x = auto_grad_checkpoint(self.base_model.blocks[index], x, y, t0, y_lens, **kwargs) + else: + for index in range(1, self.total_blocks_num): + x = auto_grad_checkpoint(self.base_model.blocks[index], x, y, t0, y_lens, **kwargs) + + x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + return x + + def forward(self, x, timesteps, context, cn_hint=None, **kwargs): + """ + Forward pass that adapts comfy input to original forward function + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + timesteps: (N,) tensor of diffusion timesteps + context: (N, 1, 120, C) conditioning + cn_hint: controlnet hint + """ + ## Still accepts the input w/o that dim but returns garbage + if len(context.shape) == 3: + context = context.unsqueeze(1) + + ## run original forward pass + out = self.forward_raw( + x = x.to(self.dtype), + timestep = timesteps.to(self.dtype), + y = context.to(self.dtype), + c = cn_hint, + ) + + ## only return EPS + out = out.to(torch.float) + eps, rest = out[:, :self.in_channels], out[:, self.in_channels:] + return eps + + def forward_with_dpmsolver(self, x, t, y, data_info, c, **kwargs): + model_out = self.forward_raw(x, t, y, data_info=data_info, c=c, **kwargs) + return model_out.chunk(2, dim=1)[0] + + # def forward_with_dpmsolver(self, x, t, y, data_info, c, **kwargs): + # return self.base_model.forward_with_dpmsolver(x, t, y, data_info=data_info, c=self.forward_c(c), **kwargs) + + def forward_with_cfg(self, x, t, y, cfg_scale, data_info, c, **kwargs): + return self.base_model.forward_with_cfg(x, t, y, cfg_scale, data_info, c=self.forward_c(c), **kwargs) + + def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): + if all((k.startswith('base_model') or k.startswith('controlnet')) for k in state_dict.keys()): + return super().load_state_dict(state_dict, strict) + else: + new_key = {} + for k in state_dict.keys(): + new_key[k] = re.sub(r"(blocks\.\d+)(.*)", r"\1.base_block\2", k) + for k, v in new_key.items(): + if k != v: + print(f"replace {k} to {v}") + state_dict[v] = state_dict.pop(k) + + return self.base_model.load_state_dict(state_dict, strict) + + def unpatchify(self, x): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.out_channels + p = self.x_embedder.patch_size[0] + assert self.h * self.w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], self.h, self.w, p, p, c)) + x = torch.einsum('nhwpqc->nchpwq', x) + imgs = x.reshape(shape=(x.shape[0], c, self.h * p, self.w * p)) + return imgs + + # @property + # def dtype(self): + ## 返回模型参数的数据类型 + # return next(self.parameters()).dtype + + +# The implementation for PixArtMS_Half + 1024 resolution +class ControlPixArtMSHalf(ControlPixArtHalf): + # support multi-scale res model (multi-scale model can also be applied to single reso training & inference) + def __init__(self, base_model: PixArtMS, copy_blocks_num: int = 13) -> None: + super().__init__(base_model=base_model, copy_blocks_num=copy_blocks_num) + + def forward_raw(self, x, timestep, y, mask=None, data_info=None, c=None, **kwargs): + # modify the original PixArtMS forward function + """ + Forward pass of PixArt. + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t: (N,) tensor of diffusion timesteps + y: (N, 1, 120, C) tensor of class labels + """ + if c is not None: + c = c.to(self.dtype) + c = self.forward_c(c) + bs = x.shape[0] + x = x.to(self.dtype) + timestep = timestep.to(self.dtype) + y = y.to(self.dtype) + c_size, ar = data_info['img_hw'].to(self.dtype), data_info['aspect_ratio'].to(self.dtype) + self.h, self.w = x.shape[-2]//self.patch_size, x.shape[-1]//self.patch_size + + pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(self.pos_embed.shape[-1], (self.h, self.w), lewei_scale=self.lewei_scale, base_size=self.base_size)).unsqueeze(0).to(x.device).to(self.dtype) + x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(timestep) # (N, D) + csize = self.csize_embedder(c_size, bs) # (N, D) + ar = self.ar_embedder(ar, bs) # (N, D) + t = t + torch.cat([csize, ar], dim=1) + t0 = self.t_block(t) + y = self.y_embedder(y, self.training) # (N, D) + if mask is not None: + if mask.shape[0] != y.shape[0]: + mask = mask.repeat(y.shape[0] // mask.shape[0], 1) + mask = mask.squeeze(1).squeeze(1) + y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1]) + y_lens = mask.sum(dim=1).tolist() + else: + y_lens = [y.shape[2]] * y.shape[0] + y = y.squeeze(1).view(1, -1, x.shape[-1]) + + # define the first layer + x = auto_grad_checkpoint(self.base_model.blocks[0], x, y, t0, y_lens, **kwargs) # (N, T, D) #support grad checkpoint + + if c is not None: + # update c + for index in range(1, self.copy_blocks_num + 1): + c, c_skip = auto_grad_checkpoint(self.controlnet[index - 1], x, y, t0, y_lens, c, **kwargs) + x = auto_grad_checkpoint(self.base_model.blocks[index], x + c_skip, y, t0, y_lens, **kwargs) + + # update x + for index in range(self.copy_blocks_num + 1, self.total_blocks_num): + x = auto_grad_checkpoint(self.base_model.blocks[index], x, y, t0, y_lens, **kwargs) + else: + for index in range(1, self.total_blocks_num): + x = auto_grad_checkpoint(self.base_model.blocks[index], x, y, t0, y_lens, **kwargs) + + x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + return x + + def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, cn_hint=None, **kwargs): + """ + Forward pass that adapts comfy input to original forward function + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + timesteps: (N,) tensor of diffusion timesteps + context: (N, 1, 120, C) conditioning + img_hw: height|width conditioning + aspect_ratio: aspect ratio conditioning + cn_hint: controlnet hint + """ + ## size/ar from cond with fallback based on the latent image shape. + bs = x.shape[0] + data_info = {} + if img_hw is None: + data_info["img_hw"] = torch.tensor( + [[x.shape[2]*8, x.shape[3]*8]], + dtype=self.dtype, + device=x.device + ).repeat(bs, 1) + else: + data_info["img_hw"] = img_hw.to(x.dtype) + if aspect_ratio is None or True: + data_info["aspect_ratio"] = torch.tensor( + [[x.shape[2]/x.shape[3]]], + dtype=self.dtype, + device=x.device + ).repeat(bs, 1) + else: + data_info["aspect_ratio"] = aspect_ratio.to(x.dtype) + + ## Still accepts the input w/o that dim but returns garbage + if len(context.shape) == 3: + context = context.unsqueeze(1) + + ## run original forward pass + out = self.forward_raw( + x = x.to(self.dtype), + timestep = timesteps.to(self.dtype), + y = context.to(self.dtype), + c = cn_hint, + data_info=data_info, + ) + + ## only return EPS + out = out.to(torch.float) + eps, rest = out[:, :self.in_channels], out[:, self.in_channels:] + return eps diff --git a/PixArt/nodes.py b/PixArt/nodes.py index 184e6c3..08eeb17 100644 --- a/PixArt/nodes.py +++ b/PixArt/nodes.py @@ -38,7 +38,7 @@ class PixArtResolutionSelect(): def INPUT_TYPES(s): return { "required": { - "model": (list(pixart_conf.keys()),), + "model": (list(pixart_res.keys()),), # keys are the same for both "ratio": (list(pixart_res["PixArtMS_XL_2"].keys()),{"default":"1.00"}), } @@ -92,6 +92,55 @@ class PixArtLoraLoader: model_lora = load_pixart_lora(model, lora, lora_path, strength,) return (model_lora,) +class PixArtResolutionCond: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "cond": ("CONDITIONING", ), + "width": ("INT", {"default": 1024.0, "min": 0, "max": 8192}), + "height": ("INT", {"default": 1024.0, "min": 0, "max": 8192}), + } + } + + RETURN_TYPES = ("CONDITIONING",) + RETURN_NAMES = ("cond",) + FUNCTION = "add_cond" + CATEGORY = "ExtraModels/PixArt" + TITLE = "PixArt Resolution Conditioning" + + def add_cond(self, cond, width, height): + for c in range(len(cond)): + cond[c][1].update({ + "img_hw": [[height, width]], + "aspect_ratio": [[height/width]], + }) + return (cond,) + +class PixArtControlNetCond: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "cond": ("CONDITIONING",), + "latent": ("LATENT",), + # "image": ("IMAGE",), + # "vae": ("VAE",), + # "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}) + } + } + + RETURN_TYPES = ("CONDITIONING",) + RETURN_NAMES = ("cond",) + FUNCTION = "add_cond" + CATEGORY = "ExtraModels/PixArt" + TITLE = "PixArt ControlNet Conditioning" + + def add_cond(self, cond, latent): + for c in range(len(cond)): + cond[c][1]["cn_hint"] = latent["samples"] * 0.18215 + return (cond,) + class PixArtDPMSampler: """ The sampler from the reference code. @@ -189,4 +238,6 @@ NODE_CLASS_MAPPINGS = { "PixArtLoraLoader" : PixArtLoraLoader, "PixArtDPMSampler" : PixArtDPMSampler, "PixArtT5TextEncode" : PixArtT5TextEncode, + "PixArtResolutionCond" : PixArtResolutionCond, + "PixArtControlNetCond" : PixArtControlNetCond, }