commit 3bdb6bdcff6dfa57755d04d5e2296243be618a2d Author: Artur Penner Date: Thu Jan 15 18:07:35 2026 +0100 Initial clean commit for ComfyUI-VIBE diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b908d4c --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +.DS_Store diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..58fa8fe --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..c273d8d --- /dev/null +++ b/nodes.py @@ -0,0 +1,128 @@ +# ------------------------------------------------------------------------------ +# VIBE: Visual Instruction Based Editor implementation for ComfyUI +# Implementation: ato-zen +# Repository: https://github.com/ato-zen/ComfyUI-VIBE +# ------------------------------------------------------------------------------ + +import torch +import numpy as np +from PIL import Image +import os +import sys +import folder_paths + +# Registration of model directory in ComfyUI structure +VIBE_MODELS_DIR = os.path.join(folder_paths.models_dir, "vibe") +if not os.path.exists(VIBE_MODELS_DIR): + os.makedirs(VIBE_MODELS_DIR, exist_ok=True) + +# Add node directory to sys.path for internal module resolution +NODE_ROOT = os.path.dirname(os.path.abspath(__file__)) +if NODE_ROOT not in sys.path: + sys.path.insert(0, NODE_ROOT) + +try: + from vibe.editor import ImageEditor +except ImportError as e: + print(f"VIBE Load Error: Internal 'vibe' module not found. {e}") + ImageEditor = None + +VIBE_MODEL_INSTANCE = None +VIBE_CURRENT_PATH = None + +class VibeEditorNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive": ("STRING", {"multiline": True, "default": "make it blue"}), + "steps": ("INT", {"default": 20, "min": 1, "max": 100}), + "cfg_text": ("FLOAT", {"default": 4.5, "min": 1.0, "max": 20.0, "step": 0.1}), + "cfg_image": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 5.0, "step": 0.1}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "image": ("IMAGE",), + "latent_image": ("LATENT",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate" + CATEGORY = "VIBE" + + def _auto_find_model(self): + """Finds the first folder in models/vibe containing model_index.json""" + if not os.path.exists(VIBE_MODELS_DIR): + return None + for entry in os.listdir(VIBE_MODELS_DIR): + full_path = os.path.join(VIBE_MODELS_DIR, entry) + if os.path.isdir(full_path) and "model_index.json" in os.listdir(full_path): + return full_path + return None + + def generate(self, positive, steps, cfg_text, cfg_image, seed, image=None, latent_image=None): + global VIBE_MODEL_INSTANCE, VIBE_CURRENT_PATH + + if ImageEditor is None: + raise ImportError("VIBE library missing in custom_nodes folder.") + + checkpoint_path = self._auto_find_model() or "iitolstykh/VIBE-Image-Edit" + + if VIBE_MODEL_INSTANCE is None or VIBE_CURRENT_PATH != checkpoint_path: + print(f"VIBE: Loading model from {checkpoint_path}") + VIBE_MODEL_INSTANCE = ImageEditor(checkpoint_path=checkpoint_path, device="cuda") + VIBE_CURRENT_PATH = checkpoint_path + + # RESOLUTION LOGIC: Priority is Latent > Image > Default + target_w, target_h = 1024, 1024 + + if latent_image is not None: + target_h = latent_image["samples"].shape[2] * 8 + target_w = latent_image["samples"].shape[3] * 8 + elif image is not None: + target_h, target_w = image.shape[1], image.shape[2] + + # IMAGE CONDITIONING LOGIC + input_pil = None + if image is not None: + batch_data = 255. * image[0].cpu().numpy() + input_pil = Image.fromarray(np.clip(batch_data, 0, 255).astype(np.uint8)).convert("RGB") + + # Auto-resize input image if latent_image provided a different target resolution + if input_pil.size != (target_w, target_h): + print(f"VIBE: Resizing input image to match latent dimensions: {target_w}x{target_h}") + input_pil = input_pil.resize((target_w, target_h), Image.Resampling.LANCZOS) + + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + print(f"VIBE: Processing '{positive}' at {target_w}x{target_h}") + + result = VIBE_MODEL_INSTANCE.generate_edited_image( + instruction=positive, + conditioning_image=input_pil, + num_inference_steps=steps, + guidance_scale=cfg_text, + image_guidance_scale=cfg_image, + seed=seed, + randomize_seed=False, + t2i_width=target_w, + t2i_height=target_h, + num_images_per_prompt=1 + ) + + final_pil = result if isinstance(result, Image.Image) else result[0] + final_pil = final_pil.convert("RGB") + + if final_pil.size != (target_w, target_h): + final_pil = final_pil.resize((target_w, target_h), Image.Resampling.LANCZOS) + + out_tensor = np.array(final_pil).astype(np.float32) / 255.0 + out_tensor = torch.from_numpy(out_tensor)[None, ] + + return (out_tensor,) + +NODE_CLASS_MAPPINGS = { "VIBE_Editor": VibeEditorNode } +NODE_DISPLAY_NAME_MAPPINGS = { "VIBE_Editor": "VIBE Image Editor" } diff --git a/vibe/__init__.py b/vibe/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/vibe/editor.py b/vibe/editor.py new file mode 100644 index 0000000..4691b30 --- /dev/null +++ b/vibe/editor.py @@ -0,0 +1,171 @@ +"""Image editor class.""" + +import random + +import numpy as np +import torch +from loguru import logger +from PIL import Image + +from vibe.generative_pipeline import VIBESanaEditingPipeline +from vibe.utils import retry_decorator +from vibe.utils.img_utils import get_multiscale_transform, postprocess_padded_image, revert_resize + +MAX_SEED = np.iinfo(np.int32).max + + +def randomize_seed_fn(seed: int, *, randomize_seed: bool = False) -> int: + """Randomize the seed. + + Args: + seed (int): Seed. + randomize_seed (bool): Whether to randomize the seed. + + Returns: + int: Randomized seed. + """ + if randomize_seed: + seed = random.randint(0, MAX_SEED) # noqa: S311 + return seed + + +class ImageEditor: + """Image editor class.""" + + def __init__( + self, + checkpoint_path: str, + image_guidance_scale: float = 1.2, + guidance_scale: float = 4.5, + num_inference_steps: int = 20, + device: str = "cuda:0", + **kwargs, + ) -> None: + """Initialize the image editor. + + Args: + checkpoint_path (str): The path to the local model checkpoint. + image_guidance_scale (float): The image guidance scale. + guidance_scale (float): The guidance scale. + num_inference_steps (int): The number of inference steps. + device (str): The device to use. + """ + self.weight_dtype = torch.bfloat16 + self.device = device + self.num_inference_steps = num_inference_steps + self.image_guidance_scale = image_guidance_scale + self.guidance_scale = guidance_scale + self.get_generation_pipe(checkpoint_path, **kwargs) + + @retry_decorator(logger=logger, delay=1) + def get_generation_pipe(self, checkpoint_path: str, **kwargs) -> None: + """Get the generation pipe. + + Args: + checkpoint_path (str): The path to the pipeline checkpoint. + """ + self.pipe = VIBESanaEditingPipeline.from_pretrained( + checkpoint_path, torch_dtype=self.weight_dtype, **kwargs + ).to(self.device) + + def prepare_image_for_diffusion(self, image: Image.Image) -> tuple[Image.Image, tuple[int, ...] | None]: + """Prepare the image for diffusion. + + Args: + image (Image.Image): The input image. + + Returns: + tuple[Image.Image, tuple[int, ...] | None]: The prepared image and the crop coordinates. + """ + inp_image = image.copy() + ori_h, ori_w = inp_image.height, inp_image.width + + # multiscale when we need to calculate the closest aspect ratio and resize & crop image + closest_size = self.pipe.get_closest_size(ori_h, ori_w) + + transform, crop_coords = get_multiscale_transform( + closest_size=closest_size, + orig_size=(ori_h, ori_w), + resize_factor=self.pipe.vae_scale_factor, + out_img_type="pil", + ) + inp_image = transform(inp_image) + return inp_image, crop_coords + + @torch.inference_mode() + def generate_edited_image( # noqa: PLR0913 + self, + instruction: str | list[str], + *, + conditioning_image: Image.Image | None = None, + randomize_seed: bool = False, + do_revert_resize: bool = True, + seed: int = 42, + num_images_per_prompt: int = 1, + t2i_height: int | None = None, + t2i_width: int | None = None, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + image_guidance_scale: float | None = None, + ) -> list[Image.Image]: + """Generate an image based on the provided image and text prompt. + + Args: + instruction (str | list[str]): Generation instruction or list of instructions. + conditioning_image (Image.Image | None): Input image to edit or None if t2i generation is enabled. + randomize_seed (bool): Whether to randomize the seed. + do_revert_resize (bool): resize generated image to initial size. + seed (int): Seed for random number generation. Defaults to 0. + num_images_per_prompt (int): number of images to generate for each prompt + t2i_height (int | None): The height of the t2i image. + t2i_width (int | None): The width of the t2i image. + image_guidance_scale (float | None): The image guidance scale. + guidance_scale (float | None): The guidance scale. + num_inference_steps (int | None): The number of inference steps. + + Returns: + list[Image.Image]: The edited images. + """ + # Randomize seed + seed = randomize_seed_fn(seed, randomize_seed=randomize_seed) + + # Prepare input image for generation + if conditioning_image is not None: + orig_height, orig_width = conditioning_image.height, conditioning_image.width + inp_image, crop_coords = self.prepare_image_for_diffusion(conditioning_image) + inp_height, inp_width = inp_image.size[1], inp_image.size[0] + else: + if t2i_height is None or t2i_width is None: + logger.warning("Height and width for t2i generation are not provided, using default value 1024x1024") + t2i_height = t2i_width = 1024 + inp_height, inp_width = self.pipe.get_closest_size(t2i_height, t2i_width) + if (inp_height, inp_width) != (t2i_height, t2i_width): + logger.warning( + f"The desired size of t2i generation {t2i_height}x{t2i_width} " + f"were adjusted to {inp_height}x{inp_width} to follow default binning scheme." + ) + + # Generate edited image + generated_images = self.pipe( # type: ignore + conditioning_image=conditioning_image, + prompt=instruction, + height=inp_height, + width=inp_width, + num_inference_steps=self.num_inference_steps if num_inference_steps is None else num_inference_steps, + num_images_per_prompt=num_images_per_prompt, + guidance_scale=self.guidance_scale if guidance_scale is None else guidance_scale, + image_guidance_scale=self.image_guidance_scale if image_guidance_scale is None else image_guidance_scale, + output_type="pil", + generator=torch.Generator(device=self.device).manual_seed(seed), + ).images + + # Return generated images if t2i generation is enabled + if conditioning_image is None: + return generated_images # type: ignore[return-value] + + # remove paddings in case of letterbox usage + generated_images = [postprocess_padded_image(im, crop_coords) for im in generated_images] # type: ignore + + if do_revert_resize: + return [revert_resize(im, (orig_height, orig_width)) for im in generated_images] + return generated_images diff --git a/vibe/generative_pipeline/__init__.py b/vibe/generative_pipeline/__init__.py new file mode 100644 index 0000000..a3e5aa1 --- /dev/null +++ b/vibe/generative_pipeline/__init__.py @@ -0,0 +1,5 @@ +"""Generative pipeline.""" + +from .vibe_sana_pipeline import VIBESanaEditingPipeline + +__all__ = ["VIBESanaEditingPipeline"] diff --git a/vibe/generative_pipeline/aspects_multiscale.py b/vibe/generative_pipeline/aspects_multiscale.py new file mode 100644 index 0000000..29056c1 --- /dev/null +++ b/vibe/generative_pipeline/aspects_multiscale.py @@ -0,0 +1,621 @@ + +ASPECT_RATIO_384 = { + "0.25": [192.0, 768.0], + "0.26": [192.0, 744.0], + "0.27": [192.0, 720.0], + "0.28": [192.0, 696.0], + "0.32": [216.0, 672.0], + "0.33": [216.0, 648.0], + "0.35": [216.0, 624.0], + "0.4": [240.0, 600.0], + "0.42": [240.0, 576.0], + "0.48": [264.0, 552.0], + "0.5": [264.0, 528.0], + "0.52": [264.0, 504.0], + "0.57": [288.0, 504.0], + "0.6": [288.0, 480.0], + "0.68": [312.0, 456.0], + "0.72": [312.0, 432.0], + "0.78": [336.0, 432.0], + "0.82": [336.0, 408.0], + "0.88": [360.0, 408.0], + "0.94": [360.0, 384.0], + "1.0": [384.0, 384.0], + "1.07": [384.0, 360.0], + "1.13": [408.0, 360.0], + "1.21": [408.0, 336.0], + "1.29": [432.0, 336.0], + "1.38": [432.0, 312.0], + "1.46": [456.0, 312.0], + "1.67": [480.0, 288.0], + "1.75": [504.0, 288.0], + "2.0": [528.0, 264.0], + "2.09": [552.0, 264.0], + "2.4": [576.0, 240.0], + "2.5": [600.0, 240.0], + "2.89": [624.0, 216.0], + "3.0": [648.0, 216.0], + "3.11": [672.0, 216.0], + "3.62": [696.0, 192.0], + "3.75": [720.0, 192.0], + "3.88": [744.0, 192.0], + "4.0": [768.0, 192.0], +} + +ASPECT_RATIO_512 = { + "0.25": [256.0, 1024.0], + "0.26": [256.0, 992.0], + "0.27": [256.0, 960.0], + "0.28": [256.0, 928.0], + "0.32": [288.0, 896.0], + "0.33": [288.0, 864.0], + "0.35": [288.0, 832.0], + "0.4": [320.0, 800.0], + "0.42": [320.0, 768.0], + "0.48": [352.0, 736.0], + "0.5": [352.0, 704.0], + "0.52": [352.0, 672.0], + "0.57": [384.0, 672.0], + "0.6": [384.0, 640.0], + "0.68": [416.0, 608.0], + "0.72": [416.0, 576.0], + "0.78": [448.0, 576.0], + "0.82": [448.0, 544.0], + "0.88": [480.0, 544.0], + "0.94": [480.0, 512.0], + "1.0": [512.0, 512.0], + "1.07": [512.0, 480.0], + "1.13": [544.0, 480.0], + "1.21": [544.0, 448.0], + "1.29": [576.0, 448.0], + "1.38": [576.0, 416.0], + "1.46": [608.0, 416.0], + "1.67": [640.0, 384.0], + "1.75": [672.0, 384.0], + "2.0": [704.0, 352.0], + "2.09": [736.0, 352.0], + "2.4": [768.0, 320.0], + "2.5": [800.0, 320.0], + "2.89": [832.0, 288.0], + "3.0": [864.0, 288.0], + "3.11": [896.0, 288.0], + "3.62": [928.0, 256.0], + "3.75": [960.0, 256.0], + "3.88": [992.0, 256.0], + "4.0": [1024.0, 256.0], +} + +ASPECT_RATIO_768 = { + "0.25": [384.0, 1536.0], + "0.26": [384.0, 1488.0], + "0.27": [384.0, 1440.0], + "0.28": [384.0, 1392.0], + "0.32": [432.0, 1344.0], + "0.33": [432.0, 1296.0], + "0.35": [432.0, 1248.0], + "0.4": [480.0, 1200.0], + "0.42": [480.0, 1152.0], + "0.48": [528.0, 1104.0], + "0.5": [528.0, 1056.0], + "0.52": [528.0, 1008.0], + "0.57": [576.0, 1008.0], + "0.6": [576.0, 960.0], + "0.68": [624.0, 912.0], + "0.72": [624.0, 864.0], + "0.78": [672.0, 864.0], + "0.82": [672.0, 816.0], + "0.88": [720.0, 816.0], + "0.94": [720.0, 768.0], + "1.0": [768.0, 768.0], + "1.07": [768.0, 720.0], + "1.13": [816.0, 720.0], + "1.21": [816.0, 672.0], + "1.29": [864.0, 672.0], + "1.38": [864.0, 624.0], + "1.46": [912.0, 624.0], + "1.67": [960.0, 576.0], + "1.75": [1008.0, 576.0], + "2.0": [1056.0, 528.0], + "2.09": [1104.0, 528.0], + "2.4": [1152.0, 480.0], + "2.5": [1200.0, 480.0], + "2.89": [1248.0, 432.0], + "3.0": [1296.0, 432.0], + "3.11": [1344.0, 432.0], + "3.62": [1392.0, 384.0], + "3.75": [1440.0, 384.0], + "3.88": [1488.0, 384.0], + "4.0": [1536.0, 384.0], +} + +ASPECT_RATIO_896 = { + "0.25": [448.0, 1792.0], + "0.26": [448.0, 1736.0], + "0.27": [448.0, 1680.0], + "0.28": [448.0, 1624.0], + "0.32": [504.0, 1568.0], + "0.33": [504.0, 1512.0], + "0.35": [504.0, 1456.0], + "0.4": [560.0, 1400.0], + "0.42": [560.0, 1344.0], + "0.48": [616.0, 1288.0], + "0.5": [616.0, 1232.0], + "0.52": [616.0, 1176.0], + "0.57": [672.0, 1176.0], + "0.6": [672.0, 1120.0], + "0.68": [728.0, 1064.0], + "0.72": [728.0, 1008.0], + "0.78": [784.0, 1008.0], + "0.82": [784.0, 952.0], + "0.88": [840.0, 952.0], + "0.94": [840.0, 896.0], + "1.0": [896.0, 896.0], + "1.07": [896.0, 840.0], + "1.13": [952.0, 840.0], + "1.21": [952.0, 784.0], + "1.29": [1008.0, 784.0], + "1.38": [1008.0, 728.0], + "1.46": [1064.0, 728.0], + "1.67": [1120.0, 672.0], + "1.75": [1176.0, 672.0], + "2.0": [1232.0, 616.0], + "2.09": [1288.0, 616.0], + "2.4": [1344.0, 560.0], + "2.5": [1400.0, 560.0], + "2.89": [1456.0, 504.0], + "3.0": [1512.0, 504.0], + "3.11": [1568.0, 504.0], + "3.62": [1624.0, 448.0], + "3.75": [1680.0, 448.0], + "3.88": [1736.0, 448.0], + "4.0": [1792.0, 448.0], +} + +ASPECT_RATIO_1024 = { + "0.25": [512.0, 2048.0], + "0.26": [512.0, 1984.0], + "0.27": [512.0, 1920.0], + "0.28": [512.0, 1856.0], + "0.32": [576.0, 1792.0], + "0.33": [576.0, 1728.0], + "0.35": [576.0, 1664.0], + "0.4": [640.0, 1600.0], + "0.42": [640.0, 1536.0], + "0.48": [704.0, 1472.0], + "0.5": [704.0, 1408.0], + "0.52": [704.0, 1344.0], + "0.57": [768.0, 1344.0], + "0.6": [768.0, 1280.0], + "0.68": [832.0, 1216.0], + "0.72": [832.0, 1152.0], + "0.78": [896.0, 1152.0], + "0.82": [896.0, 1088.0], + "0.88": [960.0, 1088.0], + "0.94": [960.0, 1024.0], + "1.0": [1024.0, 1024.0], + "1.07": [1024.0, 960.0], + "1.13": [1088.0, 960.0], + "1.21": [1088.0, 896.0], + "1.29": [1152.0, 896.0], + "1.38": [1152.0, 832.0], + "1.46": [1216.0, 832.0], + "1.67": [1280.0, 768.0], + "1.75": [1344.0, 768.0], + "2.0": [1408.0, 704.0], + "2.09": [1472.0, 704.0], + "2.4": [1536.0, 640.0], + "2.5": [1600.0, 640.0], + "2.89": [1664.0, 576.0], + "3.0": [1728.0, 576.0], + "3.11": [1792.0, 576.0], + "3.62": [1856.0, 512.0], + "3.75": [1920.0, 512.0], + "3.88": [1984.0, 512.0], + "4.0": [2048.0, 512.0], +} + +ASPECT_RATIO_1152 = { + "0.25": [576.0, 2304.0], + "0.26": [576.0, 2232.0], + "0.27": [576.0, 2160.0], + "0.28": [576.0, 2088.0], + "0.32": [648.0, 2016.0], + "0.33": [648.0, 1944.0], + "0.35": [648.0, 1872.0], + "0.4": [720.0, 1800.0], + "0.42": [720.0, 1728.0], + "0.48": [792.0, 1656.0], + "0.5": [792.0, 1584.0], + "0.52": [792.0, 1512.0], + "0.57": [864.0, 1512.0], + "0.6": [864.0, 1440.0], + "0.68": [936.0, 1368.0], + "0.72": [936.0, 1296.0], + "0.78": [1008.0, 1296.0], + "0.82": [1008.0, 1224.0], + "0.88": [1080.0, 1224.0], + "0.94": [1080.0, 1152.0], + "1.0": [1152.0, 1152.0], + "1.07": [1152.0, 1080.0], + "1.13": [1224.0, 1080.0], + "1.21": [1224.0, 1008.0], + "1.29": [1296.0, 1008.0], + "1.38": [1296.0, 936.0], + "1.46": [1368.0, 936.0], + "1.67": [1440.0, 864.0], + "1.75": [1512.0, 864.0], + "2.0": [1584.0, 792.0], + "2.09": [1656.0, 792.0], + "2.4": [1728.0, 720.0], + "2.5": [1800.0, 720.0], + "2.89": [1872.0, 648.0], + "3.0": [1944.0, 648.0], + "3.11": [2016.0, 648.0], + "3.62": [2088.0, 576.0], + "3.75": [2160.0, 576.0], + "3.88": [2232.0, 576.0], + "4.0": [2304.0, 576.0], +} + +ASPECT_RATIO_1280 = { + "0.25": [640.0, 2560.0], + "0.26": [640.0, 2480.0], + "0.27": [640.0, 2400.0], + "0.28": [640.0, 2320.0], + "0.32": [720.0, 2240.0], + "0.33": [720.0, 2160.0], + "0.35": [720.0, 2080.0], + "0.4": [800.0, 2000.0], + "0.42": [800.0, 1920.0], + "0.48": [880.0, 1840.0], + "0.5": [880.0, 1760.0], + "0.52": [880.0, 1680.0], + "0.57": [960.0, 1680.0], + "0.6": [960.0, 1600.0], + "0.68": [1040.0, 1520.0], + "0.72": [1040.0, 1440.0], + "0.78": [1120.0, 1440.0], + "0.82": [1120.0, 1360.0], + "0.88": [1200.0, 1360.0], + "0.94": [1200.0, 1280.0], + "1.0": [1280.0, 1280.0], + "1.07": [1280.0, 1200.0], + "1.13": [1360.0, 1200.0], + "1.21": [1360.0, 1120.0], + "1.29": [1440.0, 1120.0], + "1.38": [1440.0, 1040.0], + "1.46": [1520.0, 1040.0], + "1.67": [1600.0, 960.0], + "1.75": [1680.0, 960.0], + "2.0": [1760.0, 880.0], + "2.09": [1840.0, 880.0], + "2.4": [1920.0, 800.0], + "2.5": [2000.0, 800.0], + "2.89": [2080.0, 720.0], + "3.0": [2160.0, 720.0], + "3.11": [2240.0, 720.0], + "3.62": [2320.0, 640.0], + "3.75": [2400.0, 640.0], + "3.88": [2480.0, 640.0], + "4.0": [2560.0, 640.0], +} + +ASPECT_RATIO_1408 = { + "0.25": [704.0, 2816.0], + "0.26": [704.0, 2728.0], + "0.27": [704.0, 2640.0], + "0.28": [704.0, 2552.0], + "0.32": [792.0, 2464.0], + "0.33": [792.0, 2376.0], + "0.35": [792.0, 2288.0], + "0.4": [880.0, 2200.0], + "0.42": [880.0, 2112.0], + "0.48": [968.0, 2024.0], + "0.5": [968.0, 1936.0], + "0.52": [968.0, 1848.0], + "0.57": [1056.0, 1848.0], + "0.6": [1056.0, 1760.0], + "0.68": [1144.0, 1672.0], + "0.72": [1144.0, 1584.0], + "0.78": [1232.0, 1584.0], + "0.82": [1232.0, 1496.0], + "0.88": [1320.0, 1496.0], + "0.94": [1320.0, 1408.0], + "1.0": [1408.0, 1408.0], + "1.07": [1408.0, 1320.0], + "1.13": [1496.0, 1320.0], + "1.21": [1496.0, 1232.0], + "1.29": [1584.0, 1232.0], + "1.38": [1584.0, 1144.0], + "1.46": [1672.0, 1144.0], + "1.67": [1760.0, 1056.0], + "1.75": [1848.0, 1056.0], + "2.0": [1936.0, 968.0], + "2.09": [2024.0, 968.0], + "2.4": [2112.0, 880.0], + "2.5": [2200.0, 880.0], + "2.89": [2288.0, 792.0], + "3.0": [2376.0, 792.0], + "3.11": [2464.0, 792.0], + "3.62": [2552.0, 704.0], + "3.75": [2640.0, 704.0], + "3.88": [2728.0, 704.0], + "4.0": [2816.0, 704.0], +} + +ASPECT_RATIO_2048 = { + "0.25": [1024.0, 4096.0], + "0.26": [1024.0, 3968.0], + "0.27": [1024.0, 3840.0], + "0.28": [1024.0, 3712.0], + "0.32": [1152.0, 3584.0], + "0.33": [1152.0, 3456.0], + "0.35": [1152.0, 3328.0], + "0.4": [1280.0, 3200.0], + "0.42": [1280.0, 3072.0], + "0.48": [1408.0, 2944.0], + "0.5": [1408.0, 2816.0], + "0.52": [1408.0, 2688.0], + "0.57": [1536.0, 2688.0], + "0.6": [1536.0, 2560.0], + "0.68": [1664.0, 2432.0], + "0.72": [1664.0, 2304.0], + "0.78": [1792.0, 2304.0], + "0.82": [1792.0, 2176.0], + "0.88": [1920.0, 2176.0], + "0.94": [1920.0, 2048.0], + "1.0": [2048.0, 2048.0], + "1.07": [2048.0, 1920.0], + "1.13": [2176.0, 1920.0], + "1.21": [2176.0, 1792.0], + "1.29": [2304.0, 1792.0], + "1.38": [2304.0, 1664.0], + "1.46": [2432.0, 1664.0], + "1.67": [2560.0, 1536.0], + "1.75": [2688.0, 1536.0], + "2.0": [2816.0, 1408.0], + "2.09": [2944.0, 1408.0], + "2.4": [3072.0, 1280.0], + "2.5": [3200.0, 1280.0], + "2.89": [3328.0, 1152.0], + "3.0": [3456.0, 1152.0], + "3.11": [3584.0, 1152.0], + "3.62": [3712.0, 1024.0], + "3.75": [3840.0, 1024.0], + "3.88": [3968.0, 1024.0], + "4.0": [4096.0, 1024.0], +} + +ASPECT_RATIO_1440 = { + "0.25": [720.0, 2880.0], + "0.26": [720.0, 2790.0], + "0.27": [720.0, 2700.0], + "0.28": [720.0, 2610.0], + "0.32": [810.0, 2520.0], + "0.33": [810.0, 2430.0], + "0.35": [810.0, 2340.0], + "0.4": [900.0, 2250.0], + "0.42": [900.0, 2160.0], + "0.48": [990.0, 2070.0], + "0.5": [990.0, 1980.0], + "0.52": [990.0, 1890.0], + "0.57": [1080.0, 1890.0], + "0.6": [1080.0, 1800.0], + "0.68": [1170.0, 1710.0], + "0.72": [1170.0, 1620.0], + "0.78": [1260.0, 1620.0], + "0.82": [1260.0, 1530.0], + "0.88": [1350.0, 1530.0], + "0.94": [1350.0, 1440.0], + "1.0": [1440.0, 1440.0], + "1.07": [1440.0, 1350.0], + "1.13": [1530.0, 1350.0], + "1.21": [1530.0, 1260.0], + "1.29": [1620.0, 1260.0], + "1.38": [1620.0, 1170.0], + "1.46": [1710.0, 1170.0], + "1.67": [1800.0, 1080.0], + "1.75": [1890.0, 1080.0], + "2.0": [1980.0, 990.0], + "2.09": [2070.0, 990.0], + "2.4": [2160.0, 900.0], + "2.5": [2250.0, 900.0], + "2.89": [2340.0, 810.0], + "3.0": [2430.0, 810.0], + "3.11": [2520.0, 810.0], + "3.62": [2610.0, 720.0], + "3.75": [2700.0, 720.0], + "3.88": [2790.0, 720.0], + "4.0": [2880.0, 720.0], +} + +ASPECT_RATIO_1536 = { + "0.25": [768.0, 3072.0], + "0.26": [768.0, 2976.0], + "0.27": [768.0, 2880.0], + "0.28": [768.0, 2784.0], + "0.32": [864.0, 2688.0], + "0.33": [864.0, 2592.0], + "0.35": [864.0, 2496.0], + "0.4": [960.0, 2400.0], + "0.42": [960.0, 2304.0], + "0.48": [1056.0, 2208.0], + "0.5": [1056.0, 2112.0], + "0.52": [1056.0, 2016.0], + "0.57": [1152.0, 2016.0], + "0.6": [1152.0, 1920.0], + "0.68": [1248.0, 1824.0], + "0.72": [1248.0, 1728.0], + "0.78": [1344.0, 1728.0], + "0.82": [1344.0, 1632.0], + "0.88": [1440.0, 1632.0], + "0.94": [1440.0, 1536.0], + "1.0": [1536.0, 1536.0], + "1.07": [1536.0, 1440.0], + "1.13": [1632.0, 1440.0], + "1.21": [1632.0, 1344.0], + "1.29": [1728.0, 1344.0], + "1.38": [1728.0, 1248.0], + "1.46": [1824.0, 1248.0], + "1.67": [1920.0, 1152.0], + "1.75": [2016.0, 1152.0], + "2.0": [2112.0, 1056.0], + "2.09": [2208.0, 1056.0], + "2.4": [2304.0, 960.0], + "2.5": [2400.0, 960.0], + "2.89": [2496.0, 864.0], + "3.0": [2592.0, 864.0], + "3.11": [2688.0, 864.0], + "3.62": [2784.0, 768.0], + "3.75": [2880.0, 768.0], + "3.88": [2976.0, 768.0], + "4.0": [3072.0, 768.0], +} + +ASPECT_RATIO_1664 = { + "0.25": [832.0, 3328.0], + "0.26": [832.0, 3224.0], + "0.27": [832.0, 3120.0], + "0.28": [832.0, 3016.0], + "0.32": [936.0, 2912.0], + "0.33": [936.0, 2808.0], + "0.35": [936.0, 2704.0], + "0.4": [1040.0, 2600.0], + "0.42": [1040.0, 2496.0], + "0.48": [1144.0, 2392.0], + "0.5": [1144.0, 2288.0], + "0.52": [1144.0, 2184.0], + "0.57": [1248.0, 2184.0], + "0.6": [1248.0, 2080.0], + "0.68": [1352.0, 1976.0], + "0.72": [1352.0, 1872.0], + "0.78": [1456.0, 1872.0], + "0.82": [1456.0, 1768.0], + "0.88": [1560.0, 1768.0], + "0.94": [1560.0, 1664.0], + "1.0": [1664.0, 1664.0], + "1.07": [1664.0, 1560.0], + "1.13": [1768.0, 1560.0], + "1.21": [1768.0, 1456.0], + "1.29": [1872.0, 1456.0], + "1.38": [1872.0, 1352.0], + "1.46": [1976.0, 1352.0], + "1.67": [2080.0, 1248.0], + "1.75": [2184.0, 1248.0], + "2.0": [2288.0, 1144.0], + "2.09": [2392.0, 1144.0], + "2.4": [2496.0, 1040.0], + "2.5": [2600.0, 1040.0], + "2.89": [2704.0, 936.0], + "3.0": [2808.0, 936.0], + "3.11": [2912.0, 936.0], + "3.62": [3016.0, 832.0], + "3.75": [3120.0, 832.0], + "3.88": [3224.0, 832.0], + "4.0": [3328.0, 832.0], +} + +ASPECT_RATIO_1792 = { + "0.25": [896.0, 3584.0], + "0.26": [896.0, 3472.0], + "0.27": [896.0, 3360.0], + "0.28": [896.0, 3248.0], + "0.32": [1008.0, 3136.0], + "0.33": [1008.0, 3024.0], + "0.35": [1008.0, 2912.0], + "0.4": [1120.0, 2800.0], + "0.42": [1120.0, 2688.0], + "0.48": [1232.0, 2576.0], + "0.5": [1232.0, 2464.0], + "0.52": [1232.0, 2352.0], + "0.57": [1344.0, 2352.0], + "0.6": [1344.0, 2240.0], + "0.68": [1456.0, 2128.0], + "0.72": [1456.0, 2016.0], + "0.78": [1568.0, 2016.0], + "0.82": [1568.0, 1904.0], + "0.88": [1680.0, 1904.0], + "0.94": [1680.0, 1792.0], + "1.0": [1792.0, 1792.0], + "1.07": [1792.0, 1680.0], + "1.13": [1904.0, 1680.0], + "1.21": [1904.0, 1568.0], + "1.29": [2016.0, 1568.0], + "1.38": [2016.0, 1456.0], + "1.46": [2128.0, 1456.0], + "1.67": [2240.0, 1344.0], + "1.75": [2352.0, 1344.0], + "2.0": [2464.0, 1232.0], + "2.09": [2576.0, 1232.0], + "2.4": [2688.0, 1120.0], + "2.5": [2800.0, 1120.0], + "2.89": [2912.0, 1008.0], + "3.0": [3024.0, 1008.0], + "3.11": [3136.0, 1008.0], + "3.62": [3248.0, 896.0], + "3.75": [3360.0, 896.0], + "3.88": [3472.0, 896.0], + "4.0": [3584.0, 896.0], +} + +ASPECT_RATIO_1920 = { + "0.25": [960.0, 3840.0], + "0.26": [960.0, 3720.0], + "0.27": [960.0, 3600.0], + "0.28": [960.0, 3480.0], + "0.32": [1080.0, 3360.0], + "0.33": [1080.0, 3240.0], + "0.35": [1080.0, 3120.0], + "0.4": [1200.0, 3000.0], + "0.42": [1200.0, 2880.0], + "0.48": [1320.0, 2760.0], + "0.5": [1320.0, 2640.0], + "0.52": [1320.0, 2520.0], + "0.57": [1440.0, 2520.0], + "0.6": [1440.0, 2400.0], + "0.68": [1560.0, 2280.0], + "0.72": [1560.0, 2160.0], + "0.78": [1680.0, 2160.0], + "0.82": [1680.0, 2040.0], + "0.88": [1800.0, 2040.0], + "0.94": [1800.0, 1920.0], + "1.0": [1920.0, 1920.0], + "1.07": [1920.0, 1800.0], + "1.13": [2040.0, 1800.0], + "1.21": [2040.0, 1680.0], + "1.29": [2160.0, 1680.0], + "1.38": [2160.0, 1560.0], + "1.46": [2280.0, 1560.0], + "1.67": [2400.0, 1440.0], + "1.75": [2520.0, 1440.0], + "2.0": [2640.0, 1320.0], + "2.09": [2760.0, 1320.0], + "2.4": [2880.0, 1200.0], + "2.5": [3000.0, 1200.0], + "2.89": [3120.0, 1080.0], + "3.0": [3240.0, 1080.0], + "3.11": [3360.0, 1080.0], + "3.62": [3480.0, 960.0], + "3.75": [3600.0, 960.0], + "3.88": [3720.0, 960.0], + "4.0": [3840.0, 960.0], +} + +ALL_ASPECT_RATIOS = [ + ASPECT_RATIO_384, + ASPECT_RATIO_512, + ASPECT_RATIO_768, + ASPECT_RATIO_896, + ASPECT_RATIO_1024, + ASPECT_RATIO_1152, + ASPECT_RATIO_1280, + ASPECT_RATIO_1408, + # ASPECT_RATIO_1440, + ASPECT_RATIO_1536, + ASPECT_RATIO_1664, + ASPECT_RATIO_1792, + ASPECT_RATIO_1920, + ASPECT_RATIO_2048, +] + +ALL_RESOLUTIONS = [384, 512, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048] diff --git a/vibe/generative_pipeline/vibe_sana_pipeline.py b/vibe/generative_pipeline/vibe_sana_pipeline.py new file mode 100644 index 0000000..07a224f --- /dev/null +++ b/vibe/generative_pipeline/vibe_sana_pipeline.py @@ -0,0 +1,763 @@ +"""VIBE Sana Editing pipeline.""" + +from typing import Any + +import torch +from diffusers.image_processor import PixArtImageProcessor +from diffusers.models import AutoencoderDC +from diffusers.pipelines.sana.pipeline_output import SanaPipelineOutput +from diffusers.pipelines.sana.pipeline_sana import SanaPipeline, retrieve_timesteps +from diffusers.schedulers import DPMSolverMultistepScheduler +from diffusers.utils import is_torch_xla_available +from loguru import logger +from PIL import Image +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from vibe.generative_pipeline.aspects_multiscale import ASPECT_RATIO_512, ASPECT_RATIO_1024, ASPECT_RATIO_2048 +from vibe.transformer import VIBESanaEditingModel + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm # type: ignore # pylint: disable + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +class VIBESanaEditingPipeline(SanaPipeline): + """Image Editing Pipeline for Sana.""" + + edit_query_template = "what will this image be like if {instruction}" + t2i_query_template = "generate the image by description: {instruction}" + min_pixels = 352 * 352 # min size of image in Qwen3VLProcessor + max_pixels = 672 * 672 # max size of image in Qwen3VLProcessor + bin_512_space = 512 * 512 + bin_1024_space = 1024 * 1024 + bin_2048_space = 2048 * 2048 + + def __init__( + self, + tokenizer: Qwen3VLProcessor, + text_encoder: Qwen3VLForConditionalGeneration, + vae: AutoencoderDC, + transformer: VIBESanaEditingModel, + scheduler: DPMSolverMultistepScheduler, + ) -> None: + """Initialize the SanaIP2P pipeline. + + Args: + tokenizer (Qwen3VLProcessor): The tokenizer for the text encoder. + text_encoder (Qwen3VLForConditionalGeneration): The text encoder for the editing prompt. + vae (AutoencoderDC): The VAE model. + transformer (VIBESanaEditingModel): The editing transformer. + scheduler (DPMSolverMultistepScheduler): The sana scheduler. + """ + super(SanaPipeline, self).__init__() + + self.register_modules( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) + if hasattr(self, "vae") and self.vae is not None + else 32 + ) + self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) + + @property + def do_classifier_free_guidance(self) -> bool: + """Property to check if classifier free guidance is enabled. + + Returns: + bool: True if classifier free guidance is enabled, False otherwise. + """ + return self._guidance_scale > 1 + + @property + def do_image_guidance(self) -> bool: + """Property to check if image guidance is enabled. + + Returns: + bool: True if image guidance is enabled, False otherwise. + """ + return self._image_guidance_scale >= 1 + + def get_closest_size(self, height: int, width: int) -> tuple[int, int]: + """Get the closest size to the original size. + + Args: + height (int): The height of the image. + width (int): The width of the image. + + Returns: + tuple[int, int]: The closest size to the original size. + """ + sample_area = height * width + if sample_area <= self.bin_512_space * 2 + 256**2: + aspect_ratio_bin = ASPECT_RATIO_512 + elif sample_area <= self.bin_1024_space * 2 + 512**2: + aspect_ratio_bin = ASPECT_RATIO_1024 + else: + aspect_ratio_bin = ASPECT_RATIO_2048 + height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) + return height, width + + def prepare_inputs_batch( + self, + instruction: str | list[str], + image: Image.Image | None = None, + ) -> dict[str, torch.Tensor]: + """Prepare inputs for the batch. + + Args: + instruction (str | list[str]): Edit prompts. + image (Image.Image | None): The input image or None if t2i generation is enabled. + + Returns: + dict[str, torch.Tensor]: The inputs for the batch. + - input_ids: The input IDs. + - attention_mask: The attention mask. + - image_grid_thw: The image grid size. + - pixel_values: The pixel values. + """ + # Prepare instructions + instructions = [instruction] if isinstance(instruction, str) else instruction + samples = [] + for instr in instructions: + query = ( + self.edit_query_template.format(instruction=instr.strip()) + if image is not None + else self.t2i_query_template.format(instruction=instr.strip()) + ) + user_content = [{"type": "image", "image": image}] if image is not None else [] + user_content.append({"type": "text", "text": query}) + message = [{"role": "user", "content": user_content}] + text = self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=False) + sample = {"text": text} + if image is not None: + sample["image"] = image + samples.append(sample) + + # Invoke the processor to get the input ids and attention mask + texts = [sample["text"] for sample in samples] + images = [sample["image"] for sample in samples] if image is not None else None + kwargs = {"min_pixels": self.min_pixels, "max_pixels": self.max_pixels} + inputs = self.tokenizer(text=texts, images=images, padding=True, return_tensors="pt", **kwargs).data + return {key: value.to(self._execution_device) for key, value in inputs.items()} + + def prepare_inputs_for_meta_queries( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + image_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Add meta queries into input embeddings before any padding tokens. + + The meta queries (stored in `self.transformer.meta_queries`) are inserted after all valid tokens + (as indicated by attention_mask) but before padding tokens. + The attention mask is updated with ones for the meta queries. + + Args: + inputs_embeds (torch.Tensor): Input embeddings of shape (B, L, d_model). + attention_mask (torch.Tensor): Attention mask of shape (B, L), with 1 for valid tokens and 0 for padding. + image_mask (Optional[torch.Tensor]): Image mask of shape (B, L). + + Returns: + tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + - Updated input embeddings of shape (B, L + N_queries, d_model) + - Updated attention mask of shape (B, L + N_queries) + - Updated image mask of shape (B, L + N_queries) + """ + batch_size, _, _ = inputs_embeds.size() + num_meta_queries = self.transformer.meta_queries.size(0) # Number of meta query tokens + + # Lists to collect updated tensors for each sample. + updated_embeds_list = [] + updated_attention_list = [] + updated_image_mask_list = [] + + for batch_idx in range(batch_size): + # Compute the number of valid tokens (non-padding) using the attention mask. + valid_length = int(attention_mask[batch_idx].sum().item()) + + # Insert meta queries after valid tokens and before any padding tokens. + valid_embeds = inputs_embeds[batch_idx, :valid_length, :] # (valid_length, d_model) + pad_embeds = inputs_embeds[batch_idx, valid_length:, :] # (pad_length, d_model) + updated_embeds = torch.cat([valid_embeds, self.transformer.meta_queries, pad_embeds], dim=0) + updated_embeds_list.append(updated_embeds) + + # Update the attention mask: + valid_mask = attention_mask[batch_idx, :valid_length] # (valid_length,) all ones + pad_mask = attention_mask[batch_idx, valid_length:] # (pad_length,) all zeros + meta_mask = torch.ones(num_meta_queries, device=attention_mask.device, dtype=attention_mask.dtype) + updated_attention = torch.cat([valid_mask, meta_mask, pad_mask], dim=0) + updated_attention_list.append(updated_attention) + + if image_mask is not None: + valid_image_mask = image_mask[batch_idx, :valid_length] # (valid_length,) all ones + pad_image_mask = image_mask[batch_idx, valid_length:] # (pad_length,) all zeros + meta_image_mask = torch.zeros(num_meta_queries, device=image_mask.device, dtype=image_mask.dtype) + updated_image_mask = torch.cat([valid_image_mask, meta_image_mask, pad_image_mask], dim=0) + updated_image_mask_list.append(updated_image_mask) + + # Stack lists to form batch tensors. + updated_embeds = torch.stack(updated_embeds_list, dim=0) # (B, L + N_queries, d_model) + updated_attention_mask = torch.stack(updated_attention_list, dim=0) # (B, L + N_queries) + + if updated_image_mask_list: # noqa: SIM108 + updated_image_mask = torch.stack(updated_image_mask_list, dim=0) # (B, L + N_queries) + else: + updated_image_mask = None # type: ignore + + return updated_embeds, updated_attention_mask, updated_image_mask + + def get_rope_index( + self, + input_ids: torch.LongTensor, + image_grid_thw: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.LongTensor: + """Get the rope index. + + Args: + input_ids (torch.LongTensor): The input IDs of the model. + image_grid_thw (torch.LongTensor | None): The grid size of the image. + attention_mask (torch.Tensor | None): The attention mask of the model. + + Returns: + torch.LongTensor: The position IDs. + """ + attention_mask_tensor = ( + attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"] # type: ignore + ) + if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4: # noqa: PLR2004 + attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2) + # Only apply conversion for floating point tensors (inverted masks) + if attention_mask_tensor.dtype.is_floating_point: + attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min + attention_mask_tensor = (1.0 - attention_mask_tensor).int() + + # Calculate RoPE index once per generation. + position_ids, rope_deltas = self.text_encoder.model.get_rope_index( + input_ids, + image_grid_thw, + None, + attention_mask=attention_mask_tensor, + ) + self.text_encoder.model.rope_deltas = rope_deltas + return position_ids + + def prepare_initial_input_embeddings( + self, + input_ids: torch.LongTensor, + attention_mask: torch.Tensor, + image_grid_thw: torch.LongTensor | None = None, + pixel_values: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Prepare initial input embeddings for the model. + + Args: + input_ids (torch.LongTensor): The input IDs of the model. + attention_mask (torch.Tensor): The attention mask of the model. + image_grid_thw (torch.LongTensor | None): The grid size of the image. + pixel_values (torch.Tensor | None): The pixel values of the image. + + Returns: + dict[str, torch.Tensor]: The model inputs: + - inputs_embeds: The initial input embeddings. + - deepstack_visual_embeds: The deepstack visual embeds or None if no image is provided. + - position_ids: The position ids. + - attention_mask: The attention mask. + - visual_pos_masks: The visual pos masks or None if no image is provided. + """ + inputs_embeds = self.text_encoder.get_input_embeddings()(input_ids) + attention_mask = attention_mask.to(inputs_embeds.device) + attention_mask = attention_mask.bool() + + if pixel_values is not None: + image_embeds, deepstack_visual_embeds = self.text_encoder.get_image_features(pixel_values, image_grid_thw) + image_embeds = torch.cat(image_embeds, dim=0) + image_mask, _ = self.text_encoder.model.get_placeholder_mask( + input_ids, + inputs_embeds=inputs_embeds, + image_features=image_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + image_mask = image_mask[..., 0] + visual_pos_masks = image_mask + else: + visual_pos_masks = None + deepstack_visual_embeds = None + + # Add meta queries to the input embeddings + inputs_embeds, attention_mask, visual_pos_masks = self.prepare_inputs_for_meta_queries( # type: ignore + inputs_embeds, + attention_mask, + image_mask=visual_pos_masks, + ) + + # add placeholder for meta queries, it will be used to create position ids for meta queries only + img_input_ids = torch.full( + (input_ids.shape[0], self.transformer.meta_queries.size(0)), + -1, + device=input_ids.device, + dtype=torch.long, + ) + input_ids = torch.cat([input_ids, img_input_ids], dim=1) # type: ignore + + # computing rope index + position_ids = self.get_rope_index(input_ids, image_grid_thw, attention_mask) + + return { + "inputs_embeds": inputs_embeds, + "deepstack_visual_embeds": deepstack_visual_embeds, + "position_ids": position_ids, + "attention_mask": attention_mask, + "visual_pos_masks": visual_pos_masks, + } + + def _get_editing_embeddings( + self, + edit_prompt: list[str] | str, + conditioning_image: Image.Image | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Get the editing embeddings for the given prompt and image. + + Args: + edit_prompt (list[str] | str): The editing prompt. + conditioning_image (Image.Image | None): The conditioning image or None if t2i generation is enabled. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Tuple of prompt embeds and negative prompt embeds. + """ + # Prepare inputs for the text encoder call. + prepared_inputs = self.prepare_inputs_batch(instruction=edit_prompt, image=conditioning_image) + model_inputs = self.prepare_initial_input_embeddings(**prepared_inputs) # type: ignore[arg-type] + + # Forward the text encoder. + llm_model_result = self.text_encoder.language_model(input_ids=None, **model_inputs) + hidden_states = llm_model_result.last_hidden_state + attention_mask = model_inputs["attention_mask"] + + # Get the hidden states for the meta tokens + meta_hidden_states = self.transformer.get_hidden_states_for_meta_tokens(hidden_states, attention_mask) + + # Forward the edit heads + batch_size = meta_hidden_states.shape[0] + return self.transformer.forward_edit_heads(batch_size, meta_hidden_states) + + def encode_prompt( # type: ignore[override] # noqa: PLR0913 + self, + batch_size: int, + prompt: str | list[str] | None = None, + conditioning_image: Image.Image | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + num_images_per_prompt: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode the prompt into text encoder hidden states. + + Args: + batch_size (int): The batch size. + conditioning_image: Image.Image: The conditioning image to pass in prompt encoder. + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + num_images_per_prompt (`int`, *optional*, defaults to 1): + Number of images that should be generated per prompt. torch device to place the resulting embeddings on + + Returns: + tuple[torch.Tensor, torch.Tensor]: Tuple of prompt embeds and negative prompt embeds. + """ + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self._get_editing_embeddings( + edit_prompt=prompt, # type: ignore + conditioning_image=conditioning_image, + ) + + seq_len = prompt_embeds.size(1) + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + if self.do_classifier_free_guidance: + if negative_prompt_embeds is None: + msg = "Negative prompt embeds are required when classifier free guidance is enabled." + raise ValueError(msg) + seq_len = negative_prompt_embeds.size(1) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + if self.do_image_guidance: + prompt_embeds = torch.cat( # type: ignore + [prompt_embeds, negative_prompt_embeds, negative_prompt_embeds], + dim=0, + ) + else: + prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0) + + return prompt_embeds, negative_prompt_embeds # type: ignore + + def prepare_guidance_inputs( + self, + image_latents: torch.Tensor, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare the guidance inputs for the denoising process. + + Args: + image_latents (torch.Tensor): The image latents to be used for the denoising process. + prompt_embeds (torch.Tensor): The prompt embeds from edit head. + negative_prompt_embeds (torch.Tensor): The negative prompt embeds from edit head. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Tuple of image latents and prompt embeds for the denoising process. + """ + if self.do_classifier_free_guidance: + uncond_image_latents = torch.zeros_like(image_latents) + image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0) + + prompt_embeds = torch.cat( + [ + prompt_embeds, + negative_prompt_embeds, + negative_prompt_embeds, + ], + dim=0, + ) + return image_latents, prompt_embeds + + def prepare_image_latents( + self, + image: torch.Tensor, + batch_size: int, + num_images_per_prompt: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Prepare the coonditioning image latents for the denoising process. + + Args: + image (torch.Tensor): The image to process. + batch_size (int): The batch size. + num_images_per_prompt (int): The number of images to generate per prompt. + dtype (torch.dtype): The dtype of the image. + device (torch.device): The device to use. + """ + image = image.to(device=device, dtype=dtype) + batch_size = batch_size * num_images_per_prompt + + image_latents = self.vae.encode(image).latent + image_latents = image_latents * self.vae.config.scaling_factor + + image_latents = torch.cat([image_latents] * batch_size, dim=0) + + if self.do_classifier_free_guidance: + if self.do_image_guidance: + uncond_image_latents = torch.zeros_like(image_latents) + image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0) + else: + image_latents = torch.cat([image_latents, image_latents], dim=0) + return image_latents + + def _run_denoising_loop( # noqa: PLR0913 + self, + latents: torch.Tensor, + latent_channels: int, + input_image_latents: torch.Tensor | None, + timesteps: torch.Tensor, + num_inference_steps: int, + num_warmup_steps: int, + prompt_embeds: torch.Tensor, + guidance_scale: float, + image_guidance_scale: float, + extra_step_kwargs: dict, + *, + is_t2i: bool = False, + ) -> torch.Tensor: + """Run the denoising loop over the given timesteps. + + Args: + latents (torch.Tensor): Latents to denoise. + latent_channels (int): The number of latent channels. + input_image_latents (torch.Tensor | None): The input image latents or None if t2i generation is enabled. + timesteps (torch.Tensor): The timesteps for the scheduler. + num_inference_steps (int): The number of total inference steps. + num_warmup_steps (int): The number of warmup steps for the scheduler. + prompt_embeds (torch.Tensor): The prompt embeds from edit head. + guidance_scale (float): The guidance scale for the text prompt. + image_guidance_scale (float): The guidance scale for the image prompt. + extra_step_kwargs (dict): The extra step kwargs for the scheduler. + is_t2i (bool): Whether the t2i generation is enabled. + + Returns: + torch.Tensor: The final latents after the denoising loop. + """ + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + if self.do_classifier_free_guidance: + latent_model_input = torch.cat([latents] * (3 if self.do_image_guidance else 2)) + else: + latent_model_input = latents + + # concat noised and input image latents. + if input_image_latents is None: # In case of t2i generation, we don't have input image latents + scaled_latent_model_input = torch.cat([latent_model_input, latent_model_input], dim=1) + else: + scaled_latent_model_input = torch.cat([latent_model_input, input_image_latents], dim=1) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + timestep = timestep * self.transformer.config.timestep_scale + + # predict noise model_output + noise_pred = self.transformer( + hidden_states=scaled_latent_model_input.to(dtype=transformer_dtype), + encoder_hidden_states=prompt_embeds, + timestep=timestep, + return_dict=False, + attention_kwargs=self.attention_kwargs, + t2i_samples=[is_t2i] * scaled_latent_model_input.shape[0], + )[0] + noise_pred = noise_pred.float() + + # perform guidance + if self.do_classifier_free_guidance: + if self.do_image_guidance: + noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3) + noise_pred = ( + noise_pred_uncond + + guidance_scale * (noise_pred_text - noise_pred_image) + + image_guidance_scale * (noise_pred_image - noise_pred_uncond) + ) + else: + noise_pred_text, noise_pred_uncond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + return latents + + @torch.no_grad() + def __call__( # noqa: PLR0913,PLR0912,PLR0915 + self, + conditioning_image: Image.Image | None = None, + prompt: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + latents: torch.FloatTensor | None = None, + height: int | None = None, + width: int | None = None, + eta: float = 0.0, + num_inference_steps: int = 20, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + guidance_scale: float = 4.5, + image_guidance_scale: float = 1.2, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str = "pil", + attention_kwargs: dict[str, Any] | None = None, + *, + return_dict: bool = True, + use_resolution_binning: bool = True, + **_: Any, + ) -> SanaPipelineOutput | tuple: + """Function invoked when calling the pipeline for generation. + + Args: + conditioning_image (Image.Image | None): The input conditioning image or None if t2i generation is enabled. + prompt (str | list[str] | None): The editing prompt. + prompt_embeds (torch.FloatTensor): The prompt embeds from edit head. + negative_prompt_embeds (torch.FloatTensor): The negative prompt embeds from edit head. + latents (Optional[torch.FloatTensor]): The latents to use for the denoising process. + height (Optional[int]): The height of the conditioning_image. + width (Optional[int]): The width of the conditioning_image. + num_inference_steps (int): The number of inference steps. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: huggingface.co/papers/2010.02502. Only + applies to [`schedulers.DDIMScheduler`], will be ignored for others. + guidance_scale (float): The guidance scale for the text prompt. + image_guidance_scale (float): The guidance scale for the conditioning_image. + num_images_per_prompt (int): The number of images to generate per prompt. + generator (Optional[Union[torch.Generator, List[torch.Generator]]]): The generator. + output_type (str): The output type. + attention_kwargs (dict[str, Any]): The attention kwargs. + return_dict (bool): Whether to return a return_dict. + use_resolution_binning (bool): Whether to use resolution binning. + _: Additional keyword arguments. + + Returns: + SanaPipelineOutput: The output of the pipeline. + """ + # check if we need to fallback to t2i generation + is_t2i = False + if conditioning_image is None: + is_t2i = True + logger.warning("Fallback to t2i generation because conditioning_image is not provided.") + + # 0. Set pipeline attributes. + device = self._execution_device + dtype = self.dtype # type: ignore + self._guidance_scale = guidance_scale + self._image_guidance_scale = 0 if is_t2i else image_guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + if (height is None or width is None) and conditioning_image is None: + msg = "Either height and width or conditioning_image must be provided." + raise ValueError(msg) + + # 1. Check inputs. Raise error if not correct + if use_resolution_binning: + if height is None or width is None: + height, width = conditioning_image.height, conditioning_image.width # type: ignore[union-attr] + orig_height, orig_width = height, width + height, width = self.get_closest_size(height, width) + + self.check_inputs( + prompt, + height, + width, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + ) + + # Identify batch size. + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] # type: ignore + + # 1. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( # type: ignore[assignment] + batch_size=batch_size, + conditioning_image=conditioning_image, + prompt=prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + num_images_per_prompt=num_images_per_prompt, + ) + + # 2. Prepare conditioning image latents + if is_t2i: + image_latents = None + else: + processed_image = self.image_processor.preprocess(conditioning_image, height=height, width=width) + processed_image = processed_image.to(device).to(dtype) + image_latents = self.prepare_image_latents( + processed_image, + batch_size, + num_images_per_prompt, + self.dtype, + device, + ) + + # 3. Prepare latents + latent_channels = self.vae.config.latent_channels + latents = self.prepare_latents( # type: ignore + batch_size * num_images_per_prompt, + latent_channels, + height, + width, + dtype, + device, + generator, + latents, + ) + + num_channels_latents_check = ( + latent_channels * 2 if self.transformer.input_condition_type == "channel_cat" else latent_channels + ) + if num_channels_latents_check != self.transformer.config.in_channels: + msg = f"The config of `pipeline.transformer` expects {self.transformer.config.in_channels} channels," + f"but received {num_channels_latents_check}." + raise ValueError(msg) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) # type: ignore[arg-type] + self._num_timesteps = len(timesteps) # type: ignore[arg-type] + + # 5. Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 6. Denoising loop + latents = self._run_denoising_loop( # type: ignore + latents=latents, # type: ignore[arg-type] + latent_channels=latent_channels, + input_image_latents=image_latents, + timesteps=timesteps, # type: ignore[arg-type] + num_inference_steps=num_inference_steps, + num_warmup_steps=num_warmup_steps, + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + guidance_scale=guidance_scale, + image_guidance_scale=image_guidance_scale, + extra_step_kwargs=extra_step_kwargs, + is_t2i=is_t2i, + ) + + if output_type == "latent": + image = latents # type: ignore[assignment] + else: + latents = latents.to(self.dtype) # type: ignore + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + + if use_resolution_binning: + image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) + + if output_type != "latent": + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return SanaPipelineOutput(images=image) # type: ignore[arg-type] diff --git a/vibe/transformer/__init__.py b/vibe/transformer/__init__.py new file mode 100644 index 0000000..926c5bb --- /dev/null +++ b/vibe/transformer/__init__.py @@ -0,0 +1,5 @@ +"""VIBE Sana Editing model.""" + +from .vibe_sana_editing import VIBESanaEditingModel + +__all__ = ["VIBESanaEditingModel"] diff --git a/vibe/transformer/edit_head.py b/vibe/transformer/edit_head.py new file mode 100644 index 0000000..e5d0814 --- /dev/null +++ b/vibe/transformer/edit_head.py @@ -0,0 +1,200 @@ +"""Meta Connector.""" + +import torch +from torch import nn +from torch.nn import functional as func + + +class Expert(nn.Module): + """Position-wise Feed-Forward Networks. + + This consists of two linear transformations with a ReLU activation in between. + """ + + def __init__(self, d_model: int, d_expert: int) -> None: + """Initialize the Expert layer. + + Args: + d_model (int): The dimension of the model. + d_expert (int): The dimension of the expert. + """ + super().__init__() + self.d_model = d_model + self.d_expert = d_expert + + # Linear transformation y = xW+b + self.fc1 = nn.Linear(self.d_model, self.d_expert, bias=True) + self.fc2 = nn.Linear(self.d_expert, self.d_model, bias=True) + self.fc3 = nn.Linear(self.d_model, self.d_expert, bias=True) + + # Initialize weights + self._init_weights() + + def _init_weights(self) -> None: + """Initialize the weights of the Expert layer.""" + # init weights + nn.init.xavier_uniform_(self.fc1.weight) + nn.init.zeros_(self.fc1.bias) + nn.init.xavier_uniform_(self.fc2.weight) + nn.init.zeros_(self.fc2.bias) + nn.init.xavier_uniform_(self.fc3.weight) + nn.init.zeros_(self.fc3.bias) + + def forward(self, embedding: torch.Tensor) -> torch.Tensor: + """Forward pass of the Expert layer. + + Args: + embedding (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor. + """ + inner_emb = func.silu(self.fc1(embedding)) * self.fc3(embedding) + return self.fc2(inner_emb) + + +class MetaAttentionBlock(nn.Module): + """Meta Attention Block.""" + + def __init__(self, d_model: int, n_heads: int, dropout: float = 0.05) -> None: + """Initialize the Meta Attention Block. + + Args: + d_model (int): The dimension of the embeddings. + n_heads (int): The number of attention heads. + dropout (float): The dropout rate. + """ + super().__init__() + self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True) + + # Dropout layers for residual connections + self.dropout_attn = nn.Dropout(dropout) + self.dropout_ffn = nn.Dropout(dropout) + + self.mlp = Expert(d_model, d_model * 4) + + self.input_norm = nn.LayerNorm(d_model) + self.output_norm = nn.LayerNorm(d_model) + + # Initialize weights + self._init_weights() + + def _init_weights(self) -> None: + """Initialize weights for attention and MLP blocks.""" + # Initialize MultiheadAttention weights + # The in_proj_weight contains Q, K, V concatenated + nn.init.xavier_uniform_(self.attn.in_proj_weight) + if self.attn.in_proj_bias is not None: + nn.init.zeros_(self.attn.in_proj_bias) + + # Initialize the output projection of MultiheadAttention + nn.init.xavier_uniform_(self.attn.out_proj.weight) + if self.attn.out_proj.bias is not None: + nn.init.zeros_(self.attn.out_proj.bias) + + nn.init.ones_(self.input_norm.weight) + nn.init.zeros_(self.input_norm.bias) + nn.init.ones_(self.output_norm.weight) + nn.init.zeros_(self.output_norm.bias) + + def forward(self, embedding: torch.Tensor) -> torch.Tensor: + """Forward pass of the MetaAttentionBlock. + + Args: + embedding (torch.Tensor): The input embeddings. Shape: (batch_size, seq_len, hidden_dim). + + Returns: + torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim). + """ + # apply layer normalization to embeddings + normed_embedding = self.input_norm(embedding) + + # apply attention and skip connection + attn_output, _ = self.attn(normed_embedding, normed_embedding, normed_embedding, need_weights=False) + embedding = embedding + self.dropout_attn(attn_output) + + output = self.mlp(self.output_norm(embedding)) + return embedding + self.dropout_ffn(output) + + +class MetaConnector(nn.Module): + """Meta Connector.""" + + def __init__( # noqa: PLR0913 + self, + input_dim: int = 4096, + hidden_dim: int = 2048, + output_dim: int = 768, + n_stacks: int = 8, + heads: int = 16, + *, + use_input_proj: bool = True, + ) -> None: + """Initialize the MetaConnector. + + Args: + input_dim (int): The dimension of the input. + hidden_dim (int): The dimension of the hidden state. + output_dim (int): The dimension of the output. + n_stacks (int): The number of attention layers. + heads (int): The number of attention heads. + use_input_proj (bool): Whether to use input projection. + """ + super().__init__() + self.output_dim = output_dim + self.input_dim = input_dim + self.use_input_proj = use_input_proj + + # Input Projection + if use_input_proj: + self.input_proj = nn.Linear(input_dim, hidden_dim) + else: + hidden_dim = input_dim + + # Attention blocks + self.attn_blocks = nn.Sequential(*[MetaAttentionBlock(hidden_dim, heads) for _ in range(n_stacks)]) + + # Output Projection + self.output_proj = nn.Sequential(nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, output_dim)) + + # Initialize weights + self._init_weights() + + def _init_weights(self) -> None: + """Initialize weights.""" + # Initialize input_proj + if self.use_input_proj: + nn.init.xavier_uniform_(self.input_proj.weight) + if self.input_proj.bias is not None: + nn.init.zeros_(self.input_proj.bias) + + # Initialize output_proj + for module in self.output_proj: + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + def forward(self, hidden_states: torch.Tensor, meta_queries: torch.Tensor) -> dict[str, torch.Tensor]: + """Forward pass of the Meta Connector. + + Args: + hidden_states (torch.Tensor): The input hidden states. Shape: (batch_size, seq_len, hidden_dim). + meta_queries (torch.Tensor): The input meta queries. Shape: (batch_size, seq_len, hidden_dim). + + Returns: + torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim). + """ + hidden_states = hidden_states + meta_queries + if self.use_input_proj: + hidden_states = self.input_proj(hidden_states) + + # apply attention blocks + for attn_block in self.attn_blocks: + hidden_states = attn_block(hidden_states) + + # apply output projection + return self.output_proj(hidden_states) diff --git a/vibe/transformer/vibe_sana_editing.py b/vibe/transformer/vibe_sana_editing.py new file mode 100644 index 0000000..945d49e --- /dev/null +++ b/vibe/transformer/vibe_sana_editing.py @@ -0,0 +1,624 @@ +"""Transformer Sana model.""" + +from typing import Any + +import torch +from diffusers import ModelMixin, SanaTransformer2DModel # noqa: F401 +from diffusers.configuration_utils import register_to_config +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import PatchEmbed, PixArtAlphaTextProjection +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import AdaLayerNormSingle, RMSNorm +from diffusers.models.transformers.sana_transformer import ( + GLUMBConv, + SanaAttnProcessor2_0, + SanaCombinedTimestepGuidanceEmbeddings, + SanaModulatedNorm, +) +from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from torch import nn + +from vibe.transformer.edit_head import MetaConnector + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class SanaLinearAttnProcessor2_0: # noqa: N801 + r"""Processor for implementing scaled dot-product linear attention.""" + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Execute the forward pass of the SanaLinearAttnProcessor2_0. + + We need that redifinition because default implementation do not take into account the attention mask. + + Args: + attn (Attention): The attention object. + hidden_states (torch.Tensor): The hidden states. + encoder_hidden_states (torch.Tensor | None): The encoder hidden states. In case of cross attention. + attention_mask (torch.Tensor | None): The attention mask. + + Returns: + torch.Tensor: The processed hidden states. + """ + original_dtype = hidden_states.dtype + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + query = query.transpose(1, 2).unflatten(1, (attn.heads, -1)) + key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3) + value = value.transpose(1, 2).unflatten(1, (attn.heads, -1)) + + query = nn.functional.relu(query) + key = nn.functional.relu(key) + + query, key, value = query.float(), key.float(), value.float() + + if attention_mask is not None: + # attention_mask([B, S], dtype=torch.bool): True - keep, False - ignore. + # transform to (1.0 - keep, 0.0 - ignore). + mask = attention_mask if attention_mask.dtype == torch.bool else attention_mask > -1 + mask = mask.to(key.dtype) + if mask.ndim == 2: # noqa: PLR2004 + mask = mask.unsqueeze(1) # -> [B, 1, S] + + # Null masked tokens in key and value before torch.matmul operation + key = key * mask.unsqueeze(-1) # (B, _, S, _) * (B, 1, S, 1) + value = value * mask.unsqueeze(2) # (B, _, _, S) * (B, 1, 1, S) + + value = nn.functional.pad(value, (0, 0, 0, 1), mode="constant", value=1.0) + scores = torch.matmul(value, key) + hidden_states = torch.matmul(scores, query) + + hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + 1e-15) + hidden_states = hidden_states.flatten(1, 2).transpose(1, 2) + hidden_states = hidden_states.to(original_dtype) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + if original_dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return hidden_states + + +class SanaTransformerBlock(nn.Module): + """Transformer block introduced in [Sana](huggingface.co/papers/2410.10629).""" + + def __init__( # noqa: PLR0913 + self, + dim: int = 2240, + num_attention_heads: int = 70, + attention_head_dim: int = 32, + dropout: float = 0.0, + num_cross_attention_heads: int = 20, + cross_attention_head_dim: int = 112, + cross_attention_dim: int = 2240, + norm_eps: float = 1e-6, + mlp_ratio: float = 2.5, + qk_norm: str | None = None, + *, + attention_out_bias: bool = True, + attention_bias: bool = True, + norm_elementwise_affine: bool = False, + ) -> None: + """Initialize the SanaTransformerBlock. + + Args: + dim (int): The dimension of the input/output hidden states. + num_attention_heads (int): The number of heads for self-attention. + attention_head_dim (int): The dimension of each self-attention head. + dropout (float): The dropout probability for attention layers. + num_cross_attention_heads (int): The number of heads for cross-attention. + cross_attention_head_dim (int): The dimension of each cross-attention head. + cross_attention_dim (int): The dimension of the cross-attention context. + norm_eps (float): Epsilon value for layer normalization stability. + mlp_ratio (float): Expansion ratio for the feed-forward network hidden dimension. + qk_norm (str | None): Normalization method for Query/Key vectors. + attention_out_bias (bool): Whether to include bias in the attention output projection. + attention_bias (bool): Whether to include bias in the attention Q/K/V projections. + norm_elementwise_affine (bool): Whether to learn affine parameters for normalization. + """ + super().__init__() + + # 1. Self Attention + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + kv_heads=num_attention_heads if qk_norm is not None else None, + qk_norm=qk_norm, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=None, + processor=SanaLinearAttnProcessor2_0(), # type: ignore[arg-type] + ) + + # 2. Cross Attention + if cross_attention_dim is not None: + self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn2 = Attention( + query_dim=dim, + qk_norm=qk_norm, + kv_heads=num_cross_attention_heads if qk_norm is not None else None, + cross_attention_dim=cross_attention_dim, + heads=num_cross_attention_heads, + dim_head=cross_attention_head_dim, + dropout=dropout, + bias=True, + out_bias=attention_out_bias, + processor=SanaAttnProcessor2_0(), # type: ignore[arg-type] + ) + + # 3. Feed-forward + self.ff = GLUMBConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + def forward( # noqa: PLR0913 + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + encoder_hidden_states: torch.Tensor | None, + encoder_attention_mask: torch.Tensor | None, + timestep: torch.LongTensor, + height: int, + width: int, + ) -> torch.Tensor: + """Executes the forward pass of the transformer block. + + Args: + hidden_states (torch.Tensor): Input hidden states of shape `(batch_size, seq_len, dim)`. + attention_mask (torch.Tensor | None): Mask for self-attention. + encoder_hidden_states (torch.Tensor | None): Conditional embeddings (e.g., text) for cross-attention. + encoder_attention_mask (torch.Tensor | None): Mask for cross-attention. + timestep (torch.LongTensor): Timestep embeddings used for adaptive modulation of normalization layers. + height (int): Spatial height of the latent image (used for reshaping in FFN). + width (int): Spatial width of the latent image (used for reshaping in FFN). + + Returns: + torch.Tensor: Transformed hidden states of shape `(batch_size, seq_len, dim)`. + """ + batch_size = hidden_states.shape[0] + + # 1. Modulation + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + + # 2. Self Attention + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) + + attn_output = self.attn1(norm_hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states + gate_msa * attn_output + + # 3. Cross Attention + if self.attn2 is not None: + attn_output = self.attn2( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if width * height == norm_hidden_states.shape[1]: + norm_hidden_states = norm_hidden_states.unflatten(1, (height, width)).permute(0, 3, 1, 2) + ff_output = self.ff(norm_hidden_states) + ff_output = ff_output.flatten(2, 3).permute(0, 2, 1) + else: + # "seq_cat" case + # apply the same local processing to each of the parts independently of each other + # global interaction has already occurred, work with spatial features separately + norm_hidden_states_1, norm_hidden_states_2 = torch.chunk(norm_hidden_states, 2, dim=1) + + hidden_states_reshaped_1 = norm_hidden_states_1.unflatten(1, (height, width)).permute(0, 3, 1, 2) + hidden_states_reshaped_2 = norm_hidden_states_2.unflatten(1, (height, width)).permute(0, 3, 1, 2) + + ff_output_1 = self.ff(hidden_states_reshaped_1).flatten(2, 3).permute(0, 2, 1) + ff_output_2 = self.ff(hidden_states_reshaped_2).flatten(2, 3).permute(0, 2, 1) + + ff_output = torch.cat([ff_output_1, ff_output_2], dim=1) + + return hidden_states + gate_mlp * ff_output + + +class VIBESanaEditingModel(SanaTransformer2DModel): + """VIBE Sana editing model specialized for instruction-based image editing. + + This model extends `SanaTransformer2DModel` to support advanced editing capabilities. + """ + + @register_to_config + def __init__( # noqa: PLR0913 + self, + in_channels: int = 32, + out_channels: int | None = 32, + num_attention_heads: int = 70, + attention_head_dim: int = 32, + num_layers: int = 20, + num_cross_attention_heads: int = 20, + cross_attention_head_dim: int = 112, + cross_attention_dim: int = 2240, + caption_channels: int = 2304, + mlp_ratio: float = 2.5, + dropout: float = 0.0, + sample_size: int = 32, + patch_size: int = 1, + norm_eps: float = 1e-6, + interpolation_scale: int | None = None, + qk_norm: str | None = None, + timestep_scale: float = 1.0, # noqa: ARG002 + input_condition_type: str = "channel_cat", + edit_head_input_dim: int = 2048, + edit_head_stacks_num: int = 4, + num_meta_queries: int = 224, + meta_queries_dim: int = 2048, + *, + attention_bias: bool = False, + norm_elementwise_affine: bool = False, + guidance_embeds: bool = False, + ) -> None: + """Initialize the VIBESanaEditingModel. + + Args: + in_channels (int): Number of channels in the input latent representation. + out_channels (int | None): Number of output channels. Defaults to `in_channels`. + num_attention_heads (int): Number of attention heads in transformer blocks. + attention_head_dim (int): Dimension of each attention head. + num_layers (int): Number of transformer blocks in the model. + num_cross_attention_heads (int): Number of heads for cross-attention layers. + cross_attention_head_dim (int): Dimension of each cross-attention head. + cross_attention_dim (int): Dimension of the cross-attention context embedding. + caption_channels (int): Dimension of the input caption/text embeddings. + mlp_ratio (float): Expansion ratio for the feed-forward networks. + dropout (float): Dropout probability. + sample_size (int): Spatial size of the input sample (height/width). + patch_size (int): Size of the patches extracted from the input latent. + norm_eps (float): Epsilon for layer normalization. + interpolation_scale (int | None): Scale factor for positional embedding interpolation. + qk_norm (str | None): Normalization type for Query/Key (e.g., 'rms_norm'). + timestep_scale (float): Scale factor for the timestep. + input_condition_type (str): Method for conditioning on the input image. Options: 'channel_cat', 'seq_cat'. + edit_head_input_dim (int): Input dimension for the `MetaConnector` edit head. + edit_head_stacks_num (int): Number of layers in the `MetaConnector`. + num_meta_queries (int): Number of meta tokens to extract from the text embeddings for the edit head. + meta_queries_dim (int): Dimension of the meta query embeddings. + attention_bias (bool): Whether to use bias in attention projections. + norm_elementwise_affine (bool): Whether to learn affine parameters in normalization. + guidance_embeds (bool): Whether to use additional guidance embeddings. + """ + super(SanaTransformer2DModel, self).__init__() + + out_channels = out_channels or in_channels + inner_dim = num_attention_heads * attention_head_dim + + # 1. Patch Embedding + self.patch_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + pos_embed_type="sincos" if interpolation_scale is not None else None, + ) + + # 2. Additional condition embeddings + if guidance_embeds: + self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) + else: + self.time_embed = AdaLayerNormSingle(inner_dim) # type: ignore[assignment] + + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + SanaTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + attention_bias=attention_bias, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + ) + for _ in range(num_layers) + ] + ) + + # 4. Output blocks + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) + + # 5. Meta queries embedding + self.num_meta_queries = num_meta_queries + self.meta_queries = nn.Parameter(torch.randn(num_meta_queries, meta_queries_dim) / meta_queries_dim**0.5) + + # 6. Edit head + self.input_condition_type = input_condition_type + self.edit_head = MetaConnector( + input_dim=edit_head_input_dim, + n_stacks=edit_head_stacks_num, + output_dim=caption_channels, + ) + self.gradient_checkpointing = False + + def get_hidden_states_for_meta_tokens( + self, + text_encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Extracts hidden states corresponding to meta tokens from the text encoder output. + + This method identifies the positions of meta tokens in the input sequence using + the attention mask and extracts their embeddings. It assumes the meta tokens + are located at the end of the valid sequence. + + Args: + text_encoder_hidden_states (torch.Tensor): Output embeddings from the text encoder + of shape `(batch_size, seq_len, hidden_dim)`. + attention_mask (torch.Tensor): Boolean mask indicating valid tokens (True) and + padding (False), used to locate the end of the prompt. + + Returns: + torch.Tensor: Extracted meta token embeddings of shape `(batch_size, num_meta_queries, hidden_dim)`. + """ + img_hidden_states: list[torch.Tensor] = [] + for i in range(text_encoder_hidden_states.shape[0]): + hidden_states_masked = text_encoder_hidden_states[i][attention_mask[i]] + img_hidden_states.append(hidden_states_masked[-self.num_meta_queries :, :].unsqueeze(0)) + return torch.cat(img_hidden_states, dim=0) + + def forward_edit_heads( + self, + batch_size: int, + meta_tokens_hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Processes meta token embeddings through the edit head for both conditional and unconditional generation. + + This method generates editing embeddings for: + 1. **Conditional input** (`hid_edit`): Using the extracted `meta_tokens_hidden_states`. + 2. **Unconditional/Null input** (`hid_null`): Using zeroed hidden states, facilitating + Classifier-Free Guidance (CFG) during inference. + + Args: + batch_size (int): The batch size of the generation. + meta_tokens_hidden_states (torch.Tensor): Input hidden states for the meta tokens + extracted from the text prompt. + + Returns: + tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - `hid_edit` (torch.Tensor): Editing embeddings for the conditional branch. + - `hid_null` (torch.Tensor): Editing embeddings for the unconditional (null) branch. + """ + # repeat meta tokens embeddings to fit batch_size + meta_tokens_embeddings = self.meta_queries.unsqueeze(dim=0) + meta_tokens_embeddings = meta_tokens_embeddings.repeat(batch_size, 1, 1) + + # create zero hidden states for null prompts + zero_hidden_states = torch.zeros_like(meta_tokens_hidden_states) + + # get hidden states for real prompts + hid_edit = self.edit_head(meta_tokens_hidden_states, meta_tokens_embeddings) + + # get hidden states for null prompts + hid_null = self.edit_head(zero_hidden_states, meta_tokens_embeddings) + + return hid_edit, hid_null + + def forward( # noqa: PLR0913,PLR0915,PLR0912 + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples: tuple[torch.Tensor] | None = None, + *, + t2i_samples: list[bool] | None = None, + return_dict: bool = True, + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: + """Execute the forward pass of the VIBE Sana Editing Model. + + Args: + hidden_states (torch.Tensor): Noisy latent image input of shape + `(batch_size, num_channels, height, width)`. + encoder_hidden_states (torch.Tensor): Conditional text embeddings of shape + `(batch_size, seq_len, hidden_dim)`. + timestep (torch.Tensor): Denoising timesteps. + guidance (torch.Tensor | None, optional): Guidance embeddings for the diffusion process. + encoder_attention_mask (torch.Tensor | None, optional): Mask for `encoder_hidden_states`. + attention_mask (torch.Tensor | None, optional): Mask for `hidden_states`. + attention_kwargs (dict[str, Any] | None, optional): Extra arguments for attention processors. + controlnet_block_samples (tuple[torch.Tensor] | None, optional): Additional residuals + from a ControlNet model, if used. + t2i_samples (list[bool] | None, optional): Marker if the sample is a t2i sample. + return_dict (bool): If True, returns a `Transformer2DModelOutput`. Otherwise, returns a tuple. + + Returns: + tuple[torch.Tensor, ...] | Transformer2DModelOutput: The denoised latent output. + If `return_dict` is True, returns `Transformer2DModelOutput(sample=output)`. + + Raises: + ValueError: If `input_condition_type` is not 'channel_cat' or 'seq_cat'. + """ + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + elif attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: + logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") + + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: # noqa: PLR2004 + # ensure attention_mask is a bias + # assume that mask is expressed as: + # # (1 = keep, 0 = discard) + # # convert mask into a bias that can be added to attention scores: + # # (keep = +0, discard = -10000.0) + encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) # type: ignore[union-attr] + + # 1. Input + batch_size, num_channels, height, width = hidden_states.shape + + if self.input_condition_type == "channel_cat": # type: ignore + # For mixed batch of t2i and editing samples, we need to mask the input image embeddings for the t2i samples + if t2i_samples is not None and any(t2i_samples): + hidden_states, input_img_embeds = hidden_states.split(num_channels // 2, dim=1) # type: ignore + mask = torch.tensor(t2i_samples, dtype=torch.bool, device=input_img_embeds.device).view(-1, 1, 1, 1) + input_img_embeds = torch.where(mask, torch.zeros_like(input_img_embeds), input_img_embeds) + hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1) # type: ignore + hidden_states = self.patch_embed(hidden_states) + attention_mask = None # There is no need for attention mask in channel_cat case + + elif self.input_condition_type == "seq_cat": + num_channels = hidden_states.shape[1] + hidden_states, input_img_embeds = hidden_states.split(num_channels // 2, dim=1) + hidden_states = self.patch_embed(hidden_states) + + if t2i_samples is not None and all(t2i_samples): + # All samples are t2i samples, so there is no need for attention mask + # Also, we don't need to patch the input image embeddings + attention_mask = None + elif t2i_samples is None or (t2i_samples is not None and not any(t2i_samples)): + # All samples are editing samples, still no need for attention mask + # But we need to patch the input image embeddings + input_img_embeds = self.patch_embed(input_img_embeds) + hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1) + attention_mask = None + else: + # Mixed samples. + input_img_embeds = self.patch_embed(input_img_embeds) + # Create attnention masks for hidden_states and input_img_embeds + hidden_stated_mask = torch.ones( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=torch.bool, + device=hidden_states.device, + ) # Shape: [bs, seq_len] + input_img_embeds_mask = torch.ones( + input_img_embeds.shape[0], + input_img_embeds.shape[1], + dtype=torch.bool, + device=input_img_embeds.device, + ) # Shape: [bs, seq_len] + + # Set False for t2i samples in input_img_embeds_mask + if t2i_samples is not None: + t2i_samples_tensor = torch.as_tensor(t2i_samples, dtype=torch.bool, device=hidden_states.device) + input_img_embeds_mask[t2i_samples_tensor] = False + + hidden_states = torch.cat([hidden_states, input_img_embeds], dim=1) # type: ignore + attention_mask = torch.cat([hidden_stated_mask, input_img_embeds_mask], dim=1) + else: + msg = f"Invalid input condition type: {self.input_condition_type}" + raise ValueError(msg) # type: ignore + + patch_size = self.config.patch_size # type: ignore[attr-defined] + post_patch_height, post_patch_width = height // patch_size, width // patch_size + + if guidance is not None: + timestep, embedded_timestep = self.time_embed(timestep, guidance=guidance, hidden_dtype=hidden_states.dtype) + else: + timestep, embedded_timestep = self.time_embed( + timestep, + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + encoder_hidden_states = self.caption_norm(encoder_hidden_states) + + # 2. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = self._gradient_checkpointing_func( # type: ignore[attr-defined] + block, + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_height, + post_patch_width, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + else: + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_height, + post_patch_width, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + # 3. Normalization + hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, + post_patch_height, + post_patch_width, + patch_size, + patch_size, + -1, + ) + hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4) + output = hidden_states.reshape(batch_size, -1, post_patch_height * patch_size, post_patch_width * patch_size) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/vibe/utils/__init__.py b/vibe/utils/__init__.py new file mode 100644 index 0000000..6ab0a58 --- /dev/null +++ b/vibe/utils/__init__.py @@ -0,0 +1,5 @@ +"""Utility functions.""" + +from .retry_utils import retry_decorator + +__all__ = ["retry_decorator"] diff --git a/vibe/utils/img_utils.py b/vibe/utils/img_utils.py new file mode 100644 index 0000000..f2c2a54 --- /dev/null +++ b/vibe/utils/img_utils.py @@ -0,0 +1,151 @@ +"""Image utils.""" + +import torch +from PIL import Image +from torchvision import transforms as T # type: ignore[import-untyped] # noqa: N812 +from torchvision.transforms.functional import InterpolationMode # type: ignore[import-untyped] + + +def get_divisible_size(height: int, width: int, factor: int = 16, *, expand_to_fit: bool = False) -> tuple[int, int]: + """Adjusts the given dimensions (height and width) to be divisible by the specified factor. + + Args: + height (int): The original height of the input. + width (int): The original width of the input. + factor (int): The factor by which the dimensions should be divisible. Default is 16. + expand_to_fit (bool): If True, dimensions will be adjusted upwards to the nearest multiple of factor. + """ + h_target = height // factor * factor + w_target = width // factor * factor + h_target = max(h_target, factor) + w_target = max(w_target, factor) + + if expand_to_fit: + if w_target < width: + w_target += factor + if h_target < height: + h_target += factor + + return h_target, w_target + + +def get_multiscale_transform( # noqa: PLR0913 + closest_size: tuple[int, int], + orig_size: tuple[int, int], + interpolation: InterpolationMode = InterpolationMode.BICUBIC, + resize_factor: int = 8, + out_img_type: str = "pt", + *, + crop: bool = False, +) -> tuple[T.Compose, tuple[int, ...] | None]: + """Return transformations for the image and x,y,x,y coordinates to remove padding after generations. + + Args: + closest_size (tuple[int, int]): The closest size to the original size. + orig_size (tuple[int, int]): The original size of the image. + interpolation (InterpolationMode): The interpolation mode to use. + resize_factor (int): The resize factor to use. + out_img_type (str): The type of the output image. + crop (bool): Whether to crop the image. + + Returns: + tuple[T.Compose, tuple[int, ...] | None]: The transformations and the crop coordinates. + """ + transforms = [T.Lambda(lambda img: img.convert("RGB"))] + + height, width = closest_size + if height % resize_factor != 0 or width % resize_factor != 0: + closest_size = get_divisible_size(height, width, resize_factor, expand_to_fit=not crop) # type: ignore + + ori_h, ori_w = orig_size + # if crop==True => resize by biggest size; if crop==False => resize by smallest size + if (crop and closest_size[0] / ori_h > closest_size[1] / ori_w) or ( + not crop and closest_size[0] / ori_h < closest_size[1] / ori_w + ): + resize_size = closest_size[0], int(ori_w * closest_size[0] / ori_h) + else: + resize_size = int(ori_h * closest_size[1] / ori_w), closest_size[1] # type: ignore + + if not crop: + # resize to resize_size and pad to closest_size + left, top = abs(closest_size[1] - resize_size[1]) // 2, abs(closest_size[0] - resize_size[0]) // 2 + right, bottom = closest_size[1] - resize_size[1] - left, closest_size[0] - resize_size[0] - top + pad = [left, top, right, bottom] + transforms.extend( + [ + T.Resize(resize_size, interpolation=interpolation), + T.Pad(pad, fill=0, padding_mode="constant"), + ] + ) + crop_coords = (left, top, closest_size[1] - right, closest_size[0] - bottom) + else: + # resize to resize_size and crop to closest_size + crop_coords = None + transforms.extend( + [ + T.Resize(resize_size, interpolation=interpolation), + T.CenterCrop(closest_size), + ] + ) + + if out_img_type == "pt": + transforms.extend([T.ToTensor(), T.Normalize([0.5], [0.5])]) + transform = T.Compose(transforms) + return transform, crop_coords # type: ignore + + +def postprocess_padded_image( + image: Image.Image | torch.Tensor, + crop_coords: tuple[int, ...] | list[tuple[int, ...]] | None, +) -> Image.Image | torch.Tensor: + """Remove paddings from padded images.""" + if crop_coords is None: + return image + + def _crop(img: torch.Tensor, coords: tuple[int, ...]) -> torch.Tensor: + assert len(coords) == 4 # noqa: S101,PLR2004 + if img.shape[0] in [1, 3, 6]: + # for (C, H, W) tensor + left, top = coords[:2] + width, height = coords[2] - coords[0], coords[3] - coords[1] + if img.shape[0] in [1, 3]: + return T.functional.crop(img, top, left, height, width) + winner, looser = img.chunk(2, dim=0) + winner = T.functional.crop(winner, top, left, height, width) + looser = T.functional.crop(looser, top, left, height, width) + return torch.cat((winner, looser), dim=0) + # for (H, W, C) tensor + left, top, right, bottom = coords[:4] + return img[top:bottom, left:right, :] + + if isinstance(image, torch.Tensor) and image.ndim == 4: # noqa: PLR2004 + assert isinstance(crop_coords, list) and image.shape[0] == len(crop_coords) # noqa: PT018,S101 + output_images = [] + for img, coords in zip(image, crop_coords, strict=True): + output_images.append(_crop(img, coords)[None, :]) + return torch.cat(output_images, dim=0) + + if isinstance(image, torch.Tensor) and image.ndim == 3: # noqa: PLR2004 + return _crop(image, crop_coords) # type: ignore + + x1, y1, x2, y2 = crop_coords + return image.crop([x1, y1, x2, y2]) # type: ignore + + +def revert_resize(img: Image.Image, initial_hw: tuple[int, int]) -> Image.Image: + """Resize the image to the initial height and width. + + Args: + img (Image.Image): The input image. + initial_hw (tuple[int, int]): The initial height and width. + + Returns: + Image.Image: The resized image. + """ + height, width = initial_hw + if (width, height) == img.size: + return img + + if hasattr(Image, "Resampling"): + return img.resize((width, height), Image.Resampling.LANCZOS) + return img.resize((width, height), Image.LANCZOS) # type: ignore[attr-defined] diff --git a/vibe/utils/retry_utils.py b/vibe/utils/retry_utils.py new file mode 100644 index 0000000..fe3406e --- /dev/null +++ b/vibe/utils/retry_utils.py @@ -0,0 +1,49 @@ +"""Retry utils.""" + +import functools +import time +from collections.abc import Callable +from typing import Any + +import loguru # noqa: TC002 + + +def retry_decorator( + logger: "loguru.Logger", + max_retries: int = 5, + delay: int = 0, + exception: type[Exception] = Exception, +) -> Callable: + """Retry a function up to retries times with a delay between retries if a specified exception is raised. + + Args: + logger (loguru.Logger): The logger to use for logging. + max_retries (int): Number of times to retry the function. + delay (int): Time to wait between retries (in seconds). + exception (Exception): The type of exception to catch and retry on. + + Returns: + function: The wrapped function with retry logic. + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + attempts = 0 + while attempts < max_retries: + try: + return func(*args, **kwargs) + except exception as exp: + attempts += 1 + logger.debug(f"Attempt {attempts}/{max_retries} failed: {exp}") + if attempts == max_retries: + logger.exception("Max retries reached.") + raise + time.sleep(delay) + + message = "Retry loop exited without returning or raising." + raise RuntimeError(message) + + return wrapper + + return decorator