Initial clean commit for ComfyUI-VIBE
This commit is contained in:
Artur Penner
@@ -0,0 +1,3 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.DS_Store
|
||||
@@ -0,0 +1,2 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -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" }
|
||||
+171
@@ -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
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Generative pipeline."""
|
||||
|
||||
from .vibe_sana_pipeline import VIBESanaEditingPipeline
|
||||
|
||||
__all__ = ["VIBESanaEditingPipeline"]
|
||||
@@ -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]
|
||||
@@ -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]
|
||||
@@ -0,0 +1,5 @@
|
||||
"""VIBE Sana Editing model."""
|
||||
|
||||
from .vibe_sana_editing import VIBESanaEditingModel
|
||||
|
||||
__all__ = ["VIBESanaEditingModel"]
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Utility functions."""
|
||||
|
||||
from .retry_utils import retry_decorator
|
||||
|
||||
__all__ = ["retry_decorator"]
|
||||
@@ -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]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user