Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import math | |
| import re | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torchvision.utils import make_grid | |
| # 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 resize_numpy_image_area(image, area=512 * 512): | |
| h, w = image.shape[:2] | |
| k = math.sqrt(area / (h * w)) | |
| h = int(h * k) - (int(h * k) % 16) | |
| w = int(w * k) - (int(w * k) % 16) | |
| image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA) | |
| return image | |
| # reference: https://github.com/huggingface/diffusers/pull/9295/files | |
| def convert_flux_lora_to_diffusers(old_state_dict): | |
| new_state_dict = {} | |
| orig_keys = list(old_state_dict.keys()) | |
| def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None): | |
| down_weight = sds_sd.pop(sds_key) | |
| up_weight = sds_sd.pop(sds_key.replace(".down.weight", ".up.weight")) | |
| # calculate dims if not provided | |
| num_splits = len(ait_keys) | |
| if dims is None: | |
| dims = [up_weight.shape[0] // num_splits] * num_splits | |
| else: | |
| assert sum(dims) == up_weight.shape[0] | |
| # make ai-toolkit weight | |
| ait_down_keys = [k + ".lora_A.weight" for k in ait_keys] | |
| ait_up_keys = [k + ".lora_B.weight" for k in ait_keys] | |
| # down_weight is copied to each split | |
| ait_sd.update({k: down_weight for k in ait_down_keys}) | |
| # up_weight is split to each split | |
| ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))}) # noqa: C416 | |
| for old_key in orig_keys: | |
| # Handle double_blocks | |
| if 'double_blocks' in old_key: | |
| block_num = re.search(r"double_blocks_(\d+)", old_key).group(1) | |
| new_key = f"transformer.transformer_blocks.{block_num}" | |
| if "proj_lora1" in old_key: | |
| new_key += ".attn.to_out.0" | |
| elif "proj_lora2" in old_key: | |
| new_key += ".attn.to_add_out" | |
| elif "qkv_lora2" in old_key and "up" not in old_key: | |
| handle_qkv( | |
| old_state_dict, | |
| new_state_dict, | |
| old_key, | |
| [ | |
| f"transformer.transformer_blocks.{block_num}.attn.add_q_proj", | |
| f"transformer.transformer_blocks.{block_num}.attn.add_k_proj", | |
| f"transformer.transformer_blocks.{block_num}.attn.add_v_proj", | |
| ], | |
| ) | |
| # continue | |
| elif "qkv_lora1" in old_key and "up" not in old_key: | |
| handle_qkv( | |
| old_state_dict, | |
| new_state_dict, | |
| old_key, | |
| [ | |
| f"transformer.transformer_blocks.{block_num}.attn.to_q", | |
| f"transformer.transformer_blocks.{block_num}.attn.to_k", | |
| f"transformer.transformer_blocks.{block_num}.attn.to_v", | |
| ], | |
| ) | |
| # continue | |
| if "down" in old_key: | |
| new_key += ".lora_A.weight" | |
| elif "up" in old_key: | |
| new_key += ".lora_B.weight" | |
| # Handle single_blocks | |
| elif 'single_blocks' in old_key: | |
| block_num = re.search(r"single_blocks_(\d+)", old_key).group(1) | |
| new_key = f"transformer.single_transformer_blocks.{block_num}" | |
| if "proj_lora" in old_key: | |
| new_key += ".proj_out" | |
| elif "qkv_lora" in old_key and "up" not in old_key: | |
| handle_qkv( | |
| old_state_dict, | |
| new_state_dict, | |
| old_key, | |
| [ | |
| f"transformer.single_transformer_blocks.{block_num}.attn.to_q", | |
| f"transformer.single_transformer_blocks.{block_num}.attn.to_k", | |
| f"transformer.single_transformer_blocks.{block_num}.attn.to_v", | |
| ], | |
| ) | |
| if "down" in old_key: | |
| new_key += ".lora_A.weight" | |
| elif "up" in old_key: | |
| new_key += ".lora_B.weight" | |
| else: | |
| # Handle other potential key patterns here | |
| new_key = old_key | |
| # Since we already handle qkv above. | |
| if "qkv" not in old_key and 'embedding' not in old_key: | |
| new_state_dict[new_key] = old_state_dict.pop(old_key) | |
| # if len(old_state_dict) > 0: | |
| # raise ValueError(f"`old_state_dict` should be at this point but has: {list(old_state_dict.keys())}.") | |
| return new_state_dict | |