From a41d6868cb796a591ccc267e3c90c3f21fdba927 Mon Sep 17 00:00:00 2001 From: Bluear7878 Date: Fri, 30 May 2025 16:08:02 +0900 Subject: [PATCH 01/16] feat: support IP-adapter --- examples/flux.1-dev-IP-adapter.py | 33 ++ nunchaku/csrc/flux.h | 52 ++ nunchaku/csrc/pybind.cpp | 12 + nunchaku/models/IP_adapter/__init__.py | 0 .../IP_adapter/diffusers_adapters/__init__.py | 12 + .../IP_adapter/diffusers_adapters/flux.py | 35 ++ nunchaku/models/IP_adapter/utils.py | 515 ++++++++++++++++++ src/FluxModel.cpp | 490 +++++++++++++++++ src/FluxModel.h | 18 + tests/flux/test_flux_dev_IPA.py | 56 ++ 10 files changed, 1223 insertions(+) create mode 100644 examples/flux.1-dev-IP-adapter.py create mode 100644 nunchaku/models/IP_adapter/__init__.py create mode 100644 nunchaku/models/IP_adapter/diffusers_adapters/__init__.py create mode 100644 nunchaku/models/IP_adapter/diffusers_adapters/flux.py create mode 100644 nunchaku/models/IP_adapter/utils.py create mode 100644 tests/flux/test_flux_dev_IPA.py diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py new file mode 100644 index 00000000..d92b504e --- /dev/null +++ b/examples/flux.1-dev-IP-adapter.py @@ -0,0 +1,33 @@ +import torch +from diffusers import FluxPipeline +from diffusers.utils import load_image + +from nunchaku import NunchakuFluxTransformer2dModel +from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe +from nunchaku.utils import get_precision + +precision = get_precision() +transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-{precision}-flux.1-dev") + +pipeline = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 +).to("cuda") + +pipeline.load_ip_adapter( + pretrained_model_name_or_path_or_dict="XLabs-AI/flux-ip-adapter-v2", + weight_name="ip_adapter.safetensors", + image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", +) + +apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") + +IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") + +image = pipeline( + prompt="A woman holding a sign that says 'SVDQuant is fast!", + ip_adapter_image=IP_image.convert("RGB"), + num_inference_steps=50, + generator=torch.Generator("cuda"), +).images[0] + +image.save(f"flux.1-dev-IP-adapter-{precision}.png") diff --git a/nunchaku/csrc/flux.h b/nunchaku/csrc/flux.h index 135e25f5..6dfeff12 100644 --- a/nunchaku/csrc/flux.h +++ b/nunchaku/csrc/flux.h @@ -212,4 +212,56 @@ class QuantizedFluxModel : public ModuleWrapper { // : public torch:: throw std::invalid_argument(spdlog::fmt_lib::format("Invalid attention implementation {}", name)); } } + + std::tuple + forward_layer_ip_adapter(int64_t idx, + torch::Tensor hidden_states, + torch::Tensor encoder_hidden_states, + torch::Tensor temb, + torch::Tensor rotary_emb_img, + torch::Tensor rotary_emb_context, + torch::Tensor k_img, + torch::Tensor v_img, + std::optional controlnet_block_samples = std::nullopt, + std::optional controlnet_single_block_samples = std::nullopt) { + CUDADeviceContext ctx(deviceId); + + spdlog::debug("QuantizedFluxModel forward_layer {}", idx); + + hidden_states = hidden_states.contiguous(); + encoder_hidden_states = encoder_hidden_states.contiguous(); + temb = temb.contiguous(); + rotary_emb_img = rotary_emb_img.contiguous(); + rotary_emb_context = rotary_emb_context.contiguous(); + k_img = k_img.contiguous(); + v_img = v_img.contiguous(); + + auto &&[hidden_states_, encoder_hidden_states_, ip_query_] = net->forward_ip_adapter( + idx, + from_torch(hidden_states), + from_torch(encoder_hidden_states), + from_torch(temb), + from_torch(rotary_emb_img), + from_torch(rotary_emb_context), + from_torch(k_img), + from_torch(v_img), + controlnet_block_samples.has_value() ? from_torch(controlnet_block_samples.value().contiguous()) : Tensor{}, + controlnet_single_block_samples.has_value() + ? from_torch(controlnet_single_block_samples.value().contiguous()) + : Tensor{}); + /* + auto ip_attn_output_contig = + ip_attn_output_.is_contiguous() + ? ip_attn_output_ + : ip_attn_output_.copy(ip_attn_output_.device()); + */ + + hidden_states = to_torch(hidden_states_); + encoder_hidden_states = to_torch(encoder_hidden_states_); + // torch::Tensor ip_attn_output = to_torch(ip_attn_output_contig); + torch::Tensor ip_query = to_torch(ip_query_); + Tensor::synchronizeDevice(); + + return {hidden_states, encoder_hidden_states, ip_query}; + } }; diff --git a/nunchaku/csrc/pybind.cpp b/nunchaku/csrc/pybind.cpp index ee9cc965..7195b21c 100644 --- a/nunchaku/csrc/pybind.cpp +++ b/nunchaku/csrc/pybind.cpp @@ -49,6 +49,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("rotary_emb_context"), py::arg("controlnet_block_samples") = py::none(), py::arg("controlnet_single_block_samples") = py::none()) + .def("forward_layer_ip_adapter", + &QuantizedFluxModel::forward_layer_ip_adapter, + py::arg("idx"), + py::arg("hidden_states"), + py::arg("encoder_hidden_states"), + py::arg("temb"), + py::arg("rotary_emb_img"), + py::arg("rotary_emb_context"), + py::arg("k_img"), + py::arg("v_img"), + py::arg("controlnet_block_samples") = py::none(), + py::arg("controlnet_single_block_samples") = py::none()) .def("forward_single_layer", &QuantizedFluxModel::forward_single_layer) .def("norm_one_forward", &QuantizedFluxModel::norm_one_forward) .def("startDebug", &QuantizedFluxModel::startDebug) diff --git a/nunchaku/models/IP_adapter/__init__.py b/nunchaku/models/IP_adapter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py b/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py new file mode 100644 index 00000000..2064cd48 --- /dev/null +++ b/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py @@ -0,0 +1,12 @@ +from diffusers import DiffusionPipeline + + +def apply_IPA_on_pipe(pipe: DiffusionPipeline, *args, **kwargs): + assert isinstance(pipe, DiffusionPipeline) + + pipe_cls_name = pipe.__class__.__name__ + if pipe_cls_name.startswith("Flux") or pipe_cls_name.startswith("PuLID"): + from .flux import apply_IPA_on_pipe as apply_IPA_on_pipe_fn + else: + raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}") + return apply_IPA_on_pipe_fn(pipe, *args, **kwargs) diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py new file mode 100644 index 00000000..58319aab --- /dev/null +++ b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py @@ -0,0 +1,35 @@ +from diffusers import DiffusionPipeline, FluxTransformer2DModel +from torch import nn + +from ...IP_adapter import utils + + +def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_scale: float = 1.0, repo_id: str): + + IPA_transformer_blocks = nn.ModuleList( + [ + utils.IPA_TransformerBlocks( + transformer=transformer, + ip_adapter_scale=ip_adapter_scale, + return_hidden_states_first=False, + device=transformer.device, + ) + ] + ) + dummy_single_transformer_blocks = nn.ModuleList() + + IPA_transformer_blocks[0].load_ip_adapter_weights_per_layer(repo_id=repo_id) + + transformer.transformer_blocks = IPA_transformer_blocks + transformer.single_transformer_blocks = dummy_single_transformer_blocks + + transformer._is_IPA = True + + return transformer + + +def apply_IPA_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs): + if not shallow_patch: + apply_IPA_on_transformer(pipe.transformer, **kwargs) + + return pipe diff --git a/nunchaku/models/IP_adapter/utils.py b/nunchaku/models/IP_adapter/utils.py new file mode 100644 index 00000000..daa02319 --- /dev/null +++ b/nunchaku/models/IP_adapter/utils.py @@ -0,0 +1,515 @@ +import math + +import cv2 +import insightface +import numpy as np +import torch +import torch.nn.functional as F +from facexlib.parsing import init_parsing_model +from facexlib.utils.face_restoration_helper import FaceRestoreHelper +from huggingface_hub import hf_hub_download +from insightface.app import FaceAnalysis +from safetensors import safe_open +from torch import nn +from torchvision.transforms import InterpolationMode +from torchvision.transforms.functional import normalize, resize +from torchvision.utils import make_grid + +from nunchaku.models.pulid.encoders_transformer import IDFormer +from nunchaku.models.pulid.eva_clip import create_model_and_transforms +from nunchaku.models.pulid.eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD +from nunchaku.models.transformers.utils import pad_tensor + +num_transformer_blocks = 19 # FIXME +num_single_transformer_blocks = 38 # FIXME + + +class IPA_TransformerBlocks(nn.Module): + def __init__( + self, + *, + transformer: nn.Module = None, + ip_adapter_scale: float = 1.0, + return_hidden_states_first: bool = True, + return_hidden_states_only: bool = False, + verbose: bool = False, + device: str | torch.device, + ): + super().__init__() + self.transformer = transformer + self.transformer_blocks = transformer.transformer_blocks + self.return_hidden_states_first = return_hidden_states_first + self.return_hidden_states_only = return_hidden_states_only + self.verbose = verbose + self.ip_adapter_scale = ip_adapter_scale + + self.m = self.transformer_blocks[0].m + self.dtype = torch.bfloat16 if self.m.isBF16() else torch.float16 + self.device = device + + @staticmethod + def pack_rotemb(rotemb: torch.Tensor) -> torch.Tensor: + assert rotemb.dtype == torch.float32 + B = rotemb.shape[0] + M = rotemb.shape[1] + D = rotemb.shape[2] * 2 + msg_shape = "rotemb shape must be (B, M, D//2, 1, 2)" + assert rotemb.shape == (B, M, D // 2, 1, 2), msg_shape + assert M % 16 == 0 + assert D % 8 == 0 + rotemb = rotemb.reshape(B, M // 16, 16, D // 8, 8) + rotemb = rotemb.permute(0, 1, 3, 2, 4) + # 16*8 pack, FP32 accumulator (C) format + # https://docs.nvidia.com/cuda/parallel-thread-execution/#mma-16816-c + rotemb = rotemb.reshape(*rotemb.shape[0:3], 2, 8, 4, 2) + rotemb = rotemb.permute(0, 1, 2, 4, 5, 3, 6) + rotemb = rotemb.contiguous() + rotemb = rotemb.view(B, M, D) + return rotemb + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + image_rotary_emb: torch.Tensor, + id_embeddings=None, + id_weight=None, + joint_attention_kwargs=None, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + skip_first_layer=False, + ): + # batch_size = hidden_states.shape[0] + txt_tokens = encoder_hidden_states.shape[1] + img_tokens = hidden_states.shape[1] + + self.id_embeddings = id_embeddings + self.id_weight = id_weight + self.pulid_ca_idx = 0 + if self.id_embeddings is not None: + self.set_residual_callback() + + original_dtype = hidden_states.dtype + original_device = hidden_states.device + + hidden_states = hidden_states.to(self.dtype).to(self.device) + encoder_hidden_states = encoder_hidden_states.to(self.dtype).to(self.device) + temb = temb.to(self.dtype).to(self.device) + image_rotary_emb = image_rotary_emb.to(self.device) + + if controlnet_block_samples is not None: + if len(controlnet_block_samples) > 0: + controlnet_block_samples = torch.stack(controlnet_block_samples).to(self.device) + else: + controlnet_block_samples = None + + if controlnet_single_block_samples is not None: + if len(controlnet_single_block_samples) > 0: + controlnet_single_block_samples = torch.stack(controlnet_single_block_samples).to(self.device) + else: + controlnet_single_block_samples = None + + assert image_rotary_emb.ndim == 6 + assert image_rotary_emb.shape[0] == 1 + assert image_rotary_emb.shape[1] == 1 + assert image_rotary_emb.shape[2] == 1 * (txt_tokens + img_tokens) + # [1, tokens, head_dim / 2, 1, 2] (sincos) + image_rotary_emb = image_rotary_emb.reshape([1, txt_tokens + img_tokens, *image_rotary_emb.shape[3:]]) + rotary_emb_txt = image_rotary_emb[:, :txt_tokens, ...] # .to(self.dtype) + rotary_emb_img = image_rotary_emb[:, txt_tokens:, ...] # .to(self.dtype) + rotary_emb_single = image_rotary_emb # .to(self.dtype) + + rotary_emb_txt = self.pack_rotemb(pad_tensor(rotary_emb_txt, 256, 1)) + rotary_emb_img = self.pack_rotemb(pad_tensor(rotary_emb_img, 256, 1)) + rotary_emb_single = self.pack_rotemb(pad_tensor(rotary_emb_single, 256, 1)) + + if joint_attention_kwargs is not None and "ip_hidden_states" in joint_attention_kwargs: + ip_hidden_states = joint_attention_kwargs.pop("ip_hidden_states") + + remaining_kwargs = { + "temb": temb, + "rotary_emb_img": rotary_emb_img, + "rotary_emb_txt": rotary_emb_txt, + "rotary_emb_single": rotary_emb_single, + "controlnet_block_samples": controlnet_block_samples, + "controlnet_single_block_samples": controlnet_single_block_samples, + "txt_tokens": txt_tokens, + "ip_hidden_states": ip_hidden_states if ip_hidden_states is not None else None, + } + + torch._dynamo.graph_break() + updated_h, updated_enc, _, _ = self.call_multi_transformer_blocks( + hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, **remaining_kwargs + ) + + cat_hidden_states = torch.cat([updated_enc, updated_h], dim=1) + + updated_cat, _ = self.call_single_transformer_blocks( + hidden_states=cat_hidden_states, encoder_hidden_states=None, **remaining_kwargs + ) + # torch._dynamo.graph_break() + + final_enc = updated_cat[:, :txt_tokens, ...] + final_h = updated_cat[:, txt_tokens:, ...] + + final_h = final_h.to(original_dtype).to(original_device) + final_enc = final_enc.to(original_dtype).to(original_device) + + if self.return_hidden_states_only: + return final_h + if self.return_hidden_states_first: + return final_h, final_enc + return final_enc, final_h + + def call_multi_transformer_blocks( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + rotary_emb_img: torch.Tensor, + rotary_emb_txt: torch.Tensor, + rotary_emb_single: torch.Tensor, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + skip_first_layer=False, + txt_tokens=None, + ip_hidden_states=None, + ): + start_idx = 0 + original_hidden_states = hidden_states.clone() + original_encoder_hidden_states = encoder_hidden_states.clone() + + for idx in range(start_idx, num_transformer_blocks): + k_img = self.ip_k_projs[idx](ip_hidden_states[0]) + v_img = self.ip_v_projs[idx](ip_hidden_states[0]) + + hidden_states, encoder_hidden_states, ip_query = self.m.forward_layer_ip_adapter( + idx, + hidden_states, + encoder_hidden_states, + temb, + rotary_emb_img, + rotary_emb_txt, + k_img.squeeze(dim=1), + v_img.squeeze(dim=1), + controlnet_block_samples, + controlnet_single_block_samples, + ) + ip_query = ip_query.contiguous() + ip_query = ip_query.view(1, -1, 24, 128).transpose(1, 2) + + k_img = k_img.view(1, -1, 24, 128).transpose(1, 2) + v_img = v_img.view(1, -1, 24, 128).transpose(1, 2) + + real_ip_attn_output = F.scaled_dot_product_attention( + ip_query, k_img, v_img, attn_mask=None, dropout_p=0.0, is_causal=False + ) + real_ip_attn_output = real_ip_attn_output.transpose(1, 2).reshape(1, -1, 24 * 128) + + hidden_states = hidden_states + self.ip_adapter_scale * real_ip_attn_output + + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + hs_res = hidden_states - original_hidden_states + enc_res = encoder_hidden_states - original_encoder_hidden_states + return hidden_states, encoder_hidden_states, hs_res, enc_res + + def call_single_transformer_blocks( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + rotary_emb_img: torch.Tensor, + rotary_emb_txt: torch.Tensor, + rotary_emb_single: torch.Tensor, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + skip_first_layer=False, + txt_tokens=None, + ip_hidden_states=None, + ): + start_idx = 0 + original_hidden_states = hidden_states.clone() + + for idx in range(start_idx, num_single_transformer_blocks): + hidden_states = self.m.forward_single_layer( + idx, + hidden_states, + temb, + rotary_emb_single, + ) + + hidden_states = hidden_states.contiguous() + hs_res = hidden_states - original_hidden_states + return hidden_states, hs_res + + def load_ip_adapter_weights_per_layer( + self, + repo_id: str, + filename: str = "ip_adapter.safetensors", + prefix: str = "double_blocks.", + joint_attention_dim: int = 4096, + inner_dim: int = 3072, + ): + path = hf_hub_download(repo_id=repo_id, filename=filename) + raw_cpu = {} + with safe_open(path, framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith(prefix): + raw_cpu[key] = f.get_tensor(key) + + raw = {k: v.to(self.device) for k, v in raw_cpu.items()} + layer_ids = sorted({int(k.split(".")[1]) for k in raw.keys()}) + layers = [] + for i in layer_ids: + base = f"double_blocks.{i}.processor.ip_adapter_double_stream" + layers.append( + { + "k_weight": raw[f"{base}_k_proj.weight"], + "k_bias": raw[f"{base}_k_proj.bias"], + "v_weight": raw[f"{base}_v_proj.weight"], + "v_bias": raw[f"{base}_v_proj.bias"], + } + ) + + cross_dim = joint_attention_dim + hidden_dim = inner_dim + self.ip_k_projs = nn.ModuleList() + self.ip_v_projs = nn.ModuleList() + + for layer in layers: + k_proj = nn.Linear(cross_dim, hidden_dim, bias=True, device=self.device, dtype=self.dtype) + v_proj = nn.Linear(cross_dim, hidden_dim, bias=True, device=self.device, dtype=self.dtype) + + k_proj.weight.data.copy_(layer["k_weight"]) + k_proj.bias.data.copy_(layer["k_bias"]) + v_proj.weight.data.copy_(layer["v_weight"]) + v_proj.bias.data.copy_(layer["v_bias"]) + + self.ip_k_projs.append(k_proj) + self.ip_v_projs.append(v_proj) + + +def resize_numpy_image_long(image, resize_long_edge=768): + h, w = image.shape[:2] + if max(h, w) <= resize_long_edge: + return image + k = resize_long_edge / max(h, w) + h = int(h * k) + w = int(w * k) + image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LANCZOS4) + return image + + +# from basicsr +def img2tensor(imgs, bgr2rgb=True, float32=True): + """Numpy array to tensor. + + Args: + imgs (list[ndarray] | ndarray): Input images. + bgr2rgb (bool): Whether to change bgr to rgb. + float32 (bool): Whether to change to float32. + + Returns: + list[tensor] | tensor: Tensor images. If returned results only have + one element, just return tensor. + """ + + def _totensor(img, bgr2rgb, float32): + if img.shape[2] == 3 and bgr2rgb: + if img.dtype == "float64": + img = img.astype("float32") + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = torch.from_numpy(img.transpose(2, 0, 1)) + if float32: + img = img.float() + return img + + if isinstance(imgs, list): + return [_totensor(img, bgr2rgb, float32) for img in imgs] + return _totensor(imgs, bgr2rgb, float32) + + +def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): + """Convert torch Tensors into image numpy arrays. + + After clamping to [min, max], values will be normalized to [0, 1]. + + Args: + tensor (Tensor or list[Tensor]): Accept shapes: + 1) 4D mini-batch Tensor of shape (B x 3/1 x H x W); + 2) 3D Tensor of shape (3/1 x H x W); + 3) 2D Tensor of shape (H x W). + Tensor channel should be in RGB order. + rgb2bgr (bool): Whether to change rgb to bgr. + out_type (numpy type): output types. If ``np.uint8``, transform outputs + to uint8 type with range [0, 255]; otherwise, float type with + range [0, 1]. Default: ``np.uint8``. + min_max (tuple[int]): min and max values for clamp. + + Returns: + (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of + shape (H x W). The channel order is BGR. + """ + if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): + raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}") + + if torch.is_tensor(tensor): + tensor = [tensor] + result = [] + for _tensor in tensor: + _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) + _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) + + n_dim = _tensor.dim() + if n_dim == 4: + img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy() + img_np = img_np.transpose(1, 2, 0) + if rgb2bgr: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + elif n_dim == 3: + img_np = _tensor.numpy() + img_np = img_np.transpose(1, 2, 0) + if img_np.shape[2] == 1: # gray image + img_np = np.squeeze(img_np, axis=2) + else: + if rgb2bgr: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + elif n_dim == 2: + img_np = _tensor.numpy() + else: + raise TypeError(f"Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}") + if out_type == np.uint8: + # Unlike MATLAB, numpy.unit8() WILL NOT round by default. + img_np = (img_np * 255.0).round() + img_np = img_np.astype(out_type) + result.append(img_np) + if len(result) == 1: + result = result[0] + return result + + +def to_gray(img): + x = 0.299 * img[:, 0:1] + 0.587 * img[:, 1:2] + 0.114 * img[:, 2:3] + x = x.repeat(1, 3, 1, 1) + return x + + +def get_puild_embed(image, device, cal_uncond=False, weight_dtype=torch.bfloat16, onnx_provider="gpu"): + face_helper = FaceRestoreHelper( + upscale_factor=1, + face_size=512, + crop_ratio=(1, 1), + det_model="retinaface_resnet50", + save_ext="png", + device=device, + ) + face_helper.face_parse = None + face_helper.face_parse = init_parsing_model(model_name="bisenet", device=device) + + providers = ( + ["CPUExecutionProvider"] if onnx_provider == "cpu" else ["CUDAExecutionProvider", "CPUExecutionProvider"] + ) + + app = FaceAnalysis(name="antelopev2", root=".", providers=providers) + app.prepare(ctx_id=0, det_size=(640, 640)) + handler_ante = insightface.model_zoo.get_model("models/antelopev2/glintr100.onnx", providers=providers) + handler_ante.prepare(ctx_id=0) + + from safetensors.torch import load_file + + ckpt_path = hf_hub_download("guozinan/PuLID", "pulid_flux_v0.9.0.safetensors", local_dir="models") + + full_sd = load_file(ckpt_path) + + encoder_sd = {k.split("pulid_encoder.")[1]: v for k, v in full_sd.items() if k.startswith("pulid_encoder.")} + + pulid_encoder = IDFormer().to(device, weight_dtype) + pulid_encoder.load_state_dict(encoder_sd, strict=True) + pulid_encoder.eval() + + model, _, _ = create_model_and_transforms("EVA02-CLIP-L-14-336", "eva_clip", force_custom_clip=True) + model = model.visual + clip_vision_model = model.to(device, dtype=weight_dtype) + + face_helper.clean_all() + debug_img_list = [] + image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + # get antelopev2 embedding + face_info = app.get(image_bgr) + if len(face_info) > 0: + face_info = sorted(face_info, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[ + -1 + ] # only use the maximum face + id_ante_embedding = face_info["embedding"] + debug_img_list.append( + image[ + int(face_info["bbox"][1]) : int(face_info["bbox"][3]), + int(face_info["bbox"][0]) : int(face_info["bbox"][2]), + ] + ) + else: + id_ante_embedding = None + + # using facexlib to detect and align face + face_helper.read_image(image_bgr) + face_helper.get_face_landmarks_5(only_center_face=True) + face_helper.align_warp_face() + if len(face_helper.cropped_faces) == 0: + raise RuntimeError("facexlib align face fail") + align_face = face_helper.cropped_faces[0] + # incase insightface didn't detect face + if id_ante_embedding is None: + print("fail to detect face using insightface, extract embedding on align face") + id_ante_embedding = handler_ante.get_feat(align_face) + + id_ante_embedding = torch.from_numpy(id_ante_embedding).to(device, weight_dtype) + if id_ante_embedding.ndim == 1: + id_ante_embedding = id_ante_embedding.unsqueeze(0) + + # parsing + input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0 + input = input.to(device) + parsing_out = face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0] + parsing_out = parsing_out.argmax(dim=1, keepdim=True) + bg_label = [0, 16, 18, 7, 8, 9, 14, 15] + bg = sum(parsing_out == i for i in bg_label).bool() + white_image = torch.ones_like(input) + # only keep the face features + face_features_image = torch.where(bg, white_image, to_gray(input)) + debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False)) + + eva_transform_mean = getattr(clip_vision_model, "image_mean", OPENAI_DATASET_MEAN) + eva_transform_std = getattr(clip_vision_model, "image_std", OPENAI_DATASET_STD) + if not isinstance(eva_transform_mean, (list, tuple)): + eva_transform_mean = (eva_transform_mean,) * 3 + if not isinstance(eva_transform_std, (list, tuple)): + eva_transform_std = (eva_transform_std,) * 3 + eva_transform_mean = eva_transform_mean + eva_transform_std = eva_transform_std + + # transform img before sending to eva-clip-vit + face_features_image = resize(face_features_image, clip_vision_model.image_size, InterpolationMode.BICUBIC) + face_features_image = normalize(face_features_image, eva_transform_mean, eva_transform_std) + id_cond_vit, id_vit_hidden = clip_vision_model( + face_features_image.to(weight_dtype), return_all_features=False, return_hidden=True, shuffle=False + ) + id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True) + id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm) + + id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1) + + id_embedding = pulid_encoder(id_cond, id_vit_hidden) + + if not cal_uncond: + return id_embedding, None + + id_uncond = torch.zeros_like(id_cond) + id_vit_hidden_uncond = [] + for layer_idx in range(0, len(id_vit_hidden)): + id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx])) + uncond_id_embedding = pulid_encoder(id_uncond, id_vit_hidden_uncond) + + return id_embedding, uncond_id_embedding diff --git a/src/FluxModel.cpp b/src/FluxModel.cpp index 46d54aaa..c57da8ae 100644 --- a/src/FluxModel.cpp +++ b/src/FluxModel.cpp @@ -776,6 +776,390 @@ std::tuple JointTransformerBlock::forward(Tensor hidden_states, return {hidden_states, encoder_hidden_states}; } +std::tuple JointTransformerBlock::forward_ip_adapter_branch(Tensor hidden_states, + Tensor encoder_hidden_states, + Tensor temb, + Tensor rotary_emb, + Tensor rotary_emb_context, + float sparsityRatio) { + int batch_size = hidden_states.shape[0]; + assert(encoder_hidden_states.shape[0] == batch_size); + + nvtxRangePushA("JointTransformerBlock"); + + nvtxRangePushA("AdaNorm"); + + int num_tokens_img = hidden_states.shape[1]; + int num_tokens_txt = encoder_hidden_states.shape[1]; + + assert(hidden_states.shape[2] == dim); + assert(encoder_hidden_states.shape[2] == dim); + + Tensor q_heads; + + auto make_contiguous = [&](const Tensor &t) { + int B = t.shape.dataExtent[0]; + int R = t.shape.dataExtent[1]; + int C = t.shape.dataExtent[2]; + size_t E = t.scalar_size(); + size_t src_pitch = t.stride(1) * E; + + size_t dst_pitch = C * E; + size_t width = C * E; + size_t height = R; + + Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device()); + + auto stream = getCurrentCUDAStream(); + for (int b = 0; b < B; ++b) { + const void *src = (const char *)t.data_ptr() + t.stride(0) * b * E; + void *dst = (char *)out.data_ptr() + out.stride(0) * b * E; + checkCUDA( + cudaMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, cudaMemcpyDeviceToDevice, stream)); + } + return out; + }; + + spdlog::debug("hidden_states={} encoder_hidden_states={} temb={}", + hidden_states.shape.str(), + encoder_hidden_states.shape.str(), + temb.shape.str()); + spdlog::debug("batch_size={} num_tokens_img={} num_tokens_txt={}", batch_size, num_tokens_img, num_tokens_txt); + + auto norm1_output = norm1.forward(hidden_states, temb); + auto norm1_context_output = norm1_context.forward(encoder_hidden_states, temb); + +#if 0 + norm1_output.x = hidden_states; + norm1_context_output.x = encoder_hidden_states; +#endif + + debug("norm_hidden_states", norm1_output.x); + debug("norm_encoder_hidden_states", norm1_context_output.x); + + constexpr int POOL_SIZE = Attention::POOL_SIZE; + + nvtxRangePop(); + + auto stream = getCurrentCUDAStream(); + + int num_tokens_img_pad = 0, num_tokens_txt_pad = 0; + Tensor raw_attn_output; + + if (attnImpl == AttentionImpl::FlashAttention2) { + num_tokens_img_pad = num_tokens_img; + num_tokens_txt_pad = num_tokens_txt; + + Tensor concat; + Tensor pool; + + { + nvtxRangePushA("qkv_proj"); + + const bool blockSparse = sparsityRatio > 0; + + const int poolTokens = num_tokens_img / POOL_SIZE + num_tokens_txt / POOL_SIZE; + concat = Tensor::allocate({batch_size, num_tokens_img + num_tokens_txt, dim * 3}, + norm1_output.x.scalar_type(), + norm1_output.x.device()); + + pool = blockSparse ? Tensor::allocate({batch_size, poolTokens, dim * 3}, + norm1_output.x.scalar_type(), + norm1_output.x.device()) + : Tensor{}; + + for (int i = 0; i < batch_size; i++) { + // img first + Tensor qkv = concat.slice(0, i, i + 1).slice(1, 0, num_tokens_img); + Tensor qkv_context = + concat.slice(0, i, i + 1).slice(1, num_tokens_img, num_tokens_img + num_tokens_txt); + + Tensor pool_qkv = + pool.valid() ? pool.slice(0, i, i + 1).slice(1, 0, num_tokens_img / POOL_SIZE) : Tensor{}; + Tensor pool_qkv_context = pool.valid() + ? pool.slice(0, i, i + 1) + .slice(1, + num_tokens_img / POOL_SIZE, + num_tokens_img / POOL_SIZE + num_tokens_txt / POOL_SIZE) + : Tensor{}; + + // qkv_proj.forward(norm1_output.x.slice(0, i, i + 1), qkv); + // debug("qkv_raw", qkv); + + debug("rotary_emb", rotary_emb); + + qkv_proj.forward( + norm1_output.x.slice(0, i, i + 1), qkv, pool_qkv, norm_q.weight, norm_k.weight, rotary_emb); + debug("qkv", qkv); + + // qkv_proj_context.forward(norm1_context_output.x.slice(0, i, i + 1), qkv_context); + // debug("qkv_context_raw", qkv_context); + + debug("rotary_emb_context", rotary_emb_context); + + qkv_proj_context.forward(norm1_context_output.x.slice(0, i, i + 1), + qkv_context, + pool_qkv_context, + norm_added_q.weight, + norm_added_k.weight, + rotary_emb_context); + debug("qkv_context", qkv_context); + } + + nvtxRangePop(); + } + + spdlog::debug("concat={}", concat.shape.str()); + debug("concat", concat); + + assert(concat.shape[2] == num_heads * dim_head * 3); + + nvtxRangePushA("Attention"); + + if (pool.valid()) { + raw_attn_output = attn.forward(concat, pool, sparsityRatio); + } else { + raw_attn_output = attn.forward(concat); + } + + nvtxRangePop(); + + spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str()); + + raw_attn_output = raw_attn_output.view({batch_size, num_tokens_img + num_tokens_txt, num_heads, dim_head}); + + // IP_adapter + Tensor q_all = concat.slice(2, 0, num_heads * dim_head); // [B, N_total, dim] + + Tensor q_img = q_all.slice(1, 0, num_tokens_img); // [B, N_img, dim] + + q_heads = make_contiguous(q_img); + + } else if (attnImpl == AttentionImpl::NunchakuFP16) { + num_tokens_img_pad = ceilDiv(num_tokens_img, 256) * 256; + num_tokens_txt_pad = ceilDiv(num_tokens_txt, 256) * 256; + + Tensor concat_q, concat_k, concat_v; + + { + nvtxRangePushA("qkv_proj"); + + concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head}, + Tensor::FP16, + norm1_output.x.device()); + concat_k = Tensor::empty_like(concat_q); + concat_v = Tensor::empty_like(concat_q); + + for (int i = 0; i < batch_size; i++) { + // img first + auto sliceImg = [&](Tensor x) { return x.slice(0, i, i + 1).slice(2, 0, num_tokens_img_pad); }; + auto sliceTxt = [&](Tensor x) { + return x.slice(0, i, i + 1).slice(2, num_tokens_img_pad, num_tokens_img_pad + num_tokens_txt_pad); + }; + + qkv_proj.forward(norm1_output.x.slice(0, i, i + 1), + {}, + {}, + norm_q.weight, + norm_k.weight, + rotary_emb, + sliceImg(concat_q), + sliceImg(concat_k), + sliceImg(concat_v), + num_tokens_img); + + qkv_proj_context.forward(norm1_context_output.x.slice(0, i, i + 1), + {}, + {}, + norm_added_q.weight, + norm_added_k.weight, + rotary_emb_context, + sliceTxt(concat_q), + sliceTxt(concat_k), + sliceTxt(concat_v), + num_tokens_txt); + } + + debug("concat_q", concat_q); + debug("concat_k", concat_k); + debug("concat_v", concat_v); + + nvtxRangePop(); + } + + raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head}, + norm1_output.x.scalar_type(), + norm1_output.x.device()); + + nvtxRangePushA("Attention"); + + kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5))); + + nvtxRangePop(); + + raw_attn_output = + raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head}); + + Tensor q_tmp = concat_q.slice(2, 0, num_tokens_img); + Tensor q_view = q_tmp.transpose(1, 2).reshape({batch_size * num_tokens_img, num_heads * dim_head}); + + q_heads = make_contiguous(q_view); + + } else { + assert(false); + } + + debug("raw_attn_output", raw_attn_output); + + { + nvtxRangePushA("o_proj"); + + auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output; + + // raw_attn_output: [batch_size, num_tokens_img + num_tokens_txt, num_heads * dim_head] + + Tensor raw_attn_output_split; + if (batch_size == 1) { + raw_attn_output_split = + raw_attn_output.slice(1, 0, num_tokens_img).reshape({batch_size, num_tokens_img, num_heads * dim_head}); + } else { + raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_img, num_heads * dim_head}, + raw_attn_output.scalar_type(), + raw_attn_output.device()); + checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), + num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), + raw_attn_output.data_ptr(), + (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * + raw_attn_output.scalar_size(), + num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), + batch_size, + cudaMemcpyDeviceToDevice, + stream)); + } + + spdlog::debug("raw_attn_output_split={}", raw_attn_output_split.shape.str()); + debug("img.raw_attn_output_split", raw_attn_output_split); + + Tensor attn_output = + forward_fc(out_proj, raw_attn_output_split); // std::get(out_proj.forward(raw_attn_output_split)); + debug("img.attn_output", attn_output); + +#if 1 + // kernels::mul_add(attn_output, gate_msa, hidden_states); + kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, hidden_states, true); + hidden_states = std::move(attn_output); + + nvtxRangePop(); + nvtxRangePushA("MLP"); + + spdlog::debug("attn_output={}", hidden_states.shape.str()); + + Tensor norm_hidden_states = norm2.forward(hidden_states); + debug("scale_mlp", scale_mlp); + debug("shift_mlp", shift_mlp); + // kernels::mul_add(norm_hidden_states, scale_mlp, shift_mlp); + kernels::mul_add_batch(norm_hidden_states, scale_mlp, true, 0.0, shift_mlp, true); + + spdlog::debug("norm_hidden_states={}", norm_hidden_states.shape.str()); +#else + Tensor norm_hidden_states = hidden_states; +#endif + + // Tensor ff_output = mlp_fc2.forward(GELU::forward(mlp_fc1.forward(norm_hidden_states))); + debug("img.ff_input", norm_hidden_states); + Tensor ff_output = forward_mlp(mlp_fc1, mlp_fc2, norm_hidden_states); + debug("img.ff_output", ff_output); + + debug("gate_mlp", gate_mlp); + // kernels::mul_add(ff_output, gate_mlp, hidden_states); + kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, hidden_states, true); + hidden_states = std::move(ff_output); + + nvtxRangePop(); + + spdlog::debug("ff_output={}", hidden_states.shape.str()); + } + + if (context_pre_only) { + return {hidden_states, encoder_hidden_states, q_heads}; + } + + { + nvtxRangePushA("o_proj_context"); + + auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_context_output; + + Tensor raw_attn_output_split; + if (batch_size == 1) { + raw_attn_output_split = raw_attn_output.slice(1, num_tokens_img_pad, num_tokens_img_pad + num_tokens_txt) + .reshape({batch_size, num_tokens_txt, num_heads * dim_head}); + } else { + raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head}, + raw_attn_output.scalar_type(), + raw_attn_output.device()); + checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), + num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), + raw_attn_output.data_ptr() + num_tokens_img_pad * num_heads * dim_head * + raw_attn_output_split.scalar_size(), + (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * + raw_attn_output.scalar_size(), + num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), + batch_size, + cudaMemcpyDeviceToDevice, + stream)); + } + + spdlog::debug("raw_attn_output_split={}", raw_attn_output_split.shape.str()); + debug("context.raw_attn_output_split", raw_attn_output_split); + + Tensor attn_output = + forward_fc(out_proj_context, + raw_attn_output_split); // std::get(out_proj_context.forward(raw_attn_output_split)); + debug("context.attn_output", attn_output); + +#if 1 + // kernels::mul_add(attn_output, gate_msa, encoder_hidden_states); + kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, encoder_hidden_states, true); + encoder_hidden_states = std::move(attn_output); + + nvtxRangePop(); + nvtxRangePushA("MLP"); + + spdlog::debug("attn_output={}", encoder_hidden_states.shape.str()); + + Tensor norm_hidden_states = norm2_context.forward(encoder_hidden_states); + debug("c_scale_mlp", scale_mlp); + debug("c_shift_mlp", shift_mlp); + // kernels::mul_add(norm_hidden_states, scale_mlp, shift_mlp); + kernels::mul_add_batch(norm_hidden_states, scale_mlp, true, 0.0, shift_mlp, true); + + spdlog::debug("norm_hidden_states={}", norm_hidden_states.shape.str()); +#else + auto norm_hidden_states = encoder_hidden_states; +#endif + + // Tensor ff_output = mlp_context_fc2.forward(GELU::forward(mlp_context_fc1.forward(norm_hidden_states))); + // Tensor ff_output = + // mlp_context_fc2.forward_quant(quant_static_fuse_gelu(mlp_context_fc1.forward(norm_hidden_states), 1.0)); + debug("context.ff_input", norm_hidden_states); + Tensor ff_output = forward_mlp(mlp_context_fc1, mlp_context_fc2, norm_hidden_states); + debug("context.ff_output", ff_output); + + debug("c_gate_mlp", gate_mlp); + // kernels::mul_add(ff_output, gate_mlp, encoder_hidden_states); + kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, encoder_hidden_states, true); + encoder_hidden_states = std::move(ff_output); + + nvtxRangePop(); + + spdlog::debug("ff_output={}", encoder_hidden_states.shape.str()); + } + + nvtxRangePop(); + + return {hidden_states, encoder_hidden_states, q_heads}; +} + FluxModel::FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Device device) : dtype(dtype), offload(offload) { for (int i = 0; i < 19; i++) { @@ -967,6 +1351,112 @@ std::tuple FluxModel::forward_layer(size_t layer, return {hidden_states, encoder_hidden_states}; } +std::tuple FluxModel::forward_ip_adapter(size_t layer, + Tensor hidden_states, // [B, Nq, dim] + Tensor encoder_hidden_states, // [B, Nt, dim] + Tensor temb, + Tensor rotary_emb_img, // [B, Nq, dim_head] + Tensor rotary_emb_context, + Tensor k_img, // for Scaled_Dot_Product_Attention + Tensor v_img, // for Scaled_Dot_Product_Attention + Tensor controlnet_block_samples, + Tensor controlnet_single_block_samples) { + if (layer >= transformer_blocks.size()) { + auto [hs, enc] = forward_layer(layer, + hidden_states, + encoder_hidden_states, + temb, + rotary_emb_img, + rotary_emb_context, + controlnet_block_samples, + controlnet_single_block_samples); + Tensor dummy = Tensor{}; + return {hs, enc, dummy}; + } + + JointTransformerBlock *block = transformer_blocks.at(layer).get(); + + Tensor ip_query; + std::tie(hidden_states, encoder_hidden_states, ip_query) = block->forward_ip_adapter_branch( + hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f); + + return {hidden_states, encoder_hidden_states, ip_query}; + + /* + auto make_contiguous4d = [&](const Tensor &t4) { + // t4.shape.dataExtent = {B, H, N, C} + const auto &D = t4.shape.dataExtent; + size_t E = t4.scalar_size(); + + Tensor out = Tensor::allocate( + {D[0], D[1], D[2], D[3]}, + t4.scalarType, + t4.device() + ); + + out.shape.dataStride.resize(4); + int64_t acc = 1; + for (int d = 3; d >= 0; --d) { + out.shape.dataStride[d] = acc; + acc *= out.shape.dataExtent[d]; + } + + auto stream = getCurrentCUDAStream(); + + size_t width = D[3] * E; + size_t dst_pitch = width; + size_t src_pitch = t4.stride(2) * E; + + for (int b = 0; b < D[0]; ++b) { + for (int h = 0; h < D[1]; ++h) { + // src pointer = base + (b*stride(0) + h*stride(1))*E + const char *src = (const char*)t4.data_ptr() + + (t4.stride(0) * b + t4.stride(1) * h) * E; + // dst pointer = base + ((b*H + h)*N*C)*E + char *dst = (char*)out.data_ptr() + + ((b * D[1] + h) * D[2] * D[3]) * E; + + checkCUDA(cudaMemcpy2DAsync( + dst, dst_pitch, + src, src_pitch, + width, + D[2], + cudaMemcpyDeviceToDevice, + stream)); + } + } + checkCUDA(cudaStreamSynchronize(stream)); + return out; + }; + + int B = ip_query.size(0); + int Nq = ip_query.size(1); + int dim = ip_query.size(2); + constexpr int HEAD_DIM = 128; + int num_heads = dim / HEAD_DIM; + float scale = 1.0f / std::sqrt(float(HEAD_DIM)); + + Tensor q0 = ip_query + .reshape(TensorShape{B, Nq, num_heads, HEAD_DIM}) + .transpose(1, 2); + + Tensor q = make_contiguous4d(q0); + Tensor k = k_img; + Tensor v = v_img; + + Tensor ip_attn_output = Tensor::empty( + TensorShape{B, Nq, num_heads * HEAD_DIM}, + Tensor::FP16, + q.device() + ); + + //TODO + kernels::Scaled_Dot_Product_Attention(q, k, v, ip_attn_output, pow(HEAD_DIM, (-0.5))); + + return {hidden_states, encoder_hidden_states, q}; + */ +} + void FluxModel::setAttentionImpl(AttentionImpl impl) { for (auto &&block : this->transformer_blocks) { block->attnImpl = impl; diff --git a/src/FluxModel.h b/src/FluxModel.h index e4b3a340..7be53bd5 100644 --- a/src/FluxModel.h +++ b/src/FluxModel.h @@ -133,6 +133,12 @@ class JointTransformerBlock : public Module { Tensor rotary_emb, Tensor rotary_emb_context, float sparsityRatio); + std::tuple forward_ip_adapter_branch(Tensor hidden_states, + Tensor encoder_hidden_states, + Tensor temb, + Tensor rotary_emb, + Tensor rotary_emb_context, + float sparsityRatio); public: const int dim; @@ -178,6 +184,18 @@ class FluxModel : public Module { Tensor rotary_emb_context, Tensor controlnet_block_samples, Tensor controlnet_single_block_samples); + + std::tuple forward_ip_adapter(size_t layer, + Tensor hidden_states, + Tensor encoder_hidden_states, + Tensor temb, + Tensor rotary_emb_img, + Tensor rotary_emb_context, + Tensor k_img, + Tensor v_img, + Tensor controlnet_block_samples, + Tensor controlnet_single_block_samples); + void setAttentionImpl(AttentionImpl impl); void set_residual_callback(std::function cb); diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py new file mode 100644 index 00000000..ae1f26e5 --- /dev/null +++ b/tests/flux/test_flux_dev_IPA.py @@ -0,0 +1,56 @@ +import numpy as np +import pytest +import torch +import torch.nn.functional as F +from diffusers import FluxPipeline +from diffusers.utils import load_image + +from nunchaku import NunchakuFluxTransformer2dModel +from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe +from nunchaku.models.IP_adapter.utils import get_puild_embed, resize_numpy_image_long +from nunchaku.utils import get_precision, is_turing + + +@pytest.mark.skipif(is_turing(), reason="Skip tests due to using Turing GPUs") +def test_flux_dev_IPA(): + precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU + transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-{precision}-flux.1-dev") + + pipeline = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 + ).to("cuda") + + pipeline.load_ip_adapter( + pretrained_model_name_or_path_or_dict="XLabs-AI/flux-ip-adapter-v2", + weight_name="ip_adapter.safetensors", + image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", + ) + apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") + + id_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") + + image = pipeline( + prompt="A woman holding a sign that says 'SVDQuant is fast!", + ip_adapter_image=id_image.convert("RGB"), + num_inference_steps=50, + generator=torch.Generator("cuda"), + ).images[0] + + id_image = id_image.convert("RGB") + id_image_numpy = np.array(id_image) + id_image = resize_numpy_image_long(id_image_numpy, 1024) + id_embeddings, _ = get_puild_embed(id_image, "cuda") + + output_image = image.convert("RGB") + output_image_numpy = np.array(output_image) + output_image = resize_numpy_image_long(output_image_numpy, 1024) + output_id_embeddings, _ = get_puild_embed(output_image, "cuda") + cosine_similarities = ( + F.cosine_similarity(id_embeddings.view(32, 2048), output_id_embeddings.view(32, 2048), dim=1).mean().item() + ) + print(cosine_similarities) + assert cosine_similarities > 0.70 + + +if __name__ == "__main__": + test_flux_dev_IPA() From a4647e173038c7d09a806a065cfd917b57c56dc3 Mon Sep 17 00:00:00 2001 From: Bluear7878 Date: Mon, 23 Jun 2025 17:16:54 +0900 Subject: [PATCH 02/16] FBCache and comfyUI --- examples/flux.1-dev-IP-adapter.py | 8 + nunchaku/caching/diffusers_adapters/flux.py | 10 +- nunchaku/caching/utils.py | 14 +- nunchaku/csrc/flux.h | 17 +- nunchaku/csrc/pybind.cpp | 2 - .../IP_adapter/diffusers_adapters/flux.py | 40 ++- nunchaku/models/IP_adapter/utils.py | 246 +++++++++++------- nunchaku/pipeline/pipeline_flux_IPA.py | 41 +++ src/FluxModel.cpp | 178 ++++++------- src/FluxModel.h | 8 +- 10 files changed, 348 insertions(+), 216 deletions(-) create mode 100644 nunchaku/pipeline/pipeline_flux_IPA.py diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index d92b504e..dc330ed5 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -3,6 +3,7 @@ from diffusers.utils import load_image from nunchaku import NunchakuFluxTransformer2dModel +from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe from nunchaku.utils import get_precision @@ -18,9 +19,16 @@ weight_name="ip_adapter.safetensors", image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", ) +apply_cache_on_pipe( + pipeline, + use_double_fb_cache=True, + residual_diff_threshold_multi=0.09, + residual_diff_threshold_single=0.12, +) apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") + IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") image = pipeline( diff --git a/nunchaku/caching/diffusers_adapters/flux.py b/nunchaku/caching/diffusers_adapters/flux.py index 78c970c5..ce50e44b 100644 --- a/nunchaku/caching/diffusers_adapters/flux.py +++ b/nunchaku/caching/diffusers_adapters/flux.py @@ -13,8 +13,13 @@ def apply_cache_on_transformer( use_double_fb_cache: bool = False, residual_diff_threshold: float = 0.12, residual_diff_threshold_multi: float | None = None, - residual_diff_threshold_single: float = 0.1, + residual_diff_threshold_single: float | None = None, ): + if not hasattr(transformer, "_original_forward"): + transformer._original_forward = transformer.forward + if not hasattr(transformer, "_original_blocks"): + transformer._original_blocks = transformer.transformer_blocks + if residual_diff_threshold_multi is None: residual_diff_threshold_multi = residual_diff_threshold @@ -49,6 +54,9 @@ def new_forward(self, *args, **kwargs): transformer.forward = new_forward.__get__(transformer) transformer._is_cached = True + transformer.use_double_fb_cache = use_double_fb_cache + transformer.residual_diff_threshold_multi = residual_diff_threshold_multi + transformer.residual_diff_threshold_single = residual_diff_threshold_single return transformer diff --git a/nunchaku/caching/utils.py b/nunchaku/caching/utils.py index c1acd5aa..14d8a719 100644 --- a/nunchaku/caching/utils.py +++ b/nunchaku/caching/utils.py @@ -260,7 +260,7 @@ def forward( can_use_cache, _ = get_can_use_cache( first_hidden_states_residual, threshold=self.residual_diff_threshold, - parallelized=self.transformer is not None and getattr(self.transformer, "_is_parallelized", False), + parallelized=False, ) torch._dynamo.graph_break() @@ -329,7 +329,7 @@ def __init__( verbose: bool = False, ): super().__init__() - self.transformer = transformer + # self.transformer = transformer self.transformer_blocks = transformer.transformer_blocks self.single_transformer_blocks = transformer.single_transformer_blocks @@ -486,7 +486,7 @@ def forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, threshold=self.residual_diff_threshold_multi, - parallelized=(self.transformer is not None and getattr(self.transformer, "_is_parallelized", False)), + parallelized=False, mode="multi", verbose=self.verbose, call_remaining_fn=call_remaining_fn, @@ -515,7 +515,7 @@ def forward( hidden_states=cat_hidden_states, encoder_hidden_states=None, threshold=self.residual_diff_threshold_single, - parallelized=(self.transformer is not None and getattr(self.transformer, "_is_parallelized", False)), + parallelized=False, mode="single", verbose=self.verbose, call_remaining_fn=call_remaining_fn_single, @@ -591,8 +591,9 @@ def call_remaining_multi_transformer_blocks( controlnet_single_block_samples=None, skip_first_layer=False, txt_tokens=None, + start_idx=1, ): - start_idx = 1 + start_idx = start_idx original_hidden_states = hidden_states.clone() original_encoder_hidden_states = encoder_hidden_states.clone() @@ -627,8 +628,9 @@ def call_remaining_single_transformer_blocks( controlnet_single_block_samples=None, skip_first_layer=False, txt_tokens=None, + start_idx=1, ): - start_idx = 1 + start_idx = start_idx original_hidden_states = hidden_states.clone() for idx in range(start_idx, num_single_transformer_blocks): diff --git a/nunchaku/csrc/flux.h b/nunchaku/csrc/flux.h index 6dfeff12..f2488400 100644 --- a/nunchaku/csrc/flux.h +++ b/nunchaku/csrc/flux.h @@ -220,8 +220,6 @@ class QuantizedFluxModel : public ModuleWrapper { // : public torch:: torch::Tensor temb, torch::Tensor rotary_emb_img, torch::Tensor rotary_emb_context, - torch::Tensor k_img, - torch::Tensor v_img, std::optional controlnet_block_samples = std::nullopt, std::optional controlnet_single_block_samples = std::nullopt) { CUDADeviceContext ctx(deviceId); @@ -233,8 +231,6 @@ class QuantizedFluxModel : public ModuleWrapper { // : public torch:: temb = temb.contiguous(); rotary_emb_img = rotary_emb_img.contiguous(); rotary_emb_context = rotary_emb_context.contiguous(); - k_img = k_img.contiguous(); - v_img = v_img.contiguous(); auto &&[hidden_states_, encoder_hidden_states_, ip_query_] = net->forward_ip_adapter( idx, @@ -243,22 +239,13 @@ class QuantizedFluxModel : public ModuleWrapper { // : public torch:: from_torch(temb), from_torch(rotary_emb_img), from_torch(rotary_emb_context), - from_torch(k_img), - from_torch(v_img), controlnet_block_samples.has_value() ? from_torch(controlnet_block_samples.value().contiguous()) : Tensor{}, controlnet_single_block_samples.has_value() ? from_torch(controlnet_single_block_samples.value().contiguous()) : Tensor{}); - /* - auto ip_attn_output_contig = - ip_attn_output_.is_contiguous() - ? ip_attn_output_ - : ip_attn_output_.copy(ip_attn_output_.device()); - */ - hidden_states = to_torch(hidden_states_); - encoder_hidden_states = to_torch(encoder_hidden_states_); - // torch::Tensor ip_attn_output = to_torch(ip_attn_output_contig); + hidden_states = to_torch(hidden_states_); + encoder_hidden_states = to_torch(encoder_hidden_states_); torch::Tensor ip_query = to_torch(ip_query_); Tensor::synchronizeDevice(); diff --git a/nunchaku/csrc/pybind.cpp b/nunchaku/csrc/pybind.cpp index 7195b21c..f8097ae8 100644 --- a/nunchaku/csrc/pybind.cpp +++ b/nunchaku/csrc/pybind.cpp @@ -57,8 +57,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("temb"), py::arg("rotary_emb_img"), py::arg("rotary_emb_context"), - py::arg("k_img"), - py::arg("v_img"), py::arg("controlnet_block_samples") = py::none(), py::arg("controlnet_single_block_samples") = py::none()) .def("forward_single_layer", &QuantizedFluxModel::forward_single_layer) diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py index 58319aab..20f01f76 100644 --- a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py +++ b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py @@ -1,11 +1,16 @@ +import functools +import unittest + from diffusers import DiffusionPipeline, FluxTransformer2DModel from torch import nn +from nunchaku.caching.utils import cache_context, create_cache_context +from nunchaku.models.IP_adapter.utils import undo_all_mods_on_transformer + from ...IP_adapter import utils def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_scale: float = 1.0, repo_id: str): - IPA_transformer_blocks = nn.ModuleList( [ utils.IPA_TransformerBlocks( @@ -16,19 +21,52 @@ def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_ ) ] ) + if getattr(transformer, "_is_cached", False): + IPA_transformer_blocks[0].update_residual_diff_threshold( + use_double_fb_cache=transformer.use_double_fb_cache, + residual_diff_threshold_multi=transformer.residual_diff_threshold_multi, + residual_diff_threshold_single=transformer.residual_diff_threshold_single, + ) + undo_all_mods_on_transformer(transformer) + if not hasattr(transformer, "_original_forward"): + transformer._original_forward = transformer.forward + if not hasattr(transformer, "_original_blocks"): + transformer._original_blocks = transformer.transformer_blocks + dummy_single_transformer_blocks = nn.ModuleList() IPA_transformer_blocks[0].load_ip_adapter_weights_per_layer(repo_id=repo_id) transformer.transformer_blocks = IPA_transformer_blocks transformer.single_transformer_blocks = dummy_single_transformer_blocks + original_forward = transformer.forward + @functools.wraps(original_forward) + def new_forward(self, *args, **kwargs): + with ( + unittest.mock.patch.object(self, "transformer_blocks", IPA_transformer_blocks), + unittest.mock.patch.object(self, "single_transformer_blocks", dummy_single_transformer_blocks), + ): + return original_forward(*args, **kwargs) + + transformer.forward = new_forward.__get__(transformer) transformer._is_IPA = True return transformer def apply_IPA_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs): + if getattr(pipe, "_is_cached", False): + original_call = pipe.__class__.__call__ + + @functools.wraps(original_call) + def new_call(self, *args, **kwargs): + with cache_context(create_cache_context()): + return original_call(self, *args, **kwargs) + + pipe.__class__.__call__ = new_call + pipe.__class__._is_cached = True + if not shallow_patch: apply_IPA_on_transformer(pipe.transformer, **kwargs) diff --git a/nunchaku/models/IP_adapter/utils.py b/nunchaku/models/IP_adapter/utils.py index daa02319..9c6473c6 100644 --- a/nunchaku/models/IP_adapter/utils.py +++ b/nunchaku/models/IP_adapter/utils.py @@ -5,6 +5,7 @@ import numpy as np import torch import torch.nn.functional as F +from diffusers import FluxTransformer2DModel from facexlib.parsing import init_parsing_model from facexlib.utils.face_restoration_helper import FaceRestoreHelper from huggingface_hub import hf_hub_download @@ -15,6 +16,7 @@ from torchvision.transforms.functional import normalize, resize from torchvision.utils import make_grid +from nunchaku.caching.utils import FluxCachedTransformerBlocks, check_and_apply_cache from nunchaku.models.pulid.encoders_transformer import IDFormer from nunchaku.models.pulid.eva_clip import create_model_and_transforms from nunchaku.models.pulid.eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD @@ -24,7 +26,7 @@ num_single_transformer_blocks = 38 # FIXME -class IPA_TransformerBlocks(nn.Module): +class IPA_TransformerBlocks(FluxCachedTransformerBlocks): def __init__( self, *, @@ -35,37 +37,17 @@ def __init__( verbose: bool = False, device: str | torch.device, ): - super().__init__() - self.transformer = transformer - self.transformer_blocks = transformer.transformer_blocks - self.return_hidden_states_first = return_hidden_states_first - self.return_hidden_states_only = return_hidden_states_only - self.verbose = verbose + super().__init__( + transformer=transformer, + use_double_fb_cache=False, + residual_diff_threshold_multi=-1, + residual_diff_threshold_single=-1, + return_hidden_states_first=return_hidden_states_first, + return_hidden_states_only=return_hidden_states_only, + verbose=verbose, + ) self.ip_adapter_scale = ip_adapter_scale - - self.m = self.transformer_blocks[0].m - self.dtype = torch.bfloat16 if self.m.isBF16() else torch.float16 - self.device = device - - @staticmethod - def pack_rotemb(rotemb: torch.Tensor) -> torch.Tensor: - assert rotemb.dtype == torch.float32 - B = rotemb.shape[0] - M = rotemb.shape[1] - D = rotemb.shape[2] * 2 - msg_shape = "rotemb shape must be (B, M, D//2, 1, 2)" - assert rotemb.shape == (B, M, D // 2, 1, 2), msg_shape - assert M % 16 == 0 - assert D % 8 == 0 - rotemb = rotemb.reshape(B, M // 16, 16, D // 8, 8) - rotemb = rotemb.permute(0, 1, 3, 2, 4) - # 16*8 pack, FP32 accumulator (C) format - # https://docs.nvidia.com/cuda/parallel-thread-execution/#mma-16816-c - rotemb = rotemb.reshape(*rotemb.shape[0:3], 2, 8, 4, 2) - rotemb = rotemb.permute(0, 1, 2, 4, 5, 3, 6) - rotemb = rotemb.contiguous() - rotemb = rotemb.view(B, M, D) - return rotemb + self.image_embeds = None def forward( self, @@ -80,45 +62,40 @@ def forward( controlnet_single_block_samples=None, skip_first_layer=False, ): - # batch_size = hidden_states.shape[0] + batch_size = hidden_states.shape[0] txt_tokens = encoder_hidden_states.shape[1] img_tokens = hidden_states.shape[1] - self.id_embeddings = id_embeddings - self.id_weight = id_weight - self.pulid_ca_idx = 0 - if self.id_embeddings is not None: - self.set_residual_callback() - original_dtype = hidden_states.dtype original_device = hidden_states.device - hidden_states = hidden_states.to(self.dtype).to(self.device) - encoder_hidden_states = encoder_hidden_states.to(self.dtype).to(self.device) - temb = temb.to(self.dtype).to(self.device) - image_rotary_emb = image_rotary_emb.to(self.device) + hidden_states = hidden_states.to(self.dtype).to(original_device) + encoder_hidden_states = encoder_hidden_states.to(self.dtype).to(original_device) + temb = temb.to(self.dtype).to(original_device) + image_rotary_emb = image_rotary_emb.to(original_device) if controlnet_block_samples is not None: - if len(controlnet_block_samples) > 0: - controlnet_block_samples = torch.stack(controlnet_block_samples).to(self.device) - else: - controlnet_block_samples = None - + controlnet_block_samples = ( + torch.stack(controlnet_block_samples).to(original_device) if len(controlnet_block_samples) > 0 else None + ) if controlnet_single_block_samples is not None: - if len(controlnet_single_block_samples) > 0: - controlnet_single_block_samples = torch.stack(controlnet_single_block_samples).to(self.device) - else: - controlnet_single_block_samples = None + controlnet_single_block_samples = ( + torch.stack(controlnet_single_block_samples).to(original_device) + if len(controlnet_single_block_samples) > 0 + else None + ) assert image_rotary_emb.ndim == 6 assert image_rotary_emb.shape[0] == 1 assert image_rotary_emb.shape[1] == 1 - assert image_rotary_emb.shape[2] == 1 * (txt_tokens + img_tokens) - # [1, tokens, head_dim / 2, 1, 2] (sincos) + # [1, tokens, head_dim/2, 1, 2] (sincos) + total_tokens = txt_tokens + img_tokens + assert image_rotary_emb.shape[2] == 1 * total_tokens + image_rotary_emb = image_rotary_emb.reshape([1, txt_tokens + img_tokens, *image_rotary_emb.shape[3:]]) - rotary_emb_txt = image_rotary_emb[:, :txt_tokens, ...] # .to(self.dtype) - rotary_emb_img = image_rotary_emb[:, txt_tokens:, ...] # .to(self.dtype) - rotary_emb_single = image_rotary_emb # .to(self.dtype) + rotary_emb_txt = image_rotary_emb[:, :txt_tokens, ...] + rotary_emb_img = image_rotary_emb[:, txt_tokens:, ...] + rotary_emb_single = image_rotary_emb rotary_emb_txt = self.pack_rotemb(pad_tensor(rotary_emb_txt, 256, 1)) rotary_emb_img = self.pack_rotemb(pad_tensor(rotary_emb_img, 256, 1)) @@ -126,6 +103,8 @@ def forward( if joint_attention_kwargs is not None and "ip_hidden_states" in joint_attention_kwargs: ip_hidden_states = joint_attention_kwargs.pop("ip_hidden_states") + elif self.image_embeds is not None: + ip_hidden_states = self.image_embeds remaining_kwargs = { "temb": temb, @@ -139,15 +118,97 @@ def forward( } torch._dynamo.graph_break() - updated_h, updated_enc, _, _ = self.call_multi_transformer_blocks( - hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, **remaining_kwargs + + if (self.residual_diff_threshold_multi <= 0.0) or (batch_size > 1): + updated_h, updated_enc, _, _ = self.call_IPA_multi_transformer_blocks( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + skip_block=False, + **remaining_kwargs, + ) + + remaining_kwargs.pop("ip_hidden_states", None) + cat_hidden_states = torch.cat([updated_enc, updated_h], dim=1) + + updated_cat, _ = self.call_remaining_single_transformer_blocks( + hidden_states=cat_hidden_states, encoder_hidden_states=None, start_idx=0, **remaining_kwargs + ) + # torch._dynamo.graph_break() + + final_enc = updated_cat[:, :txt_tokens, ...] + final_h = updated_cat[:, txt_tokens:, ...] + + final_h = final_h.to(original_dtype).to(original_device) + final_enc = final_enc.to(original_dtype).to(original_device) + + if self.return_hidden_states_only: + return final_h + if self.return_hidden_states_first: + return final_h, final_enc + return final_enc, final_h + + original_hidden_states = hidden_states + first_hidden_states, first_encoder_hidden_states, _, _ = self.call_IPA_multi_transformer_blocks( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + first_block=True, + skip_block=False, + **remaining_kwargs, + ) + + hidden_states = first_hidden_states + encoder_hidden_states = first_encoder_hidden_states + first_hidden_states_residual_multi = hidden_states - original_hidden_states + del original_hidden_states + + call_remaining_fn = self.call_IPA_multi_transformer_blocks + + torch._dynamo.graph_break() + updated_h, updated_enc, threshold = check_and_apply_cache( + first_residual=first_hidden_states_residual_multi, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + threshold=self.residual_diff_threshold_multi, + parallelized=False, + mode="multi", + verbose=self.verbose, + call_remaining_fn=call_remaining_fn, + remaining_kwargs=remaining_kwargs, ) + self.residual_diff_threshold_multi = threshold + + # Single layer + remaining_kwargs.pop("ip_hidden_states", None) cat_hidden_states = torch.cat([updated_enc, updated_h], dim=1) + original_cat = cat_hidden_states + if not self.use_double_fb_cache: + ##NO FBCache + updated_cat, _ = self.call_remaining_single_transformer_blocks( + hidden_states=cat_hidden_states, encoder_hidden_states=None, start_idx=0, **remaining_kwargs + ) + else: + # USE FBCache + cat_hidden_states = self.m.forward_single_layer(0, cat_hidden_states, temb, rotary_emb_single) + + first_hidden_states_residual_single = cat_hidden_states - original_cat + del original_cat + + call_remaining_fn_single = self.call_remaining_single_transformer_blocks + + updated_cat, _, threshold = check_and_apply_cache( + first_residual=first_hidden_states_residual_single, + hidden_states=cat_hidden_states, + encoder_hidden_states=None, + threshold=self.residual_diff_threshold_single, + parallelized=False, + mode="single", + verbose=self.verbose, + call_remaining_fn=call_remaining_fn_single, + remaining_kwargs=remaining_kwargs, + ) + self.residual_diff_threshold_single = threshold - updated_cat, _ = self.call_single_transformer_blocks( - hidden_states=cat_hidden_states, encoder_hidden_states=None, **remaining_kwargs - ) # torch._dynamo.graph_break() final_enc = updated_cat[:, :txt_tokens, ...] @@ -162,7 +223,7 @@ def forward( return final_h, final_enc return final_enc, final_h - def call_multi_transformer_blocks( + def call_IPA_multi_transformer_blocks( self, hidden_states: torch.Tensor, temb: torch.Tensor, @@ -175,12 +236,20 @@ def call_multi_transformer_blocks( skip_first_layer=False, txt_tokens=None, ip_hidden_states=None, + first_block: bool = False, + skip_block: bool = True, ): - start_idx = 0 + if first_block and skip_block: + raise ValueError("`first_block` and `skip_block` cannot both be True.") + + start_idx = 1 if skip_block else 0 + end_idx = 1 if first_block else num_transformer_blocks + original_hidden_states = hidden_states.clone() original_encoder_hidden_states = encoder_hidden_states.clone() + ip_hidden_states[0] = ip_hidden_states[0].to(self.dtype).to(self.device) - for idx in range(start_idx, num_transformer_blocks): + for idx in range(start_idx, end_idx): k_img = self.ip_k_projs[idx](ip_hidden_states[0]) v_img = self.ip_v_projs[idx](ip_hidden_states[0]) @@ -191,12 +260,11 @@ def call_multi_transformer_blocks( temb, rotary_emb_img, rotary_emb_txt, - k_img.squeeze(dim=1), - v_img.squeeze(dim=1), controlnet_block_samples, controlnet_single_block_samples, ) - ip_query = ip_query.contiguous() + + ip_query = ip_query.contiguous().to(self.dtype) ip_query = ip_query.view(1, -1, 24, 128).transpose(1, 2) k_img = k_img.view(1, -1, 24, 128).transpose(1, 2) @@ -211,39 +279,10 @@ def call_multi_transformer_blocks( hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() - hs_res = hidden_states - original_hidden_states enc_res = encoder_hidden_states - original_encoder_hidden_states - return hidden_states, encoder_hidden_states, hs_res, enc_res - def call_single_transformer_blocks( - self, - hidden_states: torch.Tensor, - temb: torch.Tensor, - encoder_hidden_states: torch.Tensor, - rotary_emb_img: torch.Tensor, - rotary_emb_txt: torch.Tensor, - rotary_emb_single: torch.Tensor, - controlnet_block_samples=None, - controlnet_single_block_samples=None, - skip_first_layer=False, - txt_tokens=None, - ip_hidden_states=None, - ): - start_idx = 0 - original_hidden_states = hidden_states.clone() - - for idx in range(start_idx, num_single_transformer_blocks): - hidden_states = self.m.forward_single_layer( - idx, - hidden_states, - temb, - rotary_emb_single, - ) - - hidden_states = hidden_states.contiguous() - hs_res = hidden_states - original_hidden_states - return hidden_states, hs_res + return hidden_states, encoder_hidden_states, hs_res, enc_res def load_ip_adapter_weights_per_layer( self, @@ -291,6 +330,9 @@ def load_ip_adapter_weights_per_layer( self.ip_k_projs.append(k_proj) self.ip_v_projs.append(v_proj) + def set_ip_hidden_states(self, image_embeds, negative_image_embeds=None): + self.image_embeds = image_embeds + def resize_numpy_image_long(image, resize_long_edge=768): h, w = image.shape[:2] @@ -513,3 +555,13 @@ def get_puild_embed(image, device, cal_uncond=False, weight_dtype=torch.bfloat16 uncond_id_embedding = pulid_encoder(id_uncond, id_vit_hidden_uncond) return id_embedding, uncond_id_embedding + + +def undo_all_mods_on_transformer(transformer: FluxTransformer2DModel): + if hasattr(transformer, "_original_forward"): + transformer.forward = transformer._original_forward + del transformer._original_forward + if hasattr(transformer, "_original_blocks"): + transformer.transformer_blocks = transformer._original_blocks + del transformer._original_blocks + return transformer diff --git a/nunchaku/pipeline/pipeline_flux_IPA.py b/nunchaku/pipeline/pipeline_flux_IPA.py new file mode 100644 index 00000000..c4aa19d0 --- /dev/null +++ b/nunchaku/pipeline/pipeline_flux_IPA.py @@ -0,0 +1,41 @@ +from typing import Any, List, Optional + +import torch +from diffusers import FluxPipeline + + +class FluxPipelineWrapper(FluxPipeline): + @torch.no_grad() + def get_image_embeds( + self, + num_images_per_prompt: int = 1, + ip_adapter_image: Optional[Any] = None, # PipelineImageInput + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_ip_adapter_image: Optional[Any] = None, # PipelineImageInput + negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + ) -> (Optional[torch.Tensor], Optional[torch.Tensor]): + batch_size = 1 + + device = self.transformer.device + + image_embeds = None + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image=ip_adapter_image, + ip_adapter_image_embeds=ip_adapter_image_embeds, + device=device, + num_images_per_prompt=batch_size * num_images_per_prompt, + ) + image_embeds = self.transformer.encoder_hid_proj(image_embeds) + + negative_image_embeds = None + if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None: + negative_image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image=negative_ip_adapter_image, + ip_adapter_image_embeds=negative_ip_adapter_image_embeds, + device=device, + num_images_per_prompt=batch_size * num_images_per_prompt, + ) + negative_image_embeds = self.transformer.encoder_hid_proj(negative_image_embeds) + + return image_embeds, negative_image_embeds diff --git a/src/FluxModel.cpp b/src/FluxModel.cpp index c57da8ae..c796eace 100644 --- a/src/FluxModel.cpp +++ b/src/FluxModel.cpp @@ -775,6 +775,73 @@ std::tuple JointTransformerBlock::forward(Tensor hidden_states, return {hidden_states, encoder_hidden_states}; } +Tensor JointTransformerBlock::get_q_heads(Tensor hidden_states, + Tensor encoder_hidden_states, + Tensor temb, + Tensor rotary_emb, + Tensor rotary_emb_context, + float sparsityRatio) { + int batch_size = hidden_states.shape[0]; + int num_tokens_img = hidden_states.shape[1]; + int num_tokens_txt = encoder_hidden_states.shape[1]; + + // Apply AdaNorm. + auto norm1_output = norm1.forward(hidden_states, temb); + auto norm1_context_output = norm1_context.forward(encoder_hidden_states, temb); + + Tensor concat = Tensor::allocate( + {batch_size, num_tokens_img + num_tokens_txt, dim * 3}, norm1_output.x.scalar_type(), norm1_output.x.device()); + + const bool blockSparse = sparsityRatio > 0; + constexpr int POOL_SIZE = Attention::POOL_SIZE; + const int poolTokens = num_tokens_img / POOL_SIZE + num_tokens_txt / POOL_SIZE; + Tensor pool = + blockSparse + ? Tensor::allocate({batch_size, poolTokens, dim * 3}, norm1_output.x.scalar_type(), norm1_output.x.device()) + : Tensor{}; + + // QKV Projection. + for (int i = 0; i < batch_size; i++) { + Tensor qkv = concat.slice(0, i, i + 1).slice(1, 0, num_tokens_img); + Tensor qkv_context = concat.slice(0, i, i + 1).slice(1, num_tokens_img, num_tokens_img + num_tokens_txt); + Tensor pool_qkv = pool.valid() ? pool.slice(0, i, i + 1).slice(1, 0, num_tokens_img / POOL_SIZE) : Tensor{}; + Tensor pool_qkv_context = + pool.valid() ? pool.slice(0, i, i + 1).slice(1, num_tokens_img / POOL_SIZE, poolTokens) : Tensor{}; + + qkv_proj.forward(norm1_output.x.slice(0, i, i + 1), qkv, pool_qkv, norm_q.weight, norm_k.weight, rotary_emb); + qkv_proj_context.forward(norm1_context_output.x.slice(0, i, i + 1), + qkv_context, + pool_qkv_context, + norm_added_q.weight, + norm_added_k.weight, + rotary_emb_context); + } + + // Extract and return q_heads. + Tensor q_all = concat.slice(2, 0, num_heads * dim_head); + Tensor q_img = q_all.slice(1, 0, num_tokens_img); + + auto make_contiguous = [&](const Tensor &t) { + int B = t.shape.dataExtent[0]; + int R = t.shape.dataExtent[1]; + int C = t.shape.dataExtent[2]; + size_t E = t.scalar_size(); + size_t src_pitch = t.stride(1) * E; + size_t dst_pitch = C * E; + size_t width = C * E; + size_t height = R; + Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device()); + auto stream = getCurrentCUDAStream(); + for (int b = 0; b < B; ++b) { + const void *src = (const char *)t.data_ptr() + t.stride(0) * b * E; + void *dst = (char *)out.data_ptr() + out.stride(0) * b * E; + checkCUDA( + cudaMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, cudaMemcpyDeviceToDevice, stream)); + } + return out; + }; + return make_contiguous(q_img); +} std::tuple JointTransformerBlock::forward_ip_adapter_branch(Tensor hidden_states, Tensor encoder_hidden_states, @@ -1000,11 +1067,7 @@ std::tuple JointTransformerBlock::forward_ip_adapter_bra raw_attn_output = raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head}); - Tensor q_tmp = concat_q.slice(2, 0, num_tokens_img); - Tensor q_view = q_tmp.transpose(1, 2).reshape({batch_size * num_tokens_img, num_heads * dim_head}); - - q_heads = make_contiguous(q_view); - + q_heads = concat_q; } else { assert(false); } @@ -1357,104 +1420,35 @@ std::tuple FluxModel::forward_ip_adapter(size_t layer, Tensor temb, Tensor rotary_emb_img, // [B, Nq, dim_head] Tensor rotary_emb_context, - Tensor k_img, // for Scaled_Dot_Product_Attention - Tensor v_img, // for Scaled_Dot_Product_Attention Tensor controlnet_block_samples, Tensor controlnet_single_block_samples) { - if (layer >= transformer_blocks.size()) { - auto [hs, enc] = forward_layer(layer, - hidden_states, - encoder_hidden_states, - temb, - rotary_emb_img, - rotary_emb_context, - controlnet_block_samples, - controlnet_single_block_samples); - Tensor dummy = Tensor{}; - return {hs, enc, dummy}; + if (offload && layer > 0) { + if (layer < transformer_blocks.size()) { + transformer_blocks.at(layer)->loadLazyParams(); + } else { + transformer_blocks.at(layer - transformer_blocks.size())->loadLazyParams(); + } } - JointTransformerBlock *block = transformer_blocks.at(layer).get(); - - Tensor ip_query; - std::tie(hidden_states, encoder_hidden_states, ip_query) = block->forward_ip_adapter_branch( + std::tie(hidden_states, encoder_hidden_states) = transformer_blocks.at(layer)->forward( + hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f); + Tensor ip_query = transformer_blocks.at(layer)->get_q_heads( hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f); - return {hidden_states, encoder_hidden_states, ip_query}; - - /* - auto make_contiguous4d = [&](const Tensor &t4) { - // t4.shape.dataExtent = {B, H, N, C} - const auto &D = t4.shape.dataExtent; - size_t E = t4.scalar_size(); - - Tensor out = Tensor::allocate( - {D[0], D[1], D[2], D[3]}, - t4.scalarType, - t4.device() - ); - - out.shape.dataStride.resize(4); - int64_t acc = 1; - for (int d = 3; d >= 0; --d) { - out.shape.dataStride[d] = acc; - acc *= out.shape.dataExtent[d]; - } + if (controlnet_block_samples.valid()) { + const int num_controlnet_block_samples = controlnet_block_samples.shape[0]; - auto stream = getCurrentCUDAStream(); + int interval_control = ceilDiv(transformer_blocks.size(), static_cast(num_controlnet_block_samples)); + int block_index = layer / interval_control; - size_t width = D[3] * E; - size_t dst_pitch = width; - size_t src_pitch = t4.stride(2) * E; - - for (int b = 0; b < D[0]; ++b) { - for (int h = 0; h < D[1]; ++h) { - // src pointer = base + (b*stride(0) + h*stride(1))*E - const char *src = (const char*)t4.data_ptr() - + (t4.stride(0) * b + t4.stride(1) * h) * E; - // dst pointer = base + ((b*H + h)*N*C)*E - char *dst = (char*)out.data_ptr() - + ((b * D[1] + h) * D[2] * D[3]) * E; - - checkCUDA(cudaMemcpy2DAsync( - dst, dst_pitch, - src, src_pitch, - width, - D[2], - cudaMemcpyDeviceToDevice, - stream)); - } + hidden_states = kernels::add(hidden_states, controlnet_block_samples[block_index]); } - checkCUDA(cudaStreamSynchronize(stream)); - return out; - }; - int B = ip_query.size(0); - int Nq = ip_query.size(1); - int dim = ip_query.size(2); - constexpr int HEAD_DIM = 128; - int num_heads = dim / HEAD_DIM; - float scale = 1.0f / std::sqrt(float(HEAD_DIM)); - - Tensor q0 = ip_query - .reshape(TensorShape{B, Nq, num_heads, HEAD_DIM}) - .transpose(1, 2); - - Tensor q = make_contiguous4d(q0); - Tensor k = k_img; - Tensor v = v_img; - - Tensor ip_attn_output = Tensor::empty( - TensorShape{B, Nq, num_heads * HEAD_DIM}, - Tensor::FP16, - q.device() - ); - - //TODO - kernels::Scaled_Dot_Product_Attention(q, k, v, ip_attn_output, pow(HEAD_DIM, (-0.5))); + if (offload && layer > 0) { + transformer_blocks.at(layer)->releaseLazyParams(); + } - return {hidden_states, encoder_hidden_states, q}; - */ + return {hidden_states, encoder_hidden_states, ip_query}; } void FluxModel::setAttentionImpl(AttentionImpl impl) { diff --git a/src/FluxModel.h b/src/FluxModel.h index 7be53bd5..7182b800 100644 --- a/src/FluxModel.h +++ b/src/FluxModel.h @@ -139,6 +139,12 @@ class JointTransformerBlock : public Module { Tensor rotary_emb, Tensor rotary_emb_context, float sparsityRatio); + Tensor get_q_heads(Tensor hidden_states, + Tensor encoder_hidden_states, + Tensor temb, + Tensor rotary_emb, + Tensor rotary_emb_context, + float sparsityRatio); public: const int dim; @@ -191,8 +197,6 @@ class FluxModel : public Module { Tensor temb, Tensor rotary_emb_img, Tensor rotary_emb_context, - Tensor k_img, - Tensor v_img, Tensor controlnet_block_samples, Tensor controlnet_single_block_samples); From 10d2ac31d058526a7872e880bb094cd700eea27f Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Tue, 22 Jul 2025 07:13:47 +0800 Subject: [PATCH 03/16] fixing conflicts --- nunchaku/caching/diffusers_adapters/flux.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nunchaku/caching/diffusers_adapters/flux.py b/nunchaku/caching/diffusers_adapters/flux.py index cd7a728a..6c05221e 100644 --- a/nunchaku/caching/diffusers_adapters/flux.py +++ b/nunchaku/caching/diffusers_adapters/flux.py @@ -54,6 +54,11 @@ def apply_cache_on_transformer( ----- If already cached, only updates thresholds. Caching is only active within a cache context. """ + if not hasattr(transformer, "_original_forward"): + transformer._original_forward = transformer.forward + if not hasattr(transformer, "_original_blocks"): + transformer._original_blocks = transformer.transformer_blocks + if residual_diff_threshold_multi is None: residual_diff_threshold_multi = residual_diff_threshold From 4ea1a887e22d873608e3d60991e03e502812cc33 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Tue, 22 Jul 2025 07:46:14 +0800 Subject: [PATCH 04/16] update --- examples/flux.1-canny-dev.py | 2 +- examples/flux.1-depth-dev-lora.py | 2 +- examples/flux.1-depth-dev.py | 2 +- examples/flux.1-dev-IP-adapter.py | 10 ++--- examples/flux.1-dev-cache.py | 2 +- examples/flux.1-dev-controlnet-union-pro.py | 2 +- examples/flux.1-dev-double_cache.py | 2 +- .../flux.1-dev-double_cache_offloading.py | 2 +- examples/flux.1-dev-fp16attn.py | 2 +- examples/flux.1-dev-lora.py | 2 +- examples/flux.1-dev-multiple-lora.py | 2 +- examples/flux.1-dev-offload.py | 2 +- examples/flux.1-dev-pulid.py | 2 +- examples/flux.1-dev-qencoder.py | 2 +- examples/flux.1-dev-teacache.py | 2 +- examples/flux.1-dev-turing.py | 2 +- examples/flux.1-dev.py | 2 +- examples/flux.1-fill-dev.py | 2 +- examples/flux.1-kontext-dev.py | 2 +- examples/flux.1-redux-dev.py | 2 +- examples/flux.1-schnell.py | 2 +- examples/sana1.6b-cache.py | 2 +- examples/sana1.6b.py | 2 +- examples/sana1.6b_pag.py | 2 +- .../IP_adapter/diffusers_adapters/flux.py | 5 +-- nunchaku/pipeline/pipeline_flux_IPA.py | 41 ------------------- requirements.txt | 2 + 27 files changed, 32 insertions(+), 72 deletions(-) delete mode 100644 nunchaku/pipeline/pipeline_flux_IPA.py diff --git a/examples/flux.1-canny-dev.py b/examples/flux.1-canny-dev.py index 2feafc54..4fc0eac1 100644 --- a/examples/flux.1-canny-dev.py +++ b/examples/flux.1-canny-dev.py @@ -8,7 +8,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-canny-dev/svdq-{precision}_r32-flux.1-canny-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-canny-dev/svdq-{precision}_r32-flux.1-canny-dev.safetensors" ) pipe = FluxControlPipeline.from_pretrained( "black-forest-labs/FLUX.1-Canny-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-depth-dev-lora.py b/examples/flux.1-depth-dev-lora.py index 1135c985..6eb13c91 100644 --- a/examples/flux.1-depth-dev-lora.py +++ b/examples/flux.1-depth-dev-lora.py @@ -8,7 +8,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-depth-dev/svdq-{precision}_r32-flux.1-depth-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-depth-dev/svdq-{precision}_r32-flux.1-depth-dev.safetensors" ) pipe = FluxControlPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-depth-dev.py b/examples/flux.1-depth-dev.py index 3f546bb6..1571b645 100644 --- a/examples/flux.1-depth-dev.py +++ b/examples/flux.1-depth-dev.py @@ -8,7 +8,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-depth-dev/svdq-{precision}_r32-flux.1-depth-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-depth-dev/svdq-{precision}_r32-flux.1-depth-dev.safetensors" ) pipe = FluxControlPipeline.from_pretrained( diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index dc330ed5..ddfda804 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -8,8 +8,9 @@ from nunchaku.utils import get_precision precision = get_precision() -transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-{precision}-flux.1-dev") - +transformer = NunchakuFluxTransformer2dModel.from_pretrained( + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" +) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 ).to("cuda") @@ -29,13 +30,12 @@ apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") -IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") +IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/lecun.jpg?raw=true") image = pipeline( prompt="A woman holding a sign that says 'SVDQuant is fast!", ip_adapter_image=IP_image.convert("RGB"), - num_inference_steps=50, - generator=torch.Generator("cuda"), + num_inference_steps=50 ).images[0] image.save(f"flux.1-dev-IP-adapter-{precision}.png") diff --git a/examples/flux.1-dev-cache.py b/examples/flux.1-dev-cache.py index 6f39b528..8a8bca82 100644 --- a/examples/flux.1-dev-cache.py +++ b/examples/flux.1-dev-cache.py @@ -7,7 +7,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-dev-controlnet-union-pro.py b/examples/flux.1-dev-controlnet-union-pro.py index eb3b9810..b227bbfc 100644 --- a/examples/flux.1-dev-controlnet-union-pro.py +++ b/examples/flux.1-dev-controlnet-union-pro.py @@ -15,7 +15,7 @@ precision = get_precision() need_offload = get_gpu_memory() < 36 transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", torch_dtype=torch.bfloat16, offload=need_offload, ) diff --git a/examples/flux.1-dev-double_cache.py b/examples/flux.1-dev-double_cache.py index dba6ab7d..03b1ccc0 100644 --- a/examples/flux.1-dev-double_cache.py +++ b/examples/flux.1-dev-double_cache.py @@ -8,7 +8,7 @@ precision = get_precision() transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( diff --git a/examples/flux.1-dev-double_cache_offloading.py b/examples/flux.1-dev-double_cache_offloading.py index a61992bb..4066c909 100644 --- a/examples/flux.1-dev-double_cache_offloading.py +++ b/examples/flux.1-dev-double_cache_offloading.py @@ -8,7 +8,7 @@ precision = get_precision() transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", offload=True, ) diff --git a/examples/flux.1-dev-fp16attn.py b/examples/flux.1-dev-fp16attn.py index d8a5e30f..a5ab9870 100644 --- a/examples/flux.1-dev-fp16attn.py +++ b/examples/flux.1-dev-fp16attn.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) transformer.set_attention_impl("nunchaku-fp16") # set attention implementation to fp16 pipeline = FluxPipeline.from_pretrained( diff --git a/examples/flux.1-dev-lora.py b/examples/flux.1-dev-lora.py index 38651213..1f255543 100644 --- a/examples/flux.1-dev-lora.py +++ b/examples/flux.1-dev-lora.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-dev-multiple-lora.py b/examples/flux.1-dev-multiple-lora.py index 546dadaa..3d58d359 100644 --- a/examples/flux.1-dev-multiple-lora.py +++ b/examples/flux.1-dev-multiple-lora.py @@ -7,7 +7,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-dev-offload.py b/examples/flux.1-dev-offload.py index 1558bfd8..7e0e77ca 100644 --- a/examples/flux.1-dev-offload.py +++ b/examples/flux.1-dev-offload.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", offload=True + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", offload=True ) # set offload to False if you want to disable offloading pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-dev-pulid.py b/examples/flux.1-dev-pulid.py index b0c32788..5cb82558 100644 --- a/examples/flux.1-dev-pulid.py +++ b/examples/flux.1-dev-pulid.py @@ -10,7 +10,7 @@ precision = get_precision() transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = PuLIDFluxPipeline.from_pretrained( diff --git a/examples/flux.1-dev-qencoder.py b/examples/flux.1-dev-qencoder.py index 8c22ac99..266c5d67 100644 --- a/examples/flux.1-dev-qencoder.py +++ b/examples/flux.1-dev-qencoder.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors") pipeline = FluxPipeline.from_pretrained( diff --git a/examples/flux.1-dev-teacache.py b/examples/flux.1-dev-teacache.py index 05698d44..53eed088 100644 --- a/examples/flux.1-dev-teacache.py +++ b/examples/flux.1-dev-teacache.py @@ -9,7 +9,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-dev-turing.py b/examples/flux.1-dev-turing.py index a22969a3..903db4ef 100644 --- a/examples/flux.1-dev-turing.py +++ b/examples/flux.1-dev-turing.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", offload=True, torch_dtype=torch.float16, # Turing GPUs only support fp16 precision ) # set offload to False if you want to disable offloading diff --git a/examples/flux.1-dev.py b/examples/flux.1-dev.py index d386778b..78e05c63 100644 --- a/examples/flux.1-dev.py +++ b/examples/flux.1-dev.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-fill-dev.py b/examples/flux.1-fill-dev.py index 9b413a24..9afaae49 100644 --- a/examples/flux.1-fill-dev.py +++ b/examples/flux.1-fill-dev.py @@ -10,7 +10,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-fill-dev/svdq-{precision}_r32-flux.1-fill-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-fill-dev/svdq-{precision}_r32-flux.1-fill-dev.safetensors" ) pipe = FluxFillPipeline.from_pretrained( "black-forest-labs/FLUX.1-Fill-dev", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/flux.1-kontext-dev.py b/examples/flux.1-kontext-dev.py index 26257ab6..2868d582 100644 --- a/examples/flux.1-kontext-dev.py +++ b/examples/flux.1-kontext-dev.py @@ -6,7 +6,7 @@ from nunchaku.utils import get_precision transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{get_precision()}_r32-flux.1-kontext-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-kontext-dev/svdq-{get_precision()}_r32-flux.1-kontext-dev.safetensors" ) pipeline = FluxKontextPipeline.from_pretrained( diff --git a/examples/flux.1-redux-dev.py b/examples/flux.1-redux-dev.py index f104db5e..7c4859f2 100644 --- a/examples/flux.1-redux-dev.py +++ b/examples/flux.1-redux-dev.py @@ -10,7 +10,7 @@ "black-forest-labs/FLUX.1-Redux-dev", torch_dtype=torch.bfloat16 ).to("cuda") transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" ) pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", diff --git a/examples/flux.1-schnell.py b/examples/flux.1-schnell.py index 47877c98..be22bc22 100644 --- a/examples/flux.1-schnell.py +++ b/examples/flux.1-schnell.py @@ -6,7 +6,7 @@ precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{precision}_r32-flux.1-schnell.safetensors" + f"nunchaku-tech/nunchaku-flux.1-schnell/svdq-{precision}_r32-flux.1-schnell.safetensors" ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", transformer=transformer, torch_dtype=torch.bfloat16 diff --git a/examples/sana1.6b-cache.py b/examples/sana1.6b-cache.py index 0125f43c..d5f33f0b 100644 --- a/examples/sana1.6b-cache.py +++ b/examples/sana1.6b-cache.py @@ -5,7 +5,7 @@ from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe transformer = NunchakuSanaTransformer2DModel.from_pretrained( - "mit-han-lab/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors" + "nunchaku-tech/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors" ) pipe = SanaPipeline.from_pretrained( "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", diff --git a/examples/sana1.6b.py b/examples/sana1.6b.py index f56bf3e9..502e5661 100644 --- a/examples/sana1.6b.py +++ b/examples/sana1.6b.py @@ -4,7 +4,7 @@ from nunchaku import NunchakuSanaTransformer2DModel transformer = NunchakuSanaTransformer2DModel.from_pretrained( - "mit-han-lab/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors" + "nunchaku-tech/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors" ) pipe = SanaPipeline.from_pretrained( "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", diff --git a/examples/sana1.6b_pag.py b/examples/sana1.6b_pag.py index 771988d0..f4a8f8b2 100644 --- a/examples/sana1.6b_pag.py +++ b/examples/sana1.6b_pag.py @@ -4,7 +4,7 @@ from nunchaku import NunchakuSanaTransformer2DModel transformer = NunchakuSanaTransformer2DModel.from_pretrained( - "mit-han-lab/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors", pag_layers=8 + "nunchaku-tech/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors", pag_layers=8 ) pipe = SanaPAGPipeline.from_pretrained( "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py index 20f01f76..b489e82f 100644 --- a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py +++ b/nunchaku/models/IP_adapter/diffusers_adapters/flux.py @@ -55,7 +55,7 @@ def new_forward(self, *args, **kwargs): return transformer -def apply_IPA_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs): +def apply_IPA_on_pipe(pipe: DiffusionPipeline, **kwargs): if getattr(pipe, "_is_cached", False): original_call = pipe.__class__.__call__ @@ -67,7 +67,6 @@ def new_call(self, *args, **kwargs): pipe.__class__.__call__ = new_call pipe.__class__._is_cached = True - if not shallow_patch: - apply_IPA_on_transformer(pipe.transformer, **kwargs) + apply_IPA_on_transformer(pipe.transformer, **kwargs) return pipe diff --git a/nunchaku/pipeline/pipeline_flux_IPA.py b/nunchaku/pipeline/pipeline_flux_IPA.py deleted file mode 100644 index c4aa19d0..00000000 --- a/nunchaku/pipeline/pipeline_flux_IPA.py +++ /dev/null @@ -1,41 +0,0 @@ -from typing import Any, List, Optional - -import torch -from diffusers import FluxPipeline - - -class FluxPipelineWrapper(FluxPipeline): - @torch.no_grad() - def get_image_embeds( - self, - num_images_per_prompt: int = 1, - ip_adapter_image: Optional[Any] = None, # PipelineImageInput - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[Any] = None, # PipelineImageInput - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - ) -> (Optional[torch.Tensor], Optional[torch.Tensor]): - batch_size = 1 - - device = self.transformer.device - - image_embeds = None - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image=ip_adapter_image, - ip_adapter_image_embeds=ip_adapter_image_embeds, - device=device, - num_images_per_prompt=batch_size * num_images_per_prompt, - ) - image_embeds = self.transformer.encoder_hid_proj(image_embeds) - - negative_image_embeds = None - if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None: - negative_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image=negative_ip_adapter_image, - ip_adapter_image_embeds=negative_ip_adapter_image_embeds, - device=device, - num_images_per_prompt=batch_size * num_images_per_prompt, - ) - negative_image_embeds = self.transformer.encoder_hid_proj(negative_image_embeds) - - return image_embeds, negative_image_embeds diff --git a/requirements.txt b/requirements.txt index b06c3b34..8a51071b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ insightface opencv-python facexlib onnxruntime +# ip-adapter +timm \ No newline at end of file From 5bacadbe8807a377e206b9457fcae95f334f644e Mon Sep 17 00:00:00 2001 From: Bluear7878 Date: Wed, 23 Jul 2025 13:52:35 +0900 Subject: [PATCH 05/16] update example --- examples/flux.1-dev-IP-adapter.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index dc330ed5..2eb1ace5 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -19,6 +19,9 @@ weight_name="ip_adapter.safetensors", image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", ) + +apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.1, repo_id="XLabs-AI/flux-ip-adapter-v2") + apply_cache_on_pipe( pipeline, use_double_fb_cache=True, @@ -26,13 +29,10 @@ residual_diff_threshold_single=0.12, ) -apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") - - -IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") +IP_image = load_image("https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/ComfyUI-nunchaku/inputs/monalisa.jpg") image = pipeline( - prompt="A woman holding a sign that says 'SVDQuant is fast!", + prompt="holding an iconic starbucks cup of coffee", ip_adapter_image=IP_image.convert("RGB"), num_inference_steps=50, generator=torch.Generator("cuda"), From 50ddd6a44ba65f7cdc179df687cc10db4f5ceb2e Mon Sep 17 00:00:00 2001 From: Bluear7878 Date: Wed, 23 Jul 2025 13:55:46 +0900 Subject: [PATCH 06/16] update example --- examples/flux.1-dev-IP-adapter.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index 2eb1ace5..ecfcb677 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -29,7 +29,9 @@ residual_diff_threshold_single=0.12, ) -IP_image = load_image("https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/ComfyUI-nunchaku/inputs/monalisa.jpg") +IP_image = load_image( + "https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/ComfyUI-nunchaku/inputs/monalisa.jpg" +) image = pipeline( prompt="holding an iconic starbucks cup of coffee", From 62763ea5ad29c8f5c2c3162521fc99e57667c9db Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 05:06:59 +0800 Subject: [PATCH 07/16] style: make linter happy --- examples/flux.1-dev-IP-adapter.py | 4 +--- requirements.txt | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index 47a15498..a4c2cd30 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -35,9 +35,7 @@ ) image = pipeline( - prompt="holding an iconic starbucks cup of coffee", - ip_adapter_image=IP_image.convert("RGB"), - num_inference_steps=50 + prompt="holding an iconic starbucks cup of coffee", ip_adapter_image=IP_image.convert("RGB"), num_inference_steps=50 ).images[0] image.save(f"flux.1-dev-IP-adapter-{precision}.png") diff --git a/requirements.txt b/requirements.txt index 8a51071b..67f10f8c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,4 @@ opencv-python facexlib onnxruntime # ip-adapter -timm \ No newline at end of file +timm From ac6c719dbc45740f874cc14fbcb350c8a425ab3d Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 06:28:40 +0800 Subject: [PATCH 08/16] update --- examples/flux.1-dev-IP-adapter.py | 2 +- tests/flux/test_flux_dev_IPA.py | 38 +++++++++++++++++++++++-------- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index a4c2cd30..1b96a8d6 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -35,7 +35,7 @@ ) image = pipeline( - prompt="holding an iconic starbucks cup of coffee", ip_adapter_image=IP_image.convert("RGB"), num_inference_steps=50 + prompt="holding an sign saying 'SVDQuant is fast!'", ip_adapter_image=IP_image.convert("RGB"), num_inference_steps=50 ).images[0] image.save(f"flux.1-dev-IP-adapter-{precision}.png") diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index ae1f26e5..69a0ea3b 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -9,12 +9,15 @@ from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe from nunchaku.models.IP_adapter.utils import get_puild_embed, resize_numpy_image_long from nunchaku.utils import get_precision, is_turing - +import gc +from nunchaku.pipeline.pipeline_flux_pulid import PuLIDFluxPipeline @pytest.mark.skipif(is_turing(), reason="Skip tests due to using Turing GPUs") def test_flux_dev_IPA(): precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU - transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-{precision}-flux.1-dev") + transformer = NunchakuFluxTransformer2dModel.from_pretrained( + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + ) pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16 @@ -25,31 +28,48 @@ def test_flux_dev_IPA(): weight_name="ip_adapter.safetensors", image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", ) - apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2") + apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.1, repo_id="XLabs-AI/flux-ip-adapter-v2") - id_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") + id_image = load_image( + "https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/ComfyUI-nunchaku/inputs/monalisa.jpg" + ) image = pipeline( - prompt="A woman holding a sign that says 'SVDQuant is fast!", + prompt="holding an sign saying 'SVDQuant is fast!'", ip_adapter_image=id_image.convert("RGB"), num_inference_steps=50, - generator=torch.Generator("cuda"), ).images[0] + image.save('./tmp.png') + + del pipeline + del transformer + gc.collect() + torch.cuda.empty_cache() + + # use the pulid pipeline to get the id embedding + transformer = NunchakuFluxTransformer2dModel.from_pretrained( + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + ) + pipeline = PuLIDFluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + transformer=transformer, + torch_dtype=torch.bfloat16, + ).to("cuda") id_image = id_image.convert("RGB") id_image_numpy = np.array(id_image) id_image = resize_numpy_image_long(id_image_numpy, 1024) - id_embeddings, _ = get_puild_embed(id_image, "cuda") + id_embeddings, _ = pipeline.pulid_model.get_id_embedding(id_image) output_image = image.convert("RGB") output_image_numpy = np.array(output_image) output_image = resize_numpy_image_long(output_image_numpy, 1024) - output_id_embeddings, _ = get_puild_embed(output_image, "cuda") + output_id_embeddings, _ = pipeline.pulid_model.get_id_embedding(output_image) cosine_similarities = ( F.cosine_similarity(id_embeddings.view(32, 2048), output_id_embeddings.view(32, 2048), dim=1).mean().item() ) print(cosine_similarities) - assert cosine_similarities > 0.70 + assert cosine_similarities > 0.8 if __name__ == "__main__": From 799d41f82908f127b9cc7046ddf56ff995a4db5a Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 06:35:05 +0800 Subject: [PATCH 09/16] update ipa test --- examples/flux.1-dev-IP-adapter.py | 4 +- nunchaku/models/IP_adapter/utils.py | 225 ---------------------------- tests/flux/test_flux_dev_IPA.py | 15 +- 3 files changed, 12 insertions(+), 232 deletions(-) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index 1b96a8d6..77e7b036 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -35,7 +35,9 @@ ) image = pipeline( - prompt="holding an sign saying 'SVDQuant is fast!'", ip_adapter_image=IP_image.convert("RGB"), num_inference_steps=50 + prompt="holding an sign saying 'SVDQuant is fast!'", + ip_adapter_image=IP_image.convert("RGB"), + num_inference_steps=50, ).images[0] image.save(f"flux.1-dev-IP-adapter-{precision}.png") diff --git a/nunchaku/models/IP_adapter/utils.py b/nunchaku/models/IP_adapter/utils.py index 9c6473c6..d16a67ea 100644 --- a/nunchaku/models/IP_adapter/utils.py +++ b/nunchaku/models/IP_adapter/utils.py @@ -1,25 +1,12 @@ -import math - import cv2 -import insightface -import numpy as np import torch import torch.nn.functional as F from diffusers import FluxTransformer2DModel -from facexlib.parsing import init_parsing_model -from facexlib.utils.face_restoration_helper import FaceRestoreHelper from huggingface_hub import hf_hub_download -from insightface.app import FaceAnalysis from safetensors import safe_open from torch import nn -from torchvision.transforms import InterpolationMode -from torchvision.transforms.functional import normalize, resize -from torchvision.utils import make_grid from nunchaku.caching.utils import FluxCachedTransformerBlocks, check_and_apply_cache -from nunchaku.models.pulid.encoders_transformer import IDFormer -from nunchaku.models.pulid.eva_clip import create_model_and_transforms -from nunchaku.models.pulid.eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD from nunchaku.models.transformers.utils import pad_tensor num_transformer_blocks = 19 # FIXME @@ -345,218 +332,6 @@ def resize_numpy_image_long(image, resize_long_edge=768): return image -# from basicsr -def img2tensor(imgs, bgr2rgb=True, float32=True): - """Numpy array to tensor. - - Args: - imgs (list[ndarray] | ndarray): Input images. - bgr2rgb (bool): Whether to change bgr to rgb. - float32 (bool): Whether to change to float32. - - Returns: - list[tensor] | tensor: Tensor images. If returned results only have - one element, just return tensor. - """ - - def _totensor(img, bgr2rgb, float32): - if img.shape[2] == 3 and bgr2rgb: - if img.dtype == "float64": - img = img.astype("float32") - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img = torch.from_numpy(img.transpose(2, 0, 1)) - if float32: - img = img.float() - return img - - if isinstance(imgs, list): - return [_totensor(img, bgr2rgb, float32) for img in imgs] - return _totensor(imgs, bgr2rgb, float32) - - -def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): - """Convert torch Tensors into image numpy arrays. - - After clamping to [min, max], values will be normalized to [0, 1]. - - Args: - tensor (Tensor or list[Tensor]): Accept shapes: - 1) 4D mini-batch Tensor of shape (B x 3/1 x H x W); - 2) 3D Tensor of shape (3/1 x H x W); - 3) 2D Tensor of shape (H x W). - Tensor channel should be in RGB order. - rgb2bgr (bool): Whether to change rgb to bgr. - out_type (numpy type): output types. If ``np.uint8``, transform outputs - to uint8 type with range [0, 255]; otherwise, float type with - range [0, 1]. Default: ``np.uint8``. - min_max (tuple[int]): min and max values for clamp. - - Returns: - (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of - shape (H x W). The channel order is BGR. - """ - if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): - raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}") - - if torch.is_tensor(tensor): - tensor = [tensor] - result = [] - for _tensor in tensor: - _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) - _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) - - n_dim = _tensor.dim() - if n_dim == 4: - img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy() - img_np = img_np.transpose(1, 2, 0) - if rgb2bgr: - img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) - elif n_dim == 3: - img_np = _tensor.numpy() - img_np = img_np.transpose(1, 2, 0) - if img_np.shape[2] == 1: # gray image - img_np = np.squeeze(img_np, axis=2) - else: - if rgb2bgr: - img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) - elif n_dim == 2: - img_np = _tensor.numpy() - else: - raise TypeError(f"Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}") - if out_type == np.uint8: - # Unlike MATLAB, numpy.unit8() WILL NOT round by default. - img_np = (img_np * 255.0).round() - img_np = img_np.astype(out_type) - result.append(img_np) - if len(result) == 1: - result = result[0] - return result - - -def to_gray(img): - x = 0.299 * img[:, 0:1] + 0.587 * img[:, 1:2] + 0.114 * img[:, 2:3] - x = x.repeat(1, 3, 1, 1) - return x - - -def get_puild_embed(image, device, cal_uncond=False, weight_dtype=torch.bfloat16, onnx_provider="gpu"): - face_helper = FaceRestoreHelper( - upscale_factor=1, - face_size=512, - crop_ratio=(1, 1), - det_model="retinaface_resnet50", - save_ext="png", - device=device, - ) - face_helper.face_parse = None - face_helper.face_parse = init_parsing_model(model_name="bisenet", device=device) - - providers = ( - ["CPUExecutionProvider"] if onnx_provider == "cpu" else ["CUDAExecutionProvider", "CPUExecutionProvider"] - ) - - app = FaceAnalysis(name="antelopev2", root=".", providers=providers) - app.prepare(ctx_id=0, det_size=(640, 640)) - handler_ante = insightface.model_zoo.get_model("models/antelopev2/glintr100.onnx", providers=providers) - handler_ante.prepare(ctx_id=0) - - from safetensors.torch import load_file - - ckpt_path = hf_hub_download("guozinan/PuLID", "pulid_flux_v0.9.0.safetensors", local_dir="models") - - full_sd = load_file(ckpt_path) - - encoder_sd = {k.split("pulid_encoder.")[1]: v for k, v in full_sd.items() if k.startswith("pulid_encoder.")} - - pulid_encoder = IDFormer().to(device, weight_dtype) - pulid_encoder.load_state_dict(encoder_sd, strict=True) - pulid_encoder.eval() - - model, _, _ = create_model_and_transforms("EVA02-CLIP-L-14-336", "eva_clip", force_custom_clip=True) - model = model.visual - clip_vision_model = model.to(device, dtype=weight_dtype) - - face_helper.clean_all() - debug_img_list = [] - image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - # get antelopev2 embedding - face_info = app.get(image_bgr) - if len(face_info) > 0: - face_info = sorted(face_info, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[ - -1 - ] # only use the maximum face - id_ante_embedding = face_info["embedding"] - debug_img_list.append( - image[ - int(face_info["bbox"][1]) : int(face_info["bbox"][3]), - int(face_info["bbox"][0]) : int(face_info["bbox"][2]), - ] - ) - else: - id_ante_embedding = None - - # using facexlib to detect and align face - face_helper.read_image(image_bgr) - face_helper.get_face_landmarks_5(only_center_face=True) - face_helper.align_warp_face() - if len(face_helper.cropped_faces) == 0: - raise RuntimeError("facexlib align face fail") - align_face = face_helper.cropped_faces[0] - # incase insightface didn't detect face - if id_ante_embedding is None: - print("fail to detect face using insightface, extract embedding on align face") - id_ante_embedding = handler_ante.get_feat(align_face) - - id_ante_embedding = torch.from_numpy(id_ante_embedding).to(device, weight_dtype) - if id_ante_embedding.ndim == 1: - id_ante_embedding = id_ante_embedding.unsqueeze(0) - - # parsing - input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0 - input = input.to(device) - parsing_out = face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0] - parsing_out = parsing_out.argmax(dim=1, keepdim=True) - bg_label = [0, 16, 18, 7, 8, 9, 14, 15] - bg = sum(parsing_out == i for i in bg_label).bool() - white_image = torch.ones_like(input) - # only keep the face features - face_features_image = torch.where(bg, white_image, to_gray(input)) - debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False)) - - eva_transform_mean = getattr(clip_vision_model, "image_mean", OPENAI_DATASET_MEAN) - eva_transform_std = getattr(clip_vision_model, "image_std", OPENAI_DATASET_STD) - if not isinstance(eva_transform_mean, (list, tuple)): - eva_transform_mean = (eva_transform_mean,) * 3 - if not isinstance(eva_transform_std, (list, tuple)): - eva_transform_std = (eva_transform_std,) * 3 - eva_transform_mean = eva_transform_mean - eva_transform_std = eva_transform_std - - # transform img before sending to eva-clip-vit - face_features_image = resize(face_features_image, clip_vision_model.image_size, InterpolationMode.BICUBIC) - face_features_image = normalize(face_features_image, eva_transform_mean, eva_transform_std) - id_cond_vit, id_vit_hidden = clip_vision_model( - face_features_image.to(weight_dtype), return_all_features=False, return_hidden=True, shuffle=False - ) - id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True) - id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm) - - id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1) - - id_embedding = pulid_encoder(id_cond, id_vit_hidden) - - if not cal_uncond: - return id_embedding, None - - id_uncond = torch.zeros_like(id_cond) - id_vit_hidden_uncond = [] - for layer_idx in range(0, len(id_vit_hidden)): - id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx])) - uncond_id_embedding = pulid_encoder(id_uncond, id_vit_hidden_uncond) - - return id_embedding, uncond_id_embedding - - def undo_all_mods_on_transformer(transformer: FluxTransformer2DModel): if hasattr(transformer, "_original_forward"): transformer.forward = transformer._original_forward diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index 69a0ea3b..bc8d3dfe 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -1,3 +1,5 @@ +import gc + import numpy as np import pytest import torch @@ -7,10 +9,10 @@ from nunchaku import NunchakuFluxTransformer2dModel from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe -from nunchaku.models.IP_adapter.utils import get_puild_embed, resize_numpy_image_long -from nunchaku.utils import get_precision, is_turing -import gc +from nunchaku.models.IP_adapter.utils import resize_numpy_image_long from nunchaku.pipeline.pipeline_flux_pulid import PuLIDFluxPipeline +from nunchaku.utils import get_precision, is_turing + @pytest.mark.skipif(is_turing(), reason="Skip tests due to using Turing GPUs") def test_flux_dev_IPA(): @@ -28,7 +30,8 @@ def test_flux_dev_IPA(): weight_name="ip_adapter.safetensors", image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", ) - apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.1, repo_id="XLabs-AI/flux-ip-adapter-v2") + + apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.15, repo_id="XLabs-AI/flux-ip-adapter-v2") id_image = load_image( "https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/ComfyUI-nunchaku/inputs/monalisa.jpg" @@ -39,7 +42,7 @@ def test_flux_dev_IPA(): ip_adapter_image=id_image.convert("RGB"), num_inference_steps=50, ).images[0] - image.save('./tmp.png') + image.save("./tmp.png") del pipeline del transformer @@ -69,7 +72,7 @@ def test_flux_dev_IPA(): F.cosine_similarity(id_embeddings.view(32, 2048), output_id_embeddings.view(32, 2048), dim=1).mean().item() ) print(cosine_similarities) - assert cosine_similarities > 0.8 + assert cosine_similarities > 0.85 if __name__ == "__main__": From afa3bbb121aec0996ceadc8ad08cb3a22900f871 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 06:43:55 +0800 Subject: [PATCH 10/16] add docs and rename IP to ip --- examples/flux.1-dev-IP-adapter.py | 2 +- .../IP_adapter/diffusers_adapters/__init__.py | 12 -- .../{IP_adapter => ip_adapter}/__init__.py | 0 .../ip_adapter/diffusers_adapters/__init__.py | 42 +++++ .../diffusers_adapters/flux.py | 64 ++++++- .../{IP_adapter => ip_adapter}/utils.py | 171 ++++++++++++++++++ tests/flux/test_flux_dev_IPA.py | 4 +- 7 files changed, 274 insertions(+), 21 deletions(-) delete mode 100644 nunchaku/models/IP_adapter/diffusers_adapters/__init__.py rename nunchaku/models/{IP_adapter => ip_adapter}/__init__.py (100%) create mode 100644 nunchaku/models/ip_adapter/diffusers_adapters/__init__.py rename nunchaku/models/{IP_adapter => ip_adapter}/diffusers_adapters/flux.py (58%) rename nunchaku/models/{IP_adapter => ip_adapter}/utils.py (71%) diff --git a/examples/flux.1-dev-IP-adapter.py b/examples/flux.1-dev-IP-adapter.py index 77e7b036..1a5ab503 100644 --- a/examples/flux.1-dev-IP-adapter.py +++ b/examples/flux.1-dev-IP-adapter.py @@ -4,7 +4,7 @@ from nunchaku import NunchakuFluxTransformer2dModel from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe -from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe +from nunchaku.models.ip_adapter.diffusers_adapters import apply_IPA_on_pipe from nunchaku.utils import get_precision precision = get_precision() diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py b/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py deleted file mode 100644 index 2064cd48..00000000 --- a/nunchaku/models/IP_adapter/diffusers_adapters/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from diffusers import DiffusionPipeline - - -def apply_IPA_on_pipe(pipe: DiffusionPipeline, *args, **kwargs): - assert isinstance(pipe, DiffusionPipeline) - - pipe_cls_name = pipe.__class__.__name__ - if pipe_cls_name.startswith("Flux") or pipe_cls_name.startswith("PuLID"): - from .flux import apply_IPA_on_pipe as apply_IPA_on_pipe_fn - else: - raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}") - return apply_IPA_on_pipe_fn(pipe, *args, **kwargs) diff --git a/nunchaku/models/IP_adapter/__init__.py b/nunchaku/models/ip_adapter/__init__.py similarity index 100% rename from nunchaku/models/IP_adapter/__init__.py rename to nunchaku/models/ip_adapter/__init__.py diff --git a/nunchaku/models/ip_adapter/diffusers_adapters/__init__.py b/nunchaku/models/ip_adapter/diffusers_adapters/__init__.py new file mode 100644 index 00000000..fc9a62b4 --- /dev/null +++ b/nunchaku/models/ip_adapter/diffusers_adapters/__init__.py @@ -0,0 +1,42 @@ +""" +IP-Adapter integration for Diffusers pipelines. + +This module provides utilities to apply IP-Adapter modifications to compatible +Diffusers pipelines, such as Flux and PuLID pipelines. +""" + +from diffusers import DiffusionPipeline + + +def apply_IPA_on_pipe(pipe: DiffusionPipeline, *args, **kwargs): + """ + Apply IP-Adapter modifications to a supported Diffusers pipeline. + + Parameters + ---------- + pipe : DiffusionPipeline + The pipeline instance to modify. Must be a Flux or PuLID pipeline. + *args + Additional positional arguments passed to the underlying implementation. + **kwargs + Additional keyword arguments passed to the underlying implementation. + + Returns + ------- + DiffusionPipeline + The modified pipeline with IP-Adapter applied. + + Raises + ------ + ValueError + If the pipeline class is not supported. + + """ + assert isinstance(pipe, DiffusionPipeline) + + pipe_cls_name = pipe.__class__.__name__ + if pipe_cls_name.startswith("Flux") or pipe_cls_name.startswith("PuLID"): + from .flux import apply_IPA_on_pipe as apply_IPA_on_pipe_fn + else: + raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}") + return apply_IPA_on_pipe_fn(pipe, *args, **kwargs) diff --git a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py b/nunchaku/models/ip_adapter/diffusers_adapters/flux.py similarity index 58% rename from nunchaku/models/IP_adapter/diffusers_adapters/flux.py rename to nunchaku/models/ip_adapter/diffusers_adapters/flux.py index b489e82f..7703c983 100644 --- a/nunchaku/models/IP_adapter/diffusers_adapters/flux.py +++ b/nunchaku/models/ip_adapter/diffusers_adapters/flux.py @@ -1,3 +1,11 @@ +""" +IP-Adapter integration for Flux pipelines in Diffusers. + +This module provides functions to apply IP-Adapter modifications to +FluxTransformer2DModel and DiffusionPipeline objects, enabling image prompt +conditioning for generative models. +""" + import functools import unittest @@ -5,12 +13,37 @@ from torch import nn from nunchaku.caching.utils import cache_context, create_cache_context -from nunchaku.models.IP_adapter.utils import undo_all_mods_on_transformer - -from ...IP_adapter import utils - - -def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_scale: float = 1.0, repo_id: str): +from nunchaku.models.ip_adapter.utils import undo_all_mods_on_transformer + +from ...ip_adapter import utils + + +def apply_IPA_on_transformer( + transformer: FluxTransformer2DModel, + *, + ip_adapter_scale: float = 1.0, + repo_id: str, +): + """ + Apply IP-Adapter modifications to a FluxTransformer2DModel. + + This function replaces the transformer's blocks with IP-Adapter-enabled blocks, + loads per-layer IP-Adapter weights, and wraps the forward method to use the new blocks. + + Parameters + ---------- + transformer : FluxTransformer2DModel + The transformer model to modify. + ip_adapter_scale : float, optional + Scaling factor for the IP-Adapter (default is 1.0). + repo_id : str + HuggingFace Hub repository ID containing the IP-Adapter weights. + + Returns + ------- + FluxTransformer2DModel + The modified transformer with IP-Adapter support. + """ IPA_transformer_blocks = nn.ModuleList( [ utils.IPA_TransformerBlocks( @@ -56,6 +89,25 @@ def new_forward(self, *args, **kwargs): def apply_IPA_on_pipe(pipe: DiffusionPipeline, **kwargs): + """ + Apply IP-Adapter modifications to a DiffusionPipeline. + + This function modifies the pipeline's transformer to support IP-Adapter + conditioning. If the pipeline is cached, it also wraps the pipeline's + __call__ method to ensure cache context is used. + + Parameters + ---------- + pipe : DiffusionPipeline + The pipeline to modify. Must contain a FluxTransformer2DModel as its transformer. + **kwargs + Additional keyword arguments passed to `apply_IPA_on_transformer`. + + Returns + ------- + DiffusionPipeline + The modified pipeline with IP-Adapter support. + """ if getattr(pipe, "_is_cached", False): original_call = pipe.__class__.__call__ diff --git a/nunchaku/models/IP_adapter/utils.py b/nunchaku/models/ip_adapter/utils.py similarity index 71% rename from nunchaku/models/IP_adapter/utils.py rename to nunchaku/models/ip_adapter/utils.py index d16a67ea..4974d9aa 100644 --- a/nunchaku/models/IP_adapter/utils.py +++ b/nunchaku/models/ip_adapter/utils.py @@ -1,3 +1,11 @@ +""" +IP-Adapter utility functions and classes for FluxTransformer2DModel. + +This module provides the core implementation for integrating IP-Adapter +conditioning into Flux-based transformer models, including block modification, +weight loading, and image embedding support. +""" + import cv2 import torch import torch.nn.functional as F @@ -14,6 +22,35 @@ class IPA_TransformerBlocks(FluxCachedTransformerBlocks): + """ + Transformer block wrapper for IP-Adapter integration. + + This class extends FluxCachedTransformerBlocks to enable per-layer + IP-Adapter conditioning, efficient caching, and flexible output control. + + Parameters + ---------- + transformer : nn.Module, optional + The base transformer module to wrap. + ip_adapter_scale : float, default=1.0 + Scaling factor for the IP-Adapter output. + return_hidden_states_first : bool, default=True + If True, return hidden states before encoder states. + return_hidden_states_only : bool, default=False + If True, return only hidden states. + verbose : bool, default=False + If True, print verbose debug information. + device : str or torch.device + Device to use for computation. + + Attributes + ---------- + ip_adapter_scale : float + Scaling factor for IP-Adapter output. + image_embeds : torch.Tensor or None + Image embeddings for IP-Adapter conditioning. + """ + def __init__( self, *, @@ -49,6 +86,38 @@ def forward( controlnet_single_block_samples=None, skip_first_layer=False, ): + """ + Forward pass with IP-Adapter conditioning. + + Parameters + ---------- + hidden_states : torch.Tensor + Input hidden states. + temb : torch.Tensor + Temporal embedding tensor. + encoder_hidden_states : torch.Tensor + Encoder hidden states. + image_rotary_emb : torch.Tensor + Rotary embedding for image tokens. + id_embeddings : optional + Not used. + id_weight : optional + Not used. + joint_attention_kwargs : dict, optional + Additional attention arguments, may include 'ip_hidden_states'. + controlnet_block_samples : list, optional + ControlNet block samples for multi-blocks. + controlnet_single_block_samples : list, optional + ControlNet block samples for single blocks. + skip_first_layer : bool, default=False + If True, skip the first transformer block. + + Returns + ------- + tuple or torch.Tensor + Final hidden states and encoder states, or only hidden states if + `return_hidden_states_only` is True. + """ batch_size = hidden_states.shape[0] txt_tokens = encoder_hidden_states.shape[1] img_tokens = hidden_states.shape[1] @@ -226,6 +295,43 @@ def call_IPA_multi_transformer_blocks( first_block: bool = False, skip_block: bool = True, ): + """ + Apply IP-Adapter conditioning to multiple transformer blocks. + + Parameters + ---------- + hidden_states : torch.Tensor + Input hidden states. + temb : torch.Tensor + Temporal embedding tensor. + encoder_hidden_states : torch.Tensor + Encoder hidden states. + rotary_emb_img : torch.Tensor + Rotary embedding for image tokens. + rotary_emb_txt : torch.Tensor + Rotary embedding for text tokens. + rotary_emb_single : torch.Tensor + Rotary embedding for single block. + controlnet_block_samples : list, optional + ControlNet block samples for multi-blocks. + controlnet_single_block_samples : list, optional + ControlNet block samples for single blocks. + skip_first_layer : bool, default=False + If True, skip the first transformer block. + txt_tokens : int, optional + Number of text tokens. + ip_hidden_states : torch.Tensor, optional + Image prompt hidden states. + first_block : bool, default=False + If True, only process the first block. + skip_block : bool, default=True + If True, skip the first block. + + Returns + ------- + tuple + (hidden_states, encoder_hidden_states, hidden_states_residual, encoder_hidden_states_residual) + """ if first_block and skip_block: raise ValueError("`first_block` and `skip_block` cannot both be True.") @@ -279,6 +385,26 @@ def load_ip_adapter_weights_per_layer( joint_attention_dim: int = 4096, inner_dim: int = 3072, ): + """ + Load per-layer IP-Adapter weights from a HuggingFace Hub repository. + + Parameters + ---------- + repo_id : str + HuggingFace Hub repository ID. + filename : str, default="ip_adapter.safetensors" + Name of the safetensors file. + prefix : str, default="double_blocks." + Prefix for block keys in the file. + joint_attention_dim : int, default=4096 + Input dimension for joint attention. + inner_dim : int, default=3072 + Output dimension for projections. + + Returns + ------- + None + """ path = hf_hub_download(repo_id=repo_id, filename=filename) raw_cpu = {} with safe_open(path, framework="pt", device="cpu") as f: @@ -318,10 +444,39 @@ def load_ip_adapter_weights_per_layer( self.ip_v_projs.append(v_proj) def set_ip_hidden_states(self, image_embeds, negative_image_embeds=None): + """ + Set the image embeddings for IP-Adapter conditioning. + + Parameters + ---------- + image_embeds : torch.Tensor + Image embeddings to use. + negative_image_embeds : optional + Not used. + + Returns + ------- + None + """ self.image_embeds = image_embeds def resize_numpy_image_long(image, resize_long_edge=768): + """ + Resize a numpy image so its longest edge matches a target size. + + Parameters + ---------- + image : np.ndarray + Input image as a numpy array. + resize_long_edge : int, default=768 + Target size for the longest edge. + + Returns + ------- + np.ndarray + Resized image. + """ h, w = image.shape[:2] if max(h, w) <= resize_long_edge: return image @@ -333,6 +488,22 @@ def resize_numpy_image_long(image, resize_long_edge=768): def undo_all_mods_on_transformer(transformer: FluxTransformer2DModel): + """ + Restore a FluxTransformer2DModel to its original, unmodified state. + + This function undoes any modifications made for IP-Adapter integration, + restoring the original forward method and transformer blocks. + + Parameters + ---------- + transformer : FluxTransformer2DModel + The transformer model to restore. + + Returns + ------- + FluxTransformer2DModel + The restored transformer model. + """ if hasattr(transformer, "_original_forward"): transformer.forward = transformer._original_forward del transformer._original_forward diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index bc8d3dfe..9eb1e076 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -8,8 +8,8 @@ from diffusers.utils import load_image from nunchaku import NunchakuFluxTransformer2dModel -from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe -from nunchaku.models.IP_adapter.utils import resize_numpy_image_long +from nunchaku.models.ip_adapter.diffusers_adapters import apply_IPA_on_pipe +from nunchaku.models.ip_adapter.utils import resize_numpy_image_long from nunchaku.pipeline.pipeline_flux_pulid import PuLIDFluxPipeline from nunchaku.utils import get_precision, is_turing From 7ed768525285e78c95b3944a1c905e913ca3eeda Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 06:51:20 +0800 Subject: [PATCH 11/16] docs: add docs for ipa --- ...ku.models.ip_adapter.diffusers_adapters.flux.rst | 7 +++++++ ...unchaku.models.ip_adapter.diffusers_adapters.rst | 13 +++++++++++++ ...u.models.ip_adapter.diffusers_adapters.utils.rst | 7 +++++++ .../python_api/nunchaku.models.ip_adapter.rst | 8 ++++++++ .../python_api/nunchaku.models.ip_adapter.utils.rst | 7 +++++++ docs/source/python_api/nunchaku.models.rst | 1 + 6 files changed, 43 insertions(+) create mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.flux.rst create mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst create mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst create mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.rst create mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.utils.rst diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.flux.rst b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.flux.rst new file mode 100644 index 00000000..bcc8a53c --- /dev/null +++ b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.flux.rst @@ -0,0 +1,7 @@ +nunchaku.models.ip_adapter.diffusers_adapters.flux +================================================== + +.. automodule:: nunchaku.models.ip_adapter.diffusers_adapters.flux + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst new file mode 100644 index 00000000..0907fc87 --- /dev/null +++ b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst @@ -0,0 +1,13 @@ +nunchaku.models.ip_adapter +========================== + +.. automodule:: nunchaku.models.ip_adapter.diffusers_adapters + :members: + :undoc-members: + :show-inheritance: + +.. toctree:: + :maxdepth: 4 + + nunchaku.models.ip_adapter.diffusers_adapters.flux + nunchaku.models.ip_adapter.diffusers_adapters.utils diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst new file mode 100644 index 00000000..527378ed --- /dev/null +++ b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst @@ -0,0 +1,7 @@ +nunchaku.models.ip_adapter.diffusers_adapters.utils +=================================================== + +.. automodule:: nunchaku.models.ip_adapter.diffusers_adapters.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.rst b/docs/source/python_api/nunchaku.models.ip_adapter.rst new file mode 100644 index 00000000..dc1c8935 --- /dev/null +++ b/docs/source/python_api/nunchaku.models.ip_adapter.rst @@ -0,0 +1,8 @@ +nunchaku.models.ip_adapter +========================== + +.. toctree:: + :maxdepth: 4 + + nunchaku.models.ip_adapter.diffusers_adapters + nunchaku.models.ip_adapter.utils diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.utils.rst b/docs/source/python_api/nunchaku.models.ip_adapter.utils.rst new file mode 100644 index 00000000..2a292122 --- /dev/null +++ b/docs/source/python_api/nunchaku.models.ip_adapter.utils.rst @@ -0,0 +1,7 @@ +nunchaku.models.ip_adapter.utils +================================ + +.. automodule:: nunchaku.models.ip_adapter.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/python_api/nunchaku.models.rst b/docs/source/python_api/nunchaku.models.rst index 4074c31b..05d27680 100644 --- a/docs/source/python_api/nunchaku.models.rst +++ b/docs/source/python_api/nunchaku.models.rst @@ -7,4 +7,5 @@ nunchaku.models nunchaku.models.transformers nunchaku.models.text_encoders nunchaku.models.pulid + nunchaku.models.ip_adapter nunchaku.models.safety_checker From b908986f1474e223dc8cd4d56e0d9113256a948e Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 07:15:09 +0800 Subject: [PATCH 12/16] docs: add docs for ipa --- docs/source/index.rst | 1 + docs/source/links/huggingface.txt | 1 + ...u.models.ip_adapter.diffusers_adapters.rst | 5 +-- ...ls.ip_adapter.diffusers_adapters.utils.rst | 7 ---- docs/source/usage/ip_adapter.rst | 38 +++++++++++++++++++ docs/source/usage/pulid.rst | 7 +--- 6 files changed, 44 insertions(+), 15 deletions(-) delete mode 100644 docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst create mode 100644 docs/source/usage/ip_adapter.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index 1804fbf8..745858f2 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -24,6 +24,7 @@ Check out `DeepCompressor `_ for the quantization librar usage/attention.rst usage/fbcache.rst usage/pulid.rst + usage/ip_adapter.rst .. toctree:: :maxdepth: 1 diff --git a/docs/source/links/huggingface.txt b/docs/source/links/huggingface.txt index 7f5e354a..8b916221 100644 --- a/docs/source/links/huggingface.txt +++ b/docs/source/links/huggingface.txt @@ -8,3 +8,4 @@ .. _hf_nunchaku-flux1-dev-int4: https://huggingface.co/mit-han-lab/nunchaku-flux.1-dev/blob/main/svdq-int4_r32-flux.1-dev.safetensors .. _hf_depth_anything: https://huggingface.co/LiheYoung/depth-anything-large-hf .. _hf_nunchaku_wheels: https://huggingface.co/nunchaku-tech/nunchaku +.. _hf_ip-adapterv2: https://huggingface.co/XLabs-AI/flux-ip-adapter-v2 diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst index 0907fc87..2b3afcf7 100644 --- a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst +++ b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.rst @@ -1,5 +1,5 @@ -nunchaku.models.ip_adapter -========================== +nunchaku.models.ip_adapter.diffusers_adapters +============================================= .. automodule:: nunchaku.models.ip_adapter.diffusers_adapters :members: @@ -10,4 +10,3 @@ nunchaku.models.ip_adapter :maxdepth: 4 nunchaku.models.ip_adapter.diffusers_adapters.flux - nunchaku.models.ip_adapter.diffusers_adapters.utils diff --git a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst b/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst deleted file mode 100644 index 527378ed..00000000 --- a/docs/source/python_api/nunchaku.models.ip_adapter.diffusers_adapters.utils.rst +++ /dev/null @@ -1,7 +0,0 @@ -nunchaku.models.ip_adapter.diffusers_adapters.utils -=================================================== - -.. automodule:: nunchaku.models.ip_adapter.diffusers_adapters.utils - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/usage/ip_adapter.rst b/docs/source/usage/ip_adapter.rst new file mode 100644 index 00000000..cc7b49d4 --- /dev/null +++ b/docs/source/usage/ip_adapter.rst @@ -0,0 +1,38 @@ +IP Adapter +========== + +Nunchaku supports `IP Adapter `_, an adapter achieving image prompt capability for the FLUX.1-dev + +.. literalinclude:: ../../../examples/flux.1-dev-IP-adapter.py + :language: python + :caption: IP Adapter Example (`examples/flux.1-dev-IP-adapter.py `__) + :linenos: + +The IP Adapter integration in Nunchaku follows these main steps: + +**Model Initialization**: + +- Load a Nunchaku FLUX.1-dev transformer model using :meth:`~nunchaku.models.transformers.transformer_flux.NunchakuFluxTransformer2dModel.from_pretrained`. +- Initialize the FLUX pipeline with :class:`diffusers.FluxPipeline`, passing the transformer and setting the appropriate precision. + +**IP Adapter Loading**: + +- Use ``pipeline.load_ip_adapter`` to load the IP Adapter weights and the CLIP image encoder. + + - ``pretrained_model_name_or_path_or_dict``: Hugging Face repo or local path for the IP Adapter weights. + - ``weight_name``: Name of the weights file (e.g., ``ip_adapter.safetensors``). + - ``image_encoder_pretrained_model_name_or_path``: Name or path of the CLIP image encoder. + - Apply the IP Adapter to the pipeline with :func:`~nunchaku.models.ip_adapter.diffusers_adapters.apply_IPA_on_pipe`, specifying the adapter scale and repo ID. + +**Caching (Optional)**: + +Enable caching for faster inference and reduced memory usage with :func:`~nunchaku.caching.diffusers_adapters.apply_cache_on_pipe`. See :doc:`fbcache` for more details. + +**Image Generation**: + +- Load the image to be used as the image prompt (IP Adapter reference). +- Call the pipeline with: + + - ``prompt``: The text prompt for generation. + - ``ip_adapter_image``: The reference image (must be RGB). +- The output image will reflect both the text prompt and the visual style/content of the reference image. diff --git a/docs/source/usage/pulid.rst b/docs/source/usage/pulid.rst index 616cc16a..1b03298e 100644 --- a/docs/source/usage/pulid.rst +++ b/docs/source/usage/pulid.rst @@ -1,7 +1,7 @@ PuLID ===== -Nunchaku integrates `PuLID <_pulid_paper>`_, a tuning-free identity customization method for text-to-image generation. +Nunchaku integrates `PuLID `_, a tuning-free identity customization method for text-to-image generation. This feature allows you to generate images that maintain specific identity characteristics from reference photos. .. literalinclude:: ../../../examples/flux.1-dev-pulid.py @@ -9,13 +9,10 @@ This feature allows you to generate images that maintain specific identity chara :caption: PuLID Example (`examples/flux.1-dev-pulid.py `__) :linenos: -Implementation Overview ------------------------ - The PuLID integration follows these key steps: **Model Initialization** (lines 12-20): -Load a Nunchaku FLUX.1-dev model using :class:`~nunchaku.models.transformers.transformer_flux.NunchakuFluxTransformer2dModel` +Load a Nunchaku FLUX.1-dev model using :meth:`~nunchaku.models.transformers.transformer_flux.NunchakuFluxTransformer2dModel.from_pretrained` and initialize the FLUX PuLID pipeline with :class:`~nunchaku.pipeline.pipeline_flux_pulid.PuLIDFluxPipeline`. **Forward Method Override** (line 22): From 4d6f69d952cc0310de24d2df92429327e158db70 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 07:16:39 +0800 Subject: [PATCH 13/16] add an example for pulid --- docs/source/usage/pulid.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/usage/pulid.rst b/docs/source/usage/pulid.rst index 1b03298e..75af1356 100644 --- a/docs/source/usage/pulid.rst +++ b/docs/source/usage/pulid.rst @@ -1,6 +1,8 @@ PuLID ===== +.. image:: https://huggingface.co/datasets/nunchaku-tech/cdn/resolve/main/ComfyUI-nunchaku/workflows/nunchaku-flux.1-dev-pulid.png + Nunchaku integrates `PuLID `_, a tuning-free identity customization method for text-to-image generation. This feature allows you to generate images that maintain specific identity characteristics from reference photos. From ea075fe13f564915e2272afd487e105d9e44c683 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 13:41:46 +0800 Subject: [PATCH 14/16] update --- tests/flux/test_flux_dev_IPA.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index 9eb1e076..53560d92 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -51,13 +51,13 @@ def test_flux_dev_IPA(): # use the pulid pipeline to get the id embedding transformer = NunchakuFluxTransformer2dModel.from_pretrained( - f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors" + f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors", offload=True ) pipeline = PuLIDFluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16, - ).to("cuda") + ) id_image = id_image.convert("RGB") id_image_numpy = np.array(id_image) From 8f7eea1380c937f4b3be390f419141da37a994b2 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 13:42:20 +0800 Subject: [PATCH 15/16] save gpu memory --- tests/flux/test_flux_dev_IPA.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index 53560d92..c702984e 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -42,7 +42,6 @@ def test_flux_dev_IPA(): ip_adapter_image=id_image.convert("RGB"), num_inference_steps=50, ).images[0] - image.save("./tmp.png") del pipeline del transformer @@ -74,6 +73,11 @@ def test_flux_dev_IPA(): print(cosine_similarities) assert cosine_similarities > 0.85 + del pipeline + del transformer + gc.collect() + torch.cuda.empty_cache() + if __name__ == "__main__": test_flux_dev_IPA() From 67bab3f325448d784490862fdb828ac649072962 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Thu, 24 Jul 2025 14:39:42 +0800 Subject: [PATCH 16/16] change the threshold to 0.8 --- tests/flux/test_flux_dev_IPA.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flux/test_flux_dev_IPA.py b/tests/flux/test_flux_dev_IPA.py index c702984e..dbc10018 100644 --- a/tests/flux/test_flux_dev_IPA.py +++ b/tests/flux/test_flux_dev_IPA.py @@ -71,7 +71,7 @@ def test_flux_dev_IPA(): F.cosine_similarity(id_embeddings.view(32, 2048), output_id_embeddings.view(32, 2048), dim=1).mean().item() ) print(cosine_similarities) - assert cosine_similarities > 0.85 + assert cosine_similarities > 0.80 del pipeline del transformer