diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..28bebb2f2ed306eb783abf19e30ba53a25629925 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +examples/1.png filter=lfs diff=lfs merge=lfs -text +examples/12.png filter=lfs diff=lfs merge=lfs -text +examples/5.png filter=lfs diff=lfs merge=lfs -text +examples/9.png filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 7ae4d757a98ac2ace9ec7d868470d4cd3b5ebdb8..9b121842cf2706e2bb9cdb79f14a2ff080364188 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,32 @@ --- -title: Metaview -emoji: 🐢 -colorFrom: gray -colorTo: green +title: MetaView Novel View Synthesis +emoji: 🎥 +colorFrom: pink +colorTo: purple sdk: gradio -sdk_version: 6.20.0 -python_version: '3.12' +sdk_version: 5.49.1 app_file: app.py pinned: false +license: apache-2.0 +short_description: Monocular novel view synthesis from a single image +python_version: "3.10" +startup_duration_timeout: 1h --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# MetaView — Monocular Novel View Synthesis + +Interactive demo for [**MetaView**](https://huggingface.co/Kwai-Kolors/MetaView) +(ECCV 2026): *Monocular Novel View Synthesis with Scale-Aware Implicit Geometry +Priors*. + +Upload a single image, choose a target camera **yaw** (left/right) and **pitch** +(up/down), and MetaView renders the scene from the new viewpoint. + +The pipeline combines: +- **Qwen-Image-Edit** MM-DiT backbone (novel-view generation), +- **Depth-Anything-3** GIANT (implicit 3D feature priors) + NESTED (dense depth), +- **Modified RoPE (PRoPE)** for metric scale anchoring of the camera pose. + +Links: [Model](https://huggingface.co/Kwai-Kolors/MetaView) · +[Code](https://github.com/KlingAIResearch/MetaView) · +[Paper](https://arxiv.org/abs/2607.12000) diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..90c3252b6644a1ca964824cd5e6a690c854353db --- /dev/null +++ b/app.py @@ -0,0 +1,270 @@ +import os + +# Allocator config to survive transient memory spikes from the large DiT. +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +# DiffSynth defaults to ModelScope; force Hugging Face so all weights come from the Hub. +os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface") + +import glob +import math +import sys +import tempfile + +import spaces # must come before torch / any CUDA-touching import +import numpy as np +import torch +import torch.nn.functional as F +import gradio as gr +from PIL import Image +from huggingface_hub import snapshot_download + +# Local vendored packages (MetaView `src/` + `diffsynth/` and Depth-Anything-3). +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from depth_anything_3.api import DepthAnything3 +from diffsynth.core import ModelConfig +from diffsynth import load_state_dict +from src.MetaView_pipeline import MetaViewPipeline + +# --------------------------------------------------------------------------- +# Model ids +# --------------------------------------------------------------------------- +METAVIEW_REPO = "Kwai-Kolors/MetaView" +DA3_GIANT_REPO = "depth-anything/DA3-GIANT-1.1" +DA3_NESTED_REPO = "depth-anything/DA3NESTED-GIANT-LARGE-1.1" +QWEN_EDIT_REPO = "Qwen/Qwen-Image-Edit" + +# MetaView global config (matches the reference inference script). +EXPORT_3D_FEAT_LAYERS = [19, 27, 33, 39] +PROPE_DIM_ARRANGE = [64, 20, 20, 24] +ADD_DEPTH = len(PROPE_DIM_ARRANGE) == 4 +MERGE_3D = True +PROMPT = ["镜头视角转到指定位置"] # "move the camera view to the target position" +GEN_W, GEN_H = 960, 528 +NUM_STEPS = 40 + +device = "cuda" +dtype = torch.bfloat16 + +# --------------------------------------------------------------------------- +# Load everything at module scope. `import spaces` above has hijacked +# torch.cuda.*, so `.to("cuda")` here packs weights to disk and streams them +# into VRAM on the first @spaces.GPU call. +# --------------------------------------------------------------------------- +print("[*] Downloading MetaView checkpoint...") +metaview_dir = snapshot_download(METAVIEW_REPO) +metaview_ckpt = glob.glob(os.path.join(metaview_dir, "*.safetensors"))[0] + +print("[*] Downloading Depth-Anything-3 models...") +da3_giant_dir = snapshot_download(DA3_GIANT_REPO) +da3_nested_dir = snapshot_download(DA3_NESTED_REPO) + +# Download the Qwen-Image-Edit backbone components from the Hub explicitly so we +# can hand DiffSynth concrete local file paths (avoids the ModelScope default +# download path and the processor-dir globbing quirk). +print("[*] Downloading Qwen-Image-Edit backbone...") +qwen_edit_dir = snapshot_download( + QWEN_EDIT_REPO, + allow_patterns=["transformer/diffusion_pytorch_model*.safetensors", + "transformer/*.json", "processor/*"], +) +qwen_image_dir = snapshot_download( + "Qwen/Qwen-Image", + allow_patterns=["text_encoder/model*.safetensors", "text_encoder/*.json", + "vae/diffusion_pytorch_model.safetensors", "vae/*.json", + "tokenizer/*"], +) + +transformer_files = sorted(glob.glob( + os.path.join(qwen_edit_dir, "transformer", "diffusion_pytorch_model*.safetensors"))) +text_encoder_files = sorted(glob.glob( + os.path.join(qwen_image_dir, "text_encoder", "model*.safetensors"))) +vae_file = os.path.join(qwen_image_dir, "vae", "diffusion_pytorch_model.safetensors") +processor_path = os.path.join(qwen_edit_dir, "processor") +tokenizer_path = os.path.join(qwen_image_dir, "tokenizer") + +print("[*] Loading Depth-Anything-3 GIANT (3D feature extractor)...") +da3_giant = DepthAnything3.from_pretrained(da3_giant_dir).to(device=device).eval() + +print("[*] Loading Depth-Anything-3 NESTED (dense depth)...") +da3_nested = DepthAnything3.from_pretrained(da3_nested_dir).to(device=device).eval() + +print("[*] Loading MetaView / Qwen-Image-Edit pipeline...") +pipe = MetaViewPipeline.from_pretrained( + torch_dtype=dtype, + device=device, + model_configs=[ + ModelConfig(path=transformer_files), + ModelConfig(path=text_encoder_files), + ModelConfig(path=vae_file), + ], + tokenizer_config=ModelConfig(path=tokenizer_path), + processor_config=ModelConfig(path=processor_path), +) + +print(f"[*] Applying MetaView weights from {metaview_ckpt}...") +state_dict = load_state_dict(metaview_ckpt) +pipe.dit.load_state_dict(state_dict, strict=False) +del state_dict +print("[*] All models loaded.") + + +def compute_target_extrinsic(yaw_deg, pitch_deg, radius): + """Camera World-to-Camera extrinsic for a rotation around a sphere center + in front of the camera (yaw = left/right, pitch = up/down).""" + yaw = math.radians(yaw_deg) + pitch = math.radians(pitch_deg) + R_y = np.array([[np.cos(yaw), 0, np.sin(yaw)], + [0, 1, 0], + [-np.sin(yaw), 0, np.cos(yaw)]]) + R_x = np.array([[1, 0, 0], + [0, np.cos(pitch), -np.sin(pitch)], + [0, np.sin(pitch), np.cos(pitch)]]) + R = R_y @ R_x + C = np.array([0.0, 0.0, radius]) + t = C - R @ C + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + return T + + +@spaces.GPU(duration=150, size="xlarge") +def synthesize(image, yaw, pitch, radius, progress=gr.Progress(track_tqdm=True)): + """Synthesize a novel view of a single input image at a target camera pose. + + Args: + image: Source image (a single monocular view). + yaw: Horizontal camera rotation in degrees (positive = right, negative = left). + pitch: Vertical camera rotation in degrees (positive = up, negative = down). + radius: Rotation radius. If 0, it is auto-derived from the scene center depth. + + Returns: + The synthesized novel-view image at the requested camera pose. + """ + if image is None: + raise gr.Error("Please provide an input image.") + + original_image = image.convert("RGB") + edit_image = original_image.resize((GEN_W, GEN_H)) + + with torch.inference_mode(): + # --- 1. 3D feature extraction (DA3 GIANT) + intrinsics --- + feat_out = da3_giant.inference( + [edit_image], export_feat_layers=EXPORT_3D_FEAT_LAYERS, process_res=840 + ) + intri = feat_out.intrinsics[0] + width = intri[0, 2] * 2 + height = intri[1, 2] * 2 + Ks_matrix = [ + [intri[0, 0] / width, 0.0, 0.0], + [0.0, intri[1, 1] / height, 0.0], + [0.0, 0.0, 1.0], + ] + Ks = torch.Tensor(Ks_matrix) + Ks = torch.stack([Ks, Ks], dim=0).unsqueeze(0) # (1, 2, 3, 3) + + feats = [torch.from_numpy(feat_out.aux[f"feat_layer_{layer}"]) + for layer in EXPORT_3D_FEAT_LAYERS] + feat_3D = torch.cat(feats, dim=-1).to(dtype=dtype, device=device) + + # --- 2. Dense depth estimation (DA3 NESTED) --- + prediction = da3_nested.inference([edit_image], process_res=840) + depth_edit = torch.Tensor(prediction.depth).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(GEN_H, GEN_W), + mode="bilinear", align_corners=False)[0] + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0).unsqueeze(0) # (1, 2, H, W) + + # --- 3. Target pose --- + r = float(radius) + if r <= 0: + depth_squeeze = depth[0, 1] + r = depth_squeeze[depth_squeeze.shape[0] // 2, + depth_squeeze.shape[1] // 2].item() + extrinsic_target = compute_target_extrinsic(float(yaw), float(pitch), r) + extrinsic_source = np.eye(4) + viewmats = torch.Tensor( + np.stack((extrinsic_target, extrinsic_source), axis=0) + ).unsqueeze(0) # (1, 2, 4, 4) -> [target, source] + + # --- 4. Novel view generation (MetaView DiT) --- + generated_image = pipe( + PROMPT, edit_image=edit_image, edit_image_auto_resize=False, + seed=0, + viewmats=viewmats.to(device=device, dtype=dtype), + Ks=Ks.to(device=device, dtype=dtype), + prope_dim_arrange=PROPE_DIM_ARRANGE, + add_attn=True, + add_3D=True, + feat_3D=feat_3D, + depth=depth.to(device=device, dtype=dtype) if ADD_DEPTH else None, + merge_3D=MERGE_3D, + val=True, + num_inference_steps=NUM_STEPS, + height=GEN_H, width=GEN_W, + ) + + return generated_image + + +CSS = """ +#col-container { max-width: 1100px; margin: 0 auto; } +.dark .gradio-container { color: var(--body-text-color); } +""" + +with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: + with gr.Column(elem_id="col-container"): + gr.Markdown( + """ + # MetaView — Monocular Novel View Synthesis + Synthesize a **novel camera view** from a single image. Upload an image, + pick a target camera **yaw** (left/right) and **pitch** (up/down), and + MetaView renders the scene from that new viewpoint. + + Built on Qwen-Image-Edit + Depth-Anything-3 geometry priors. + [Model](https://huggingface.co/Kwai-Kolors/MetaView) · + [Code](https://github.com/KlingAIResearch/MetaView) · + [Paper](https://arxiv.org/abs/2607.12000) + """ + ) + with gr.Row(): + with gr.Column(): + image = gr.Image(label="Input image", type="pil", height=340) + with gr.Row(): + yaw = gr.Slider(-60, 60, value=-30, step=1, + label="Yaw (°) ← left | right →") + pitch = gr.Slider(-45, 45, value=10, step=1, + label="Pitch (°) ↓ down | up ↑") + run = gr.Button("Synthesize novel view", variant="primary") + with gr.Accordion("Advanced settings", open=False): + radius = gr.Slider( + 0.0, 10.0, value=0.0, step=0.1, + label="Rotation radius (0 = auto from center depth)", + ) + with gr.Column(): + output = gr.Image(label="Novel view", height=340) + + gr.Examples( + examples=[ + ["examples/1.png", -30, 10], + ["examples/5.png", 30, 0], + ["examples/9.png", -25, 15], + ["examples/12.png", 40, -10], + ], + inputs=[image, yaw, pitch], + outputs=output, + fn=synthesize, + cache_examples=True, + cache_mode="lazy", + ) + + run.click( + synthesize, + inputs=[image, yaw, pitch, radius], + outputs=output, + api_name="synthesize", + ) + +if __name__ == "__main__": + demo.launch(mcp_server=True) diff --git a/depth_anything_3/api.py b/depth_anything_3/api.py new file mode 100644 index 0000000000000000000000000000000000000000..b672f39c89efa809d10fcb2fe0d44955c33c930a --- /dev/null +++ b/depth_anything_3/api.py @@ -0,0 +1,451 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Depth Anything 3 API module. + +This module provides the main API for Depth Anything 3, including model loading, +inference, and export capabilities. It supports both single and nested model architectures. +""" + +from __future__ import annotations + +import time +from typing import Optional, Sequence +import numpy as np +import torch +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin +from PIL import Image + +from depth_anything_3.cfg import create_object, load_config +from depth_anything_3.registry import MODEL_REGISTRY +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.geometry import affine_inverse +from depth_anything_3.utils.io.input_processor import InputProcessor +from depth_anything_3.utils.io.output_processor import OutputProcessor +from depth_anything_3.utils.logger import logger + +# NOTE: `export` (gsplat / moviepy / open3d / pycolmap) and `align_poses_umeyama` +# (evo) are only used in the export / pose-alignment code paths that this demo +# never exercises. Import them lazily inside the functions that use them so the +# heavy, hard-to-build deps aren't required at module import time. + +torch.backends.cudnn.benchmark = False +# logger.info("CUDNN Benchmark Disabled") + +SAFETENSORS_NAME = "model.safetensors" +CONFIG_NAME = "config.json" + + +class DepthAnything3(nn.Module, PyTorchModelHubMixin): + """ + Depth Anything 3 main API class. + + This class provides a high-level interface for depth estimation using Depth Anything 3. + It supports both single and nested model architectures with metric scaling capabilities. + + Features: + - Hugging Face Hub integration via PyTorchModelHubMixin + - Support for multiple model presets (vitb, vitg, nested variants) + - Automatic mixed precision inference + - Export capabilities for various formats (GLB, PLY, NPZ, etc.) + - Camera pose estimation and metric depth scaling + + Usage: + # Load from Hugging Face Hub + model = DepthAnything3.from_pretrained("huggingface/model-name") + + # Or create with specific preset + model = DepthAnything3(preset="vitg") + + # Run inference + prediction = model.inference(images, export_dir="output", export_format="glb") + """ + + _commit_hash: str | None = None # Set by mixin when loading from Hub + + def __init__(self, model_name: str = "da3-large", **kwargs): + """ + Initialize DepthAnything3 with specified preset. + + Args: + model_name: The name of the model preset to use. + Examples: 'da3-giant', 'da3-large', 'da3metric-large', 'da3nested-giant-large'. + **kwargs: Additional keyword arguments (currently unused). + """ + super().__init__() + self.model_name = model_name + + # Build the underlying network + self.config = load_config(MODEL_REGISTRY[self.model_name]) + self.model = create_object(self.config) + self.model.eval() + + # Initialize processors + self.input_processor = InputProcessor() + self.output_processor = OutputProcessor() + + # Device management (set by user) + self.device = None + + @torch.inference_mode() + def forward( + self, + image: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = None, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, torch.Tensor]: + """ + Forward pass through the model. + + Args: + image: Input batch with shape ``(B, N, 3, H, W)`` on the model device. + extrinsics: Optional camera extrinsics with shape ``(B, N, 4, 4)``. + intrinsics: Optional camera intrinsics with shape ``(B, N, 3, 3)``. + export_feat_layers: Layer indices to return intermediate features for. + infer_gs: Enable Gaussian Splatting branch. + use_ray_pose: Use ray-based pose estimation instead of camera decoder. + ref_view_strategy: Strategy for selecting reference view from multiple views. + + Returns: + Dictionary containing model predictions + """ + # Determine optimal autocast dtype + autocast_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + with torch.no_grad(): + with torch.autocast(device_type=image.device.type, dtype=autocast_dtype): + return self.model( + image, extrinsics, intrinsics, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy + ) + + def inference( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + align_to_input_ext_scale: bool = True, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + render_exts: np.ndarray | None = None, + render_ixts: np.ndarray | None = None, + render_hw: tuple[int, int] | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_dir: str | None = None, + export_format: str = "mini_npz", + export_feat_layers: Sequence[int] | None = None, + # GLB export parameters + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + # Feat_vis export parameters + feat_vis_fps: int = 15, + # Other export parameters, e.g., gs_ply, gs_video + export_kwargs: Optional[dict] = {}, + ) -> Prediction: + """ + Run inference on input images. + + Args: + image: List of input images (numpy arrays, PIL Images, or file paths) + extrinsics: Camera extrinsics (N, 4, 4) + intrinsics: Camera intrinsics (N, 3, 3) + align_to_input_ext_scale: whether to align the input pose scale to the prediction + infer_gs: Enable the 3D Gaussian branch (needed for `gs_ply`/`gs_video` exports) + use_ray_pose: Use ray-based pose estimation instead of camera decoder (default: False) + ref_view_strategy: Strategy for selecting reference view from multiple views. + Options: "first", "middle", "saddle_balanced", "saddle_sim_range". + Default: "saddle_balanced". For single view input (S ≤ 2), no reordering is performed. + render_exts: Optional render extrinsics for Gaussian video export + render_ixts: Optional render intrinsics for Gaussian video export + render_hw: Optional render resolution for Gaussian video export + process_res: Processing resolution + process_res_method: Resize method for processing + export_dir: Directory to export results + export_format: Export format (mini_npz, npz, glb, ply, gs, gs_video) + export_feat_layers: Layer indices to export intermediate features from + conf_thresh_percentile: [GLB] Lower percentile for adaptive confidence threshold (default: 40.0) # noqa: E501 + num_max_points: [GLB] Maximum number of points in the point cloud (default: 1,000,000) + show_cameras: [GLB] Show camera wireframes in the exported scene (default: True) + feat_vis_fps: [FEAT_VIS] Frame rate for output video (default: 15) + export_kwargs: additional arguments to export functions. + + Returns: + Prediction object containing depth maps and camera parameters + """ + if "gs" in export_format: + assert infer_gs, "must set `infer_gs=True` to perform gs-related export." + + if "colmap" in export_format: + assert isinstance(image[0], str), "`image` must be image paths for COLMAP export." + + # Preprocess images + imgs_cpu, extrinsics, intrinsics = self._preprocess_inputs( + image, extrinsics, intrinsics, process_res, process_res_method + ) + + # Prepare tensors for model + imgs, ex_t, in_t = self._prepare_model_inputs(imgs_cpu, extrinsics, intrinsics) + + # Normalize extrinsics + ex_t_norm = self._normalize_extrinsics(ex_t.clone() if ex_t is not None else None) + + # Run model forward pass + export_feat_layers = list(export_feat_layers) if export_feat_layers is not None else [] + + raw_output = self._run_model_forward( + imgs, ex_t_norm, in_t, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy + ) + + # Convert raw output to prediction + prediction = self._convert_to_prediction(raw_output) + + # Align prediction to extrinsincs + prediction = self._align_to_input_extrinsics_intrinsics( + extrinsics, intrinsics, prediction, align_to_input_ext_scale + ) + + # Add processed images for visualization + prediction = self._add_processed_images(prediction, imgs_cpu) + + # Export if requested + if export_dir is not None: + + if "gs" in export_format: + if infer_gs and "gs_video" not in export_format: + export_format = f"{export_format}-gs_video" + if "gs_video" in export_format: + if "gs_video" not in export_kwargs: + export_kwargs["gs_video"] = {} + export_kwargs["gs_video"].update( + { + "extrinsics": render_exts, + "intrinsics": render_ixts, + "out_image_hw": render_hw, + } + ) + # Add GLB export parameters + if "glb" in export_format: + if "glb" not in export_kwargs: + export_kwargs["glb"] = {} + export_kwargs["glb"].update( + { + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + } + ) + # Add Feat_vis export parameters + if "feat_vis" in export_format: + if "feat_vis" not in export_kwargs: + export_kwargs["feat_vis"] = {} + export_kwargs["feat_vis"].update( + { + "fps": feat_vis_fps, + } + ) + # Add COLMAP export parameters + if "colmap" in export_format: + if "colmap" not in export_kwargs: + export_kwargs["colmap"] = {} + export_kwargs["colmap"].update( + { + "image_paths": image, + "conf_thresh_percentile": conf_thresh_percentile, + "process_res_method": process_res_method, + } + ) + self._export_results(prediction, export_format, export_dir, **export_kwargs) + + return prediction + + def _preprocess_inputs( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Preprocess input images using input processor.""" + start_time = time.time() + imgs_cpu, extrinsics, intrinsics = self.input_processor( + image, + extrinsics.copy() if extrinsics is not None else None, + intrinsics.copy() if intrinsics is not None else None, + process_res, + process_res_method, + ) + end_time = time.time() + logger.info( + "Processed Images Done taking", + end_time - start_time, + "seconds. Shape: ", + imgs_cpu.shape, + ) + return imgs_cpu, extrinsics, intrinsics + + def _prepare_model_inputs( + self, + imgs_cpu: torch.Tensor, + extrinsics: torch.Tensor | None, + intrinsics: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Prepare tensors for model input.""" + device = self._get_model_device() + + # Move images to model device + imgs = imgs_cpu.to(device, non_blocking=True)[None].float() + + # Convert camera parameters to tensors + ex_t = ( + extrinsics.to(device, non_blocking=True)[None].float() + if extrinsics is not None + else None + ) + in_t = ( + intrinsics.to(device, non_blocking=True)[None].float() + if intrinsics is not None + else None + ) + + return imgs, ex_t, in_t + + def _normalize_extrinsics(self, ex_t: torch.Tensor | None) -> torch.Tensor | None: + """Normalize extrinsics""" + if ex_t is None: + return None + transform = affine_inverse(ex_t[:, :1]) + ex_t_norm = ex_t @ transform + c2ws = affine_inverse(ex_t_norm) + translations = c2ws[..., :3, 3] + dists = translations.norm(dim=-1) + median_dist = torch.median(dists) + median_dist = torch.clamp(median_dist, min=1e-1) + ex_t_norm[..., :3, 3] = ex_t_norm[..., :3, 3] / median_dist + return ex_t_norm + + def _align_to_input_extrinsics_intrinsics( + self, + extrinsics: torch.Tensor | None, + intrinsics: torch.Tensor | None, + prediction: Prediction, + align_to_input_ext_scale: bool = True, + ransac_view_thresh: int = 10, + ) -> Prediction: + """Align depth map to input extrinsics""" + if extrinsics is None: + return prediction + from depth_anything_3.utils.pose_align import align_poses_umeyama + prediction.intrinsics = intrinsics.numpy() + _, _, scale, aligned_extrinsics = align_poses_umeyama( + prediction.extrinsics, + extrinsics.numpy(), + ransac=len(extrinsics) >= ransac_view_thresh, + return_aligned=True, + random_state=42, + ) + if align_to_input_ext_scale: + prediction.extrinsics = extrinsics[..., :3, :].numpy() + prediction.depth /= scale + else: + prediction.extrinsics = aligned_extrinsics + return prediction + + def _run_model_forward( + self, + imgs: torch.Tensor, + ex_t: torch.Tensor | None, + in_t: torch.Tensor | None, + export_feat_layers: Sequence[int] | None = None, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, torch.Tensor]: + """Run model forward pass.""" + device = imgs.device + need_sync = device.type == "cuda" + if need_sync: + torch.cuda.synchronize(device) + start_time = time.time() + feat_layers = list(export_feat_layers) if export_feat_layers is not None else None + output = self.forward(imgs, ex_t, in_t, feat_layers, infer_gs, use_ray_pose, ref_view_strategy) + if need_sync: + torch.cuda.synchronize(device) + end_time = time.time() + logger.info(f"Model Forward Pass Done. Time: {end_time - start_time} seconds") + return output + + def _convert_to_prediction(self, raw_output: dict[str, torch.Tensor]) -> Prediction: + """Convert raw model output to Prediction object.""" + start_time = time.time() + output = self.output_processor(raw_output) + end_time = time.time() + logger.info(f"Conversion to Prediction Done. Time: {end_time - start_time} seconds") + return output + + def _add_processed_images(self, prediction: Prediction, imgs_cpu: torch.Tensor) -> Prediction: + """Add processed images to prediction for visualization.""" + # Convert from (N, 3, H, W) to (N, H, W, 3) and denormalize + processed_imgs = imgs_cpu.permute(0, 2, 3, 1).cpu().numpy() # (N, H, W, 3) + + # Denormalize from ImageNet normalization + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + processed_imgs = processed_imgs * std + mean + processed_imgs = np.clip(processed_imgs, 0, 1) + processed_imgs = (processed_imgs * 255).astype(np.uint8) + + prediction.processed_images = processed_imgs + return prediction + + def _export_results( + self, prediction: Prediction, export_format: str, export_dir: str, **kwargs + ) -> None: + """Export results to specified format and directory.""" + from depth_anything_3.utils.export import export + start_time = time.time() + export(prediction, export_format, export_dir, **kwargs) + end_time = time.time() + logger.info(f"Export Results Done. Time: {end_time - start_time} seconds") + + def _get_model_device(self) -> torch.device: + """ + Get the device where the model is located. + + Returns: + Device where the model parameters are located + + Raises: + ValueError: If no tensors are found in the model + """ + if self.device is not None: + return self.device + + # Find device from parameters + for param in self.parameters(): + self.device = param.device + return param.device + + # Find device from buffers + for buffer in self.buffers(): + self.device = buffer.device + return buffer.device + + raise ValueError("No tensor found in model") diff --git a/depth_anything_3/app/css_and_html.py b/depth_anything_3/app/css_and_html.py new file mode 100644 index 0000000000000000000000000000000000000000..d414df9db72a4a395bc15d87ff4edbeb213f18bb --- /dev/null +++ b/depth_anything_3/app/css_and_html.py @@ -0,0 +1,594 @@ +# flake8: noqa: E501 + +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +CSS and HTML content for the Depth Anything 3 Gradio application. +This module contains all the CSS styles and HTML content blocks +used in the Gradio interface. +""" + +# CSS Styles for the Gradio interface +GRADIO_CSS = """ +/* Add Font Awesome CDN with all styles including brands and colors */ +@import url('https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css'); + +/* Add custom styles for colored icons */ +.fa-color-blue { + color: #3b82f6; +} + +.fa-color-purple { + color: #8b5cf6; +} + +.fa-color-cyan { + color: #06b6d4; +} + +.fa-color-green { + color: #10b981; +} + +.fa-color-yellow { + color: #f59e0b; +} + +.fa-color-red { + color: #ef4444; +} + +.link-btn { + display: inline-flex; + align-items: center; + gap: 8px; + text-decoration: none; + padding: 12px 24px; + border-radius: 50px; + font-weight: 500; + transition: all 0.3s ease; +} + +/* Dark mode tech theme */ +@media (prefers-color-scheme: dark) { + html, body { + background: #1e293b; + color: #ffffff; + } + + .gradio-container { + background: #1e293b; + color: #ffffff; + } + + .link-btn { + background: rgba(255, 255, 255, 0.2); + color: white; + backdrop-filter: blur(10px); + border: 1px solid rgba(255, 255, 255, 0.3); + } + + .link-btn:hover { + background: rgba(255, 255, 255, 0.3); + transform: translateY(-2px); + box-shadow: 0 8px 25px rgba(0, 0, 0, 0.2); + } + + .tech-bg { + background: linear-gradient(135deg, #0f172a, #1e293b); /* Darker colors */ + position: relative; + overflow: hidden; + } + + .tech-bg::before { + content: ''; + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: + radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */ + radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */ + radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.1) 0%, transparent 50%); /* Reduced opacity */ + animation: techPulse 8s ease-in-out infinite; + } + + .gradio-container .panel, + .gradio-container .block, + .gradio-container .form { + background: rgba(0, 0, 0, 0.3); + border: 1px solid rgba(59, 130, 246, 0.2); + border-radius: 10px; + } + + .gradio-container * { + color: #ffffff; + } + + .gradio-container label { + color: #e0e0e0; + } + + .gradio-container .markdown { + color: #e0e0e0; + } +} + +/* Light mode tech theme */ +@media (prefers-color-scheme: light) { + html, body { + background: #ffffff; + color: #1e293b; + } + + .gradio-container { + background: #ffffff; + color: #1e293b; + } + + .tech-bg { + background: linear-gradient(135deg, #ffffff, #f1f5f9); + position: relative; + overflow: hidden; + } + + .link-btn { + background: rgba(59, 130, 246, 0.15); + color: var(--body-text-color); + border: 1px solid rgba(59, 130, 246, 0.3); + } + + .link-btn:hover { + background: rgba(59, 130, 246, 0.25); + transform: translateY(-2px); + box-shadow: 0 8px 25px rgba(59, 130, 246, 0.2); + } + + .tech-bg::before { + content: ''; + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: + radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.1) 0%, transparent 50%), + radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.1) 0%, transparent 50%), + radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.08) 0%, transparent 50%); + animation: techPulse 8s ease-in-out infinite; + } + + .gradio-container .panel, + .gradio-container .block, + .gradio-container .form { + background: rgba(255, 255, 255, 0.8); + border: 1px solid rgba(59, 130, 246, 0.3); + border-radius: 10px; + box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); + } + + .gradio-container * { + color: #1e293b; + } + + .gradio-container label { + color: #334155; + } + + .gradio-container .markdown { + color: #334155; + } +} + + + + +@keyframes techPulse { + 0%, 100% { opacity: 0.5; } + 50% { opacity: 0.8; } +} + +/* Custom log with tech gradient */ +.custom-log * { + font-style: italic; + font-size: 22px !important; + background: linear-gradient(135deg, #3b82f6, #8b5cf6); + background-size: 400% 400%; + -webkit-background-clip: text; + background-clip: text; + font-weight: bold !important; + color: transparent !important; + text-align: center !important; + animation: techGradient 3s ease infinite; +} + +@keyframes techGradient { + 0% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } + 100% { background-position: 0% 50%; } +} + +@keyframes metricPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes pointcloudPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes camerasPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes gaussiansPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +/* Special colors for key terms - Global styles */ +.metric-text { + background: linear-gradient(45deg, #ff6b6b, #ff8e53, #ff6b6b); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: metricPulse 2s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(255, 107, 107, 0.5); +} + +.pointcloud-text { + background: linear-gradient(45deg, #4ecdc4, #44a08d, #4ecdc4); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: pointcloudPulse 2.5s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(78, 205, 196, 0.5); +} + +.cameras-text { + background: linear-gradient(45deg, #667eea, #764ba2, #667eea); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: camerasPulse 3s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(102, 126, 234, 0.5); +} + +.gaussians-text { + background: linear-gradient(45deg, #f093fb, #f5576c, #f093fb); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: gaussiansPulse 2.2s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(240, 147, 251, 0.5); +} + +.example-log * { + font-style: italic; + font-size: 16px !important; + background: linear-gradient(135deg, #3b82f6, #8b5cf6); + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; +} + +#my_radio .wrap { + display: flex; + flex-wrap: nowrap; + justify-content: center; + align-items: center; +} + +#my_radio .wrap label { + display: flex; + width: 50%; + justify-content: center; + align-items: center; + margin: 0; + padding: 10px 0; + box-sizing: border-box; +} + +/* Align navigation buttons with dropdown bottom */ +.navigation-row { + display: flex !important; + align-items: flex-end !important; + gap: 8px !important; +} + +.navigation-row > div:nth-child(1), +.navigation-row > div:nth-child(3) { + align-self: flex-end !important; +} + +.navigation-row > div:nth-child(2) { + flex: 1 !important; +} + +/* Make thumbnails clickable with pointer cursor */ +.clickable-thumbnail img { + cursor: pointer !important; +} + +.clickable-thumbnail:hover img { + cursor: pointer !important; + opacity: 0.8; + transition: opacity 0.3s ease; +} + +/* Make thumbnail containers narrower horizontally */ +.clickable-thumbnail { + padding: 5px 2px !important; + margin: 0 2px !important; +} + +.clickable-thumbnail .image-container { + margin: 0 !important; + padding: 0 !important; +} + +.scene-info { + text-align: center !important; + padding: 5px 2px !important; + margin: 0 !important; +} +""" + + +def get_header_html(logo_base64=None): + """ + Generate the main header HTML with logo and title. + + Args: + logo_base64 (str, optional): Base64 encoded logo image + + Returns: + str: HTML string for the header + """ + return """ +
+
+

+ Depth Anything 3 +

+

+ Recovering the Visual Space from Any Views +

+
+ + + Project Page + + + Paper + + + Code + +
+
+
+ + + """ + + +def get_description_html(): + """ + Generate the main description and getting started HTML. + + Returns: + str: HTML string for the description + """ + return """ +
+

+ What This Demo Does +

+
+

+ Upload images or videos → Get Metric Point Clouds, Cameras and Novel Views → Explore in 3D +

+
+ +
+

+ Tip: Landscape-oriented images or videos are preferred for best 3D recovering. +

+
+
+ + + """ + + +def get_acknowledgements_html(): + """ + Generate the acknowledgements section HTML. + + Returns: + str: HTML string for the acknowledgements + """ + return """ +
+

+ Research Credits & Acknowledgments +

+ +
+ +
+

Original Research

+

+ + Depth Anything 3 + +

+
+ + +
+

Previous Versions

+
+

+ + Depth-Anything + +

+ • +

+ + Depth-Anything-V2 + +

+
+
+
+ + +
+

+ HF demo adapted from Map Anything +

+
+
+ """ + + +def get_gradio_theme(): + """ + Get the configured Gradio theme with adaptive tech colors. + + Returns: + gr.themes.Base: Configured Gradio theme + """ + import gradio as gr + + return gr.themes.Base( + primary_hue=gr.themes.Color( + c50="#eff6ff", + c100="#dbeafe", + c200="#bfdbfe", + c300="#93c5fd", + c400="#60a5fa", + c500="#3b82f6", + c600="#2563eb", + c700="#1d4ed8", + c800="#1e40af", + c900="#1e3a8a", + c950="#172554", + ), + secondary_hue=gr.themes.Color( + c50="#f5f3ff", + c100="#ede9fe", + c200="#ddd6fe", + c300="#c4b5fd", + c400="#a78bfa", + c500="#8b5cf6", + c600="#7c3aed", + c700="#6d28d9", + c800="#5b21b6", + c900="#4c1d95", + c950="#2e1065", + ), + neutral_hue=gr.themes.Color( + c50="#f8fafc", + c100="#f1f5f9", + c200="#e2e8f0", + c300="#cbd5e1", + c400="#94a3b8", + c500="#64748b", + c600="#475569", + c700="#334155", + c800="#1e293b", + c900="#0f172a", + c950="#020617", + ), + ) + + +# Measure tab instructions HTML +MEASURE_INSTRUCTIONS_HTML = """ +### Click points on the image to compute distance. +> Metric scale estimation is difficult on aerial/drone images. +""" diff --git a/depth_anything_3/app/gradio_app.py b/depth_anything_3/app/gradio_app.py new file mode 100644 index 0000000000000000000000000000000000000000..ef188c0994c3562ef0671e2ff918a534566d38ca --- /dev/null +++ b/depth_anything_3/app/gradio_app.py @@ -0,0 +1,724 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Refactored Gradio App for Depth Anything 3. + +This is the main application file that orchestrates all components. +The original functionality has been split into modular components for better maintainability. +""" + +import argparse +import os +from typing import Any, Dict, List +import gradio as gr + +from depth_anything_3.app.css_and_html import GRADIO_CSS, get_gradio_theme +from depth_anything_3.app.modules.event_handlers import EventHandlers +from depth_anything_3.app.modules.ui_components import UIComponents + +# Set environment variables +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + + +class DepthAnything3App: + """ + Main application class for Depth Anything 3 Gradio app. + """ + + def __init__(self, model_dir: str = None, workspace_dir: str = None, gallery_dir: str = None): + """ + Initialize the application. + + Args: + model_dir: Path to the model directory + workspace_dir: Path to the workspace directory + gallery_dir: Path to the gallery directory + """ + self.model_dir = model_dir + self.workspace_dir = workspace_dir + self.gallery_dir = gallery_dir + + # Set environment variables for directories + if self.model_dir: + os.environ["DA3_MODEL_DIR"] = self.model_dir + if self.workspace_dir: + os.environ["DA3_WORKSPACE_DIR"] = self.workspace_dir + if self.gallery_dir: + os.environ["DA3_GALLERY_DIR"] = self.gallery_dir + + self.event_handlers = EventHandlers() + self.ui_components = UIComponents() + + def cache_examples( + self, + show_cam: bool = True, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + save_percentage: float = 20.0, + num_max_points: int = 1000, + cache_gs_tag: str = "", + gs_trj_mode: str = "smooth", + gs_video_quality: str = "low", + ) -> None: + """ + Pre-cache all example scenes at startup. + + Args: + show_cam: Whether to show camera in visualization + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + save_percentage: Filter percentage for point cloud + num_max_points: Maximum number of points + cache_gs_tag: Tag to match scene names for high-res+3DGS caching (e.g., "dl3dv") + gs_trj_mode: Trajectory mode for 3DGS + gs_video_quality: Video quality for 3DGS + """ + from depth_anything_3.app.modules.utils import get_scene_info + + examples_dir = os.path.join(self.workspace_dir, "examples") + if not os.path.exists(examples_dir): + print(f"Examples directory not found: {examples_dir}") + return + + scenes = get_scene_info(examples_dir) + if not scenes: + print("No example scenes found to cache.") + return + + print(f"\n{'='*60}") + print(f"Caching {len(scenes)} example scenes...") + print(f"{'='*60}\n") + + for i, scene in enumerate(scenes, 1): + scene_name = scene["name"] + + # Check if scene name matches the gs tag for high-res+3DGS caching + use_high_res_gs = cache_gs_tag and cache_gs_tag.lower() in scene_name.lower() + + if use_high_res_gs: + print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (HIGH-RES + 3DGS)") + print(f" - Number of images: {scene['num_images']}") + print(f" - Matched tag: '{cache_gs_tag}' - using high_res + 3DGS") + else: + print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (LOW-RES)") + print(f" - Number of images: {scene['num_images']}") + + try: + # Load example scene + _, target_dir, _, _, _, _, _, _, _ = self.event_handlers.load_example_scene( + scene_name + ) + + if target_dir and target_dir != "None": + # Run reconstruction with appropriate settings + print(" - Running reconstruction...") + result = self.event_handlers.gradio_demo( + target_dir=target_dir, + show_cam=show_cam, + filter_black_bg=filter_black_bg, + filter_white_bg=filter_white_bg, + process_res_method="high_res" if use_high_res_gs else "low_res", + save_percentage=save_percentage, + num_max_points=num_max_points, + infer_gs=use_high_res_gs, + ref_view_strategy="saddle_balanced", + gs_trj_mode=gs_trj_mode, + gs_video_quality=gs_video_quality, + ) + + # Check if successful + if result[0] is not None: # reconstruction_output + print(f" ✓ Scene '{scene_name}' cached successfully") + else: + print(f" ✗ Scene '{scene_name}' caching failed: {result[1]}") + else: + print(f" ✗ Scene '{scene_name}' loading failed") + + except Exception as e: + print(f" ✗ Error caching scene '{scene_name}': {str(e)}") + + print() + + print("=" * 60) + print("Example scene caching completed!") + print("=" * 60 + "\n") + + def create_app(self) -> gr.Blocks: + """ + Create and configure the Gradio application. + + Returns: + Configured Gradio Blocks interface + """ + + # Initialize theme + def get_theme(): + return get_gradio_theme() + + with gr.Blocks(theme=get_theme(), css=GRADIO_CSS) as demo: + # State variables for the tabbed interface + is_example = gr.Textbox(label="is_example", visible=False, value="None") + processed_data_state = gr.State(value=None) + measure_points_state = gr.State(value=[]) + selected_image_index_state = gr.State(value=0) # Track selected image index + # current_view_index = gr.State(value=0) # noqa: F841 Track current view index + + # Header and description + self.ui_components.create_header_section() + self.ui_components.create_description_section() + + target_dir_output = gr.Textbox(label="Target Dir", visible=False, value="None") + + # Main content area + with gr.Row(): + with gr.Column(scale=2): + # Upload section + ( + input_video, + s_time_interval, + input_images, + image_gallery, + ) = self.ui_components.create_upload_section() + + with gr.Column(scale=4): + with gr.Column(): + # gr.Markdown("**Metric 3D Reconstruction (Point Cloud and Camera Poses)**") + # Reconstruction control section (buttons) - moved below tabs + + log_output = gr.Markdown( + "Please upload a video or images, then click Reconstruct.", + elem_classes=["custom-log"], + ) + + # Tabbed interface + with gr.Tabs(): + with gr.Tab("Point Cloud & Cameras"): + reconstruction_output = ( + self.ui_components.create_3d_viewer_section() + ) + + with gr.Tab("Metric Depth"): + ( + prev_measure_btn, + measure_view_selector, + next_measure_btn, + measure_image, + measure_depth_image, + measure_text, + ) = self.ui_components.create_measure_section() + + with gr.Tab("3DGS Rendered Novel Views"): + gs_video, gs_info = self.ui_components.create_nvs_video() + + # Inference control section (before inference) + (process_res_method_dropdown, infer_gs, ref_view_strategy_dropdown) = ( + self.ui_components.create_inference_control_section() + ) + + # Display control section - includes 3DGS options, buttons, and Visualization Options # noqa: E501 + ( + show_cam, + filter_black_bg, + filter_white_bg, + save_percentage, + num_max_points, + gs_trj_mode, + gs_video_quality, + submit_btn, + clear_btn, + ) = self.ui_components.create_display_control_section() + + # bind visibility of gs_trj_mode to infer_gs + infer_gs.change( + fn=lambda checked: ( + gr.update(visible=checked), + gr.update(visible=checked), + gr.update(visible=checked), + gr.update(visible=(not checked)), + ), + inputs=infer_gs, + outputs=[gs_trj_mode, gs_video_quality, gs_video, gs_info], + ) + + # Example scenes section + gr.Markdown("## Example Scenes") + + scenes = self.ui_components.create_example_scenes_section() + scene_components = self.ui_components.create_example_scene_grid(scenes) + + # Set up event handlers + self._setup_event_handlers( + demo, + is_example, + processed_data_state, + measure_points_state, + target_dir_output, + input_video, + input_images, + s_time_interval, + image_gallery, + reconstruction_output, + log_output, + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + save_percentage, + submit_btn, + clear_btn, + num_max_points, + infer_gs, + ref_view_strategy_dropdown, + selected_image_index_state, + measure_view_selector, + measure_image, + measure_depth_image, + measure_text, + prev_measure_btn, + next_measure_btn, + scenes, + scene_components, + gs_video, + gs_info, + gs_trj_mode, + gs_video_quality, + ) + + # Acknowledgements + self.ui_components.create_acknowledgements_section() + + return demo + + def _setup_event_handlers( + self, + demo: gr.Blocks, + is_example: gr.Textbox, + processed_data_state: gr.State, + measure_points_state: gr.State, + target_dir_output: gr.Textbox, + input_video: gr.Video, + input_images: gr.File, + s_time_interval: gr.Slider, + image_gallery: gr.Gallery, + reconstruction_output: gr.Model3D, + log_output: gr.Markdown, + show_cam: gr.Checkbox, + filter_black_bg: gr.Checkbox, + filter_white_bg: gr.Checkbox, + process_res_method_dropdown: gr.Dropdown, + save_percentage: gr.Slider, + submit_btn: gr.Button, + clear_btn: gr.ClearButton, + num_max_points: gr.Slider, + infer_gs: gr.Checkbox, + ref_view_strategy_dropdown: gr.Dropdown, + selected_image_index_state: gr.State, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + measure_text: gr.Markdown, + prev_measure_btn: gr.Button, + next_measure_btn: gr.Button, + scenes: List[Dict[str, Any]], + scene_components: List[gr.Image], + gs_video: gr.Video, + gs_info: gr.Markdown, + gs_trj_mode: gr.Dropdown, + gs_video_quality: gr.Dropdown, + ) -> None: + """ + Set up all event handlers for the application. + + Args: + demo: Gradio Blocks interface + All other arguments: Gradio components to connect + """ + # Configure clear button + clear_btn.add( + [ + input_video, + input_images, + reconstruction_output, + log_output, + target_dir_output, + image_gallery, + gs_video, + ] + ) + + # Main reconstruction button + submit_btn.click( + fn=self.event_handlers.clear_fields, inputs=[], outputs=[reconstruction_output] + ).then(fn=self.event_handlers.update_log, inputs=[], outputs=[log_output]).then( + fn=self.event_handlers.gradio_demo, + inputs=[ + target_dir_output, + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + save_percentage, + # pass num_max_points + num_max_points, + infer_gs, + ref_view_strategy_dropdown, + gs_trj_mode, + gs_video_quality, + ], + outputs=[ + reconstruction_output, + log_output, + processed_data_state, + measure_image, + measure_depth_image, + measure_text, + measure_view_selector, + gs_video, + gs_video, # gs_video visibility + gs_info, # gs_info visibility + ], + ).then( + fn=lambda: "False", + inputs=[], + outputs=[is_example], # set is_example to "False" + ) + + # Real-time visualization updates + self._setup_visualization_handlers( + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + target_dir_output, + is_example, + reconstruction_output, + log_output, + ) + + # File upload handlers + input_video.change( + fn=self.event_handlers.handle_uploads, + inputs=[input_video, input_images, s_time_interval], + outputs=[reconstruction_output, target_dir_output, image_gallery, log_output], + ) + input_images.change( + fn=self.event_handlers.handle_uploads, + inputs=[input_video, input_images, s_time_interval], + outputs=[reconstruction_output, target_dir_output, image_gallery, log_output], + ) + + # Navigation handlers + self._setup_navigation_handlers( + prev_measure_btn, + next_measure_btn, + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + processed_data_state, + ) + + # Measurement handler + measure_image.select( + fn=self.event_handlers.measure, + inputs=[processed_data_state, measure_points_state, measure_view_selector], + outputs=[measure_image, measure_depth_image, measure_points_state, measure_text], + ) + + # Example scene handlers + self._setup_example_scene_handlers( + scenes, + scene_components, + reconstruction_output, + target_dir_output, + image_gallery, + log_output, + is_example, + processed_data_state, + measure_view_selector, + measure_image, + measure_depth_image, + gs_video, + gs_info, + ) + + def _setup_visualization_handlers( + self, + show_cam: gr.Checkbox, + filter_black_bg: gr.Checkbox, + filter_white_bg: gr.Checkbox, + process_res_method_dropdown: gr.Dropdown, + target_dir_output: gr.Textbox, + is_example: gr.Textbox, + reconstruction_output: gr.Model3D, + log_output: gr.Markdown, + ) -> None: + """Set up visualization update handlers.""" + # Common inputs for visualization updates + viz_inputs = [ + target_dir_output, + show_cam, + is_example, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + ] + + # Set up change handlers for all visualization controls + for component in [show_cam, filter_black_bg, filter_white_bg]: + component.change( + fn=self.event_handlers.update_visualization, + inputs=viz_inputs, + outputs=[reconstruction_output, log_output], + ) + + def _setup_navigation_handlers( + self, + prev_measure_btn: gr.Button, + next_measure_btn: gr.Button, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + measure_points_state: gr.State, + processed_data_state: gr.State, + ) -> None: + """Set up navigation handlers for measure tab.""" + # Measure tab navigation + prev_measure_btn.click( + fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view( + processed_data, current_selector, -1 + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[ + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + ], + ) + + next_measure_btn.click( + fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view( + processed_data, current_selector, 1 + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[ + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + ], + ) + + measure_view_selector.change( + fn=lambda processed_data, selector_value: ( + self.event_handlers.update_measure_view( + processed_data, int(selector_value.split()[1]) - 1 + ) + if selector_value + else (None, None, []) + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[measure_image, measure_depth_image, measure_points_state], + ) + + def _setup_example_scene_handlers( + self, + scenes: List[Dict[str, Any]], + scene_components: List[gr.Image], + reconstruction_output: gr.Model3D, + target_dir_output: gr.Textbox, + image_gallery: gr.Gallery, + log_output: gr.Markdown, + is_example: gr.Textbox, + processed_data_state: gr.State, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + gs_video: gr.Video, + gs_info: gr.Markdown, + ) -> None: + """Set up example scene handlers.""" + + def load_and_update_measure(name): + result = self.event_handlers.load_example_scene(name) + # result = (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501 + + # Update measure view if processed_data is available + measure_img = None + measure_depth = None + if result[4] is not None: # processed_data exists + measure_img, measure_depth, _ = ( + self.event_handlers.visualization_handler.update_measure_view(result[4], 0) + ) + + return result + ("True", measure_img, measure_depth) + + for i, scene in enumerate(scenes): + if i < len(scene_components): + scene_components[i].select( + fn=lambda name=scene["name"]: load_and_update_measure(name), + outputs=[ + reconstruction_output, + target_dir_output, + image_gallery, + log_output, + processed_data_state, + measure_view_selector, + gs_video, + gs_video, # gs_video_visibility + gs_info, # gs_info_visibility + is_example, + measure_image, + measure_depth_image, + ], + ) + + def launch(self, host: str = "127.0.0.1", port: int = 7860, **kwargs) -> None: + """ + Launch the application. + + Args: + host: Host address to bind to + port: Port number to bind to + **kwargs: Additional arguments for demo.launch() + """ + demo = self.create_app() + demo.queue(max_size=20).launch( + show_error=True, ssr_mode=False, server_name=host, server_port=port, **kwargs + ) + + +def main(): + """Main function to run the application.""" + parser = argparse.ArgumentParser( + description="Depth Anything 3 Gradio Application", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Basic usage + python gradio_app.py --help + python gradio_app.py --host 0.0.0.0 --port 8080 + python gradio_app.py --model-dir /path/to/model --workspace-dir /path/to/workspace + + # Cache examples at startup (all low-res) + python gradio_app.py --cache-examples + + # Cache with selective high-res+3DGS for scenes matching tag + python gradio_app.py --cache-examples --cache-gs-tag dl3dv + # This will use high-res + 3DGS for scenes containing "dl3dv" in their name, + # and low-res only for other scenes + """, + ) + + # Server configuration + parser.add_argument( + "--host", default="127.0.0.1", help="Host address to bind to (default: 127.0.0.1)" + ) + parser.add_argument( + "--port", type=int, default=7860, help="Port number to bind to (default: 7860)" + ) + + # Directory configuration + parser.add_argument( + "--model-dir", + default="depth-anything/DA3NESTED-GIANT-LARGE", + help="Path to the model directory (default: depth-anything/DA3NESTED-GIANT-LARGE)", + ) + parser.add_argument( + "--workspace-dir", + default="workspace/gradio", # noqa: E501 + help="Path to the workspace directory (default: workspace/gradio)", # noqa: E501 + ) + parser.add_argument( + "--gallery-dir", + default="workspace/gallery", + help="Path to the gallery directory (default: workspace/gallery)", # noqa: E501 + ) + + # Additional Gradio options + parser.add_argument("--share", action="store_true", help="Create a public link for the app") + parser.add_argument("--debug", action="store_true", help="Enable debug mode") + + # Example caching options + parser.add_argument( + "--cache-examples", + action="store_true", + help="Pre-cache all example scenes at startup for faster loading", + ) + parser.add_argument( + "--cache-gs-tag", + type=str, + default="", + help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.", # noqa: E501 + ) + + args = parser.parse_args() + + # Create directories if they don't exist + os.makedirs(args.workspace_dir, exist_ok=True) + os.makedirs(args.gallery_dir, exist_ok=True) + + # Initialize and launch the application + app = DepthAnything3App( + model_dir=args.model_dir, workspace_dir=args.workspace_dir, gallery_dir=args.gallery_dir + ) + + # Prepare launch arguments + launch_kwargs = {"share": args.share, "debug": args.debug} + + print("Starting Depth Anything 3 Gradio App...") + print(f"Host: {args.host}") + print(f"Port: {args.port}") + print(f"Model Directory: {args.model_dir}") + print(f"Workspace Directory: {args.workspace_dir}") + print(f"Gallery Directory: {args.gallery_dir}") + print(f"Share: {args.share}") + print(f"Debug: {args.debug}") + print(f"Cache Examples: {args.cache_examples}") + if args.cache_examples: + if args.cache_gs_tag: + print( + f"Cache GS Tag: '{args.cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)" # noqa: E501 + ) # noqa: E501 + else: + print("Cache GS Tag: None (all scenes will use low-res only)") + + # Pre-cache examples if requested + if args.cache_examples: + print("\n" + "=" * 60) + print("Pre-caching mode enabled") + if args.cache_gs_tag: + print(f"Scenes containing '{args.cache_gs_tag}' will use HIGH-RES + 3DGS") + print("Other scenes will use LOW-RES only") + else: + print("All scenes will use LOW-RES only") + print("=" * 60) + app.cache_examples( + show_cam=True, + filter_black_bg=False, + filter_white_bg=False, + save_percentage=5.0, + num_max_points=1000, + cache_gs_tag=args.cache_gs_tag, + gs_trj_mode="smooth", + gs_video_quality="low", + ) + + app.launch(host=args.host, port=args.port, **launch_kwargs) + + +if __name__ == "__main__": + main() diff --git a/depth_anything_3/app/modules/__init__.py b/depth_anything_3/app/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd0b71780214eeadbc4fc44f4ac070e0e5fa7795 --- /dev/null +++ b/depth_anything_3/app/modules/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Modules package for Depth Anything 3 Gradio app. + +This package contains all the modular components for the Gradio application. +""" + +from depth_anything_3.app.modules.event_handlers import EventHandlers +from depth_anything_3.app.modules.file_handlers import FileHandler +from depth_anything_3.app.modules.model_inference import ModelInference +from depth_anything_3.app.modules.ui_components import UIComponents +from depth_anything_3.app.modules.utils import ( + create_depth_visualization, + get_logo_base64, + get_scene_info, + save_to_gallery_func, +) +from depth_anything_3.app.modules.visualization import VisualizationHandler + +__all__ = [ + "ModelInference", + "FileHandler", + "VisualizationHandler", + "EventHandlers", + "UIComponents", + "create_depth_visualization", + "save_to_gallery_func", + "get_scene_info", + "get_logo_base64", +] diff --git a/depth_anything_3/app/modules/event_handlers.py b/depth_anything_3/app/modules/event_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..74684cca2ef641e5489ec4bcc5cdb28a301023b0 --- /dev/null +++ b/depth_anything_3/app/modules/event_handlers.py @@ -0,0 +1,622 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Event handling module for Depth Anything 3 Gradio app. + +This module handles all event callbacks and user interactions. +""" + +import os +import time +from glob import glob +from typing import Any, Dict, List, Optional, Tuple +import gradio as gr +import numpy as np +import torch + +from depth_anything_3.app.modules.file_handlers import FileHandler +from depth_anything_3.app.modules.model_inference import ModelInference +from depth_anything_3.utils.memory import cleanup_cuda_memory +from depth_anything_3.app.modules.visualization import VisualizationHandler + + +class EventHandlers: + """ + Handles all event callbacks and user interactions for the Gradio app. + """ + + def __init__(self): + """Initialize the event handlers.""" + self.model_inference = ModelInference() + self.file_handler = FileHandler() + self.visualization_handler = VisualizationHandler() + + def clear_fields(self) -> None: + """ + Clears the 3D viewer, the stored target_dir, and empties the gallery. + """ + return None + + def update_log(self) -> str: + """ + Display a quick log message while waiting. + """ + return "Loading and Reconstructing..." + + def save_current_visualization( + self, + target_dir: str, + save_percentage: float, + show_cam: bool, + filter_black_bg: bool, + filter_white_bg: bool, + processed_data: Optional[Dict], + scene_name: str = "", + ) -> str: + """ + Save current visualization results to gallery with specified save percentage. + + Args: + target_dir: Directory containing results + save_percentage: Percentage of points to save (0-100) + show_cam: Whether to show cameras + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + processed_data: Processed data from reconstruction + + Returns: + Status message + """ + if not self.file_handler.is_valid_session_dir(target_dir): + return "No reconstruction available. Please run 'Reconstruct' first." + + if processed_data is None: + return "No processed data available. Please run 'Reconstruct' first." + + try: + # Add debug information + print("[DEBUG] save_current_visualization called with:") + print(f" target_dir: {target_dir}") + print(f" save_percentage: {save_percentage}") + print(f" show_cam: {show_cam}") + print(f" filter_black_bg: {filter_black_bg}") + print(f" filter_white_bg: {filter_white_bg}") + print(f" processed_data: {processed_data is not None}") + + # Import the gallery save function + # Create gallery name with user input or auto-generated + import datetime + + from .utils import save_to_gallery_func + + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + if scene_name and scene_name.strip(): + gallery_name = f"{scene_name.strip()}_{timestamp}_pct{save_percentage:.0f}" + else: + gallery_name = f"save_{timestamp}_pct{save_percentage:.0f}" + + print(f"[DEBUG] Saving to gallery with name: {gallery_name}") + + # Save entire process folder to gallery + success, message = save_to_gallery_func( + target_dir=target_dir, processed_data=processed_data, gallery_name=gallery_name + ) + + if success: + print(f"[DEBUG] Gallery save completed successfully: {message}") + return ( + "Successfully saved to gallery!\n" + f"Gallery name: {gallery_name}\n" + f"Save percentage: {save_percentage}%\n" + f"Show cameras: {show_cam}\n" + f"Filter black bg: {filter_black_bg}\n" + f"Filter white bg: {filter_white_bg}\n\n" + f"{message}" + ) + else: + print(f"[DEBUG] Gallery save failed: {message}") + return f"Failed to save to gallery: {message}" + + except Exception as e: + return f"Error saving visualization: {str(e)}" + + def gradio_demo( + self, + target_dir: str, + show_cam: bool = True, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + save_percentage: float = 30.0, + num_max_points: int = 1_000_000, + infer_gs: bool = False, + ref_view_strategy: str = "saddle_balanced", + gs_trj_mode: str = "extend", + gs_video_quality: str = "high", + ) -> Tuple[ + Optional[str], + str, + Optional[Dict], + Optional[np.ndarray], + Optional[np.ndarray], + str, + gr.Dropdown, + Optional[str], # gs video path + gr.update, # gs video visibility update + gr.update, # gs info visibility update + ]: + """ + Perform reconstruction using the already-created target_dir/images. + + Args: + target_dir: Directory containing images + show_cam: Whether to show camera + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + save_percentage: Filter percentage for point cloud + num_max_points: Maximum number of points + infer_gs: Whether to infer 3D Gaussian Splatting + ref_view_strategy: Reference view selection strategy + + Returns: + Tuple of reconstruction results + """ + if not self.file_handler.is_valid_session_dir(target_dir): + return ( + None, + "No valid target directory found. Please upload first.", + None, + None, + None, + "", + None, + None, + gr.update(visible=False), # gs_video + gr.update(visible=True), # gs_info + ) + + start_time = time.time() + cleanup_cuda_memory() + + # Get image files for logging + target_dir_images = os.path.join(target_dir, "images") + all_files = ( + sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else [] + ) + + print("Running DepthAnything3 model...") + print(f"Reference view strategy: {ref_view_strategy}") + + with torch.no_grad(): + prediction, processed_data = self.model_inference.run_inference( + target_dir, + process_res_method=process_res_method, + show_camera=show_cam, + save_percentage=save_percentage, + num_max_points=int(num_max_points * 1000), # Convert K to actual count + infer_gs=infer_gs, + ref_view_strategy=ref_view_strategy, + gs_trj_mode=gs_trj_mode, + gs_video_quality=gs_video_quality, + ) + + # The GLB file is already generated by the API + glbfile = os.path.join(target_dir, "scene.glb") + + # Handle 3DGS video based on infer_gs flag + gsvideo_path = None + gs_video_visible = False + gs_info_visible = True + + if infer_gs: + try: + gsvideo_path = sorted(glob(os.path.join(target_dir, "gs_video", "*.mp4")))[-1] + gs_video_visible = True + gs_info_visible = False + except IndexError: + gsvideo_path = None + print("3DGS video not found, but infer_gs was enabled") + + # Cleanup + cleanup_cuda_memory() + + end_time = time.time() + print(f"Total time: {end_time - start_time:.2f} seconds") + log_msg = f"Reconstruction Success ({len(all_files)} frames). Waiting for visualization." + + # Populate visualization tabs with processed data + depth_vis, measure_img, measure_depth_vis, measure_pts = ( + self.visualization_handler.populate_visualization_tabs(processed_data) + ) + + # Update view selectors based on available views + depth_selector, measure_selector = self.visualization_handler.update_view_selectors( + processed_data + ) + + return ( + glbfile, + log_msg, + processed_data, + measure_img, # measure_image + measure_depth_vis, # measure_depth_image + "", # measure_text (empty initially) + measure_selector, # measure_view_selector + gsvideo_path, + gr.update(visible=gs_video_visible), # gs_video visibility + gr.update(visible=gs_info_visible), # gs_info visibility + ) + + def update_visualization( + self, + target_dir: str, + show_cam: bool, + is_example: str, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + ) -> Tuple[gr.update, str]: + """ + Reload saved predictions from npz, create (or reuse) the GLB for new parameters, + and return it for the 3D viewer. + + Args: + target_dir: Directory containing results + show_cam: Whether to show camera + is_example: Whether this is an example scene + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + + Returns: + Tuple of (glb_file, log_message) + """ + if not self.file_handler.is_valid_session_dir(target_dir): + return ( + gr.update(), + "No reconstruction available. Please click the Reconstruct button first.", + ) + + # Check if GLB exists (could be cached example or reconstructed scene) + glbfile = os.path.join(target_dir, "scene.glb") + if os.path.exists(glbfile): + return ( + glbfile, + ( + "Visualization loaded from cache." + if is_example == "True" + else "Visualization updated." + ), + ) + + # If no GLB but it's an example that hasn't been reconstructed yet + if is_example == "True": + return ( + gr.update(), + "No reconstruction available. Please click the Reconstruct button first.", + ) + + # For non-examples, check predictions.npz + predictions_path = os.path.join(target_dir, "predictions.npz") + if not os.path.exists(predictions_path): + error_message = ( + f"No reconstruction available at {predictions_path}. " + "Please run 'Reconstruct' first." + ) + return gr.update(), error_message + + try: + loaded = np.load(predictions_path, allow_pickle=False) + predictions = {key: loaded[key] for key in loaded.keys()} # noqa: F841 + except Exception as e: + return gr.update(), f"Cached results could not be read: {e}" + + return ( + glbfile, + "Visualization updated.", + ) + + def handle_uploads( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]: + """ + Handle file uploads and update gallery. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + return self.file_handler.update_gallery_on_upload( + input_video, input_images, s_time_interval + ) + + def load_example_scene(self, scene_name: str, examples_dir: str = None) -> Tuple[ + Optional[str], + Optional[str], + Optional[List], + str, + Optional[Dict], + gr.Dropdown, + Optional[str], + gr.update, + gr.update, + ]: + """ + Load a scene from examples directory. + + Args: + scene_name: Name of the scene to load + examples_dir: Path to examples directory (if None, uses workspace_dir/examples) + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501 + """ + if examples_dir is None: + # Get workspace directory from environment variable + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + examples_dir = os.path.join(workspace_dir, "examples") + + reconstruction_output, target_dir, image_paths, log_message = ( + self.file_handler.load_example_scene(scene_name, examples_dir) + ) + + # Try to load cached processed data if available + processed_data = None + measure_view_selector = gr.Dropdown(choices=["View 1"], value="View 1") + gs_video_path = None + gs_video_visible = False + gs_info_visible = True + + if self.file_handler.is_valid_session_dir(target_dir): + predictions_path = os.path.join(target_dir, "predictions.npz") + if os.path.exists(predictions_path): + try: + # Load predictions from cache + loaded = np.load(predictions_path, allow_pickle=False) + predictions = {key: loaded[key] for key in loaded.keys()} + + # Reconstruct processed_data structure + num_images = len(predictions.get("images", [])) + processed_data = {} + + for i in range(num_images): + processed_data[i] = { + "image": predictions["images"][i] if "images" in predictions else None, + "depth": predictions["depths"][i] if "depths" in predictions else None, + "depth_image": os.path.join( + target_dir, "depth_vis", f"{i:04d}.jpg" # Fixed: use .jpg not .png + ), + "intrinsics": ( + predictions["intrinsics"][i] + if "intrinsics" in predictions + and i < len(predictions["intrinsics"]) + else None + ), + "mask": None, + } + + # Update measure view selector + choices = [f"View {i + 1}" for i in range(num_images)] + measure_view_selector = gr.Dropdown(choices=choices, value=choices[0]) + + except Exception as e: + print(f"Error loading cached data: {e}") + + # Check for cached 3DGS video + gs_video_dir = os.path.join(target_dir, "gs_video") + if os.path.exists(gs_video_dir): + try: + from glob import glob + + gs_videos = sorted(glob(os.path.join(gs_video_dir, "*.mp4"))) + if gs_videos: + gs_video_path = gs_videos[-1] + gs_video_visible = True + gs_info_visible = False + print(f"Loaded cached 3DGS video: {gs_video_path}") + except Exception as e: + print(f"Error loading cached 3DGS video: {e}") + + return ( + reconstruction_output, + target_dir, + image_paths, + log_message, + processed_data, + measure_view_selector, + gs_video_path, + gr.update(visible=gs_video_visible), + gr.update(visible=gs_info_visible), + ) + + def navigate_depth_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector: str, + direction: int, + ) -> Tuple[str, Optional[str]]: + """ + Navigate depth view. + + Args: + processed_data: Processed data dictionary + current_selector: Current selector value + direction: Direction to navigate + + Returns: + Tuple of (new_selector_value, depth_vis) + """ + return self.visualization_handler.navigate_depth_view( + processed_data, current_selector, direction + ) + + def update_depth_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[str]: + """ + Update depth view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Path to depth visualization image or None + """ + return self.visualization_handler.update_depth_view(processed_data, view_index) + + def navigate_measure_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector: str, + direction: int, + ) -> Tuple[str, Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Navigate measure view. + + Args: + processed_data: Processed data dictionary + current_selector: Current selector value + direction: Direction to navigate + + Returns: + Tuple of (new_selector_value, measure_image, depth_right_half, measure_points) + """ + return self.visualization_handler.navigate_measure_view( + processed_data, current_selector, direction + ) + + def update_measure_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Update measure view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Tuple of (measure_image, depth_right_half, measure_points) + """ + return self.visualization_handler.update_measure_view(processed_data, view_index) + + def measure( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + measure_points: List, + current_view_selector: str, + event: gr.SelectData, + ) -> List: + """ + Handle measurement on images. + + Args: + processed_data: Processed data dictionary + measure_points: List of current measure points + current_view_selector: Current view selector value + event: Gradio select event + + Returns: + List of [image, depth_right_half, measure_points, text] + """ + return self.visualization_handler.measure( + processed_data, measure_points, current_view_selector, event + ) + + def select_first_frame( + self, image_gallery: List, selected_index: int = 0 + ) -> Tuple[List, str, str]: + """ + Select the first frame from the image gallery. + + Args: + image_gallery: List of images in the gallery + selected_index: Index of the selected image (default: 0) + + Returns: + Tuple of (updated_image_gallery, log_message, selected_frame_path) + """ + try: + if not image_gallery or len(image_gallery) == 0: + return image_gallery, "No images available to select as first frame.", "" + + # Handle None or invalid selected_index + if ( + selected_index is None + or selected_index < 0 + or selected_index >= len(image_gallery) + ): + selected_index = 0 + print(f"Invalid selected_index: {selected_index}, using default: 0") + + # Get the selected image based on index + selected_image = image_gallery[selected_index] + print(f"Selected image index: {selected_index}") + print(f"Total images: {len(image_gallery)}") + + # Extract the file path from the selected image + selected_frame_path = "" + print(f"Selected image type: {type(selected_image)}") + print(f"Selected image: {selected_image}") + + if isinstance(selected_image, tuple): + # Gradio Gallery returns tuple (path, None) + selected_frame_path = selected_image[0] + elif isinstance(selected_image, str): + selected_frame_path = selected_image + elif hasattr(selected_image, "name"): + selected_frame_path = selected_image.name + elif isinstance(selected_image, dict): + if "name" in selected_image: + selected_frame_path = selected_image["name"] + elif "path" in selected_image: + selected_frame_path = selected_image["path"] + elif "src" in selected_image: + selected_frame_path = selected_image["src"] + else: + # Try to convert to string + selected_frame_path = str(selected_image) + + print(f"Extracted path: {selected_frame_path}") + + # Extract filename from the path for matching + import os + + selected_filename = os.path.basename(selected_frame_path) + print(f"Selected filename: {selected_filename}") + + # Move the selected image to the front + updated_gallery = [selected_image] + [ + img for img in image_gallery if img != selected_image + ] + + log_message = ( + f"Selected frame: {selected_filename}. " + f"Moved to first position. Total frames: {len(updated_gallery)}" + ) + return updated_gallery, log_message, selected_filename + + except Exception as e: + print(f"Error selecting first frame: {e}") + return image_gallery, f"Error selecting first frame: {e}", "" diff --git a/depth_anything_3/app/modules/file_handlers.py b/depth_anything_3/app/modules/file_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..dd97eba2def821ebebdc770ef910c79db5de2f4a --- /dev/null +++ b/depth_anything_3/app/modules/file_handlers.py @@ -0,0 +1,355 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +File handling module for Depth Anything 3 Gradio app. + +This module handles file uploads, video processing, and file operations. +""" + +import os +import shutil +import time +from datetime import datetime +from typing import List, Optional, Tuple +import cv2 +from PIL import Image +from pillow_heif import register_heif_opener + +register_heif_opener() + + +ALLOWED_IMAGE_EXTENSIONS = { + ".png", + ".jpg", + ".jpeg", + ".bmp", + ".tiff", + ".tif", + ".webp", + ".heic", + ".heif", +} + + +def _looks_like_image(file_path: str) -> bool: + """Best-effort check that a file is actually a decodable image, not just named like one.""" + try: + with Image.open(file_path) as img: + img.verify() + return True + except Exception: + return False + + +class FileHandler: + """ + Handles file uploads and processing for the Gradio app. + """ + + def __init__(self): + """Initialize the file handler.""" + + def is_valid_session_dir(self, target_dir: Optional[str]) -> bool: + """Check that target_dir is a session/example directory this handler itself + creates under the configured workspace, rather than an arbitrary path.""" + if not target_dir or target_dir == "None": + return False + + try: + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + input_images_dir = os.path.realpath(os.path.join(workspace_dir, "input_images")) + real_target = os.path.realpath(target_dir) + + if real_target != input_images_dir and not real_target.startswith( + input_images_dir + os.sep + ): + return False + + rel = os.path.relpath(real_target, input_images_dir) + if os.sep in rel or not (rel.startswith("session_") or rel.startswith("example_")): + return False + + return os.path.isdir(real_target) + except (OSError, ValueError): + return False + + def handle_uploads( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[str, List[str]]: + """ + Create a new 'target_dir' + 'images' subfolder, and place user-uploaded + images or extracted frames from video into it. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (target_dir, image_paths) + """ + start_time = time.time() + + # Get workspace directory from environment variable or use default + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + if not os.path.exists(workspace_dir): + os.makedirs(workspace_dir) + + # Create input_images subdirectory + input_images_dir = os.path.join(workspace_dir, "input_images") + if not os.path.exists(input_images_dir): + os.makedirs(input_images_dir) + + # Create a unique folder name within input_images + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + target_dir = os.path.join(input_images_dir, f"session_{timestamp}") + target_dir_images = os.path.join(target_dir, "images") + + # Clean up if somehow that folder already exists + if os.path.exists(target_dir): + shutil.rmtree(target_dir) + os.makedirs(target_dir) + os.makedirs(target_dir_images) + + image_paths = [] + + # Handle images + if input_images is not None: + image_paths.extend(self._process_images(input_images, target_dir_images)) + + # Handle video + if input_video is not None: + image_paths.extend( + self._process_video(input_video, target_dir_images, s_time_interval) + ) + + # Sort final images for gallery + image_paths = sorted(image_paths) + + end_time = time.time() + print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds") + return target_dir, image_paths + + def _process_images(self, input_images: List, target_dir_images: str) -> List[str]: + """ + Process uploaded images. + + Args: + input_images: List of input image files + target_dir_images: Target directory for images + + Returns: + List of processed image paths + """ + image_paths = [] + + for file_data in input_images: + if isinstance(file_data, dict) and "name" in file_data: + file_path = file_data["name"] + else: + file_path = file_data + + file_ext = os.path.splitext(file_path)[1].lower() + if file_ext not in ALLOWED_IMAGE_EXTENSIONS or not _looks_like_image(file_path): + print(f"Skipping non-image upload: {os.path.basename(file_path)}") + continue + + # Check if the file is a HEIC image + if file_ext in [".heic", ".heif"]: + # Convert HEIC to JPEG for better gallery compatibility + try: + with Image.open(file_path) as img: + # Convert to RGB if necessary (HEIC can have different color modes) + if img.mode not in ("RGB", "L"): + img = img.convert("RGB") + + # Create JPEG filename + base_name = os.path.splitext(os.path.basename(file_path))[0] + dst_path = os.path.join(target_dir_images, f"{base_name}.jpg") + + # Save as JPEG with high quality + img.save(dst_path, "JPEG", quality=95) + image_paths.append(dst_path) + print( + f"Converted HEIC to JPEG: {os.path.basename(file_path)} -> " + f"{os.path.basename(dst_path)}" + ) + except Exception as e: + print(f"Error converting HEIC file {file_path}: {e}") + # Fall back to copying as is + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + else: + # Regular image files - copy as is + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + + return image_paths + + def _process_video( + self, input_video: str, target_dir_images: str, s_time_interval: float + ) -> List[str]: + """ + Process video file and extract frames. + + Args: + input_video: Path to input video file + target_dir_images: Target directory for extracted frames + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + List of extracted frame paths + """ + image_paths = [] + + if isinstance(input_video, dict) and "name" in input_video: + video_path = input_video["name"] + else: + video_path = input_video + + vs = cv2.VideoCapture(video_path) + fps = vs.get(cv2.CAP_PROP_FPS) + frame_interval = max(1, int(fps / s_time_interval)) # Convert FPS to frame interval + + count = 0 + video_frame_num = 0 + while True: + gotit, frame = vs.read() + if not gotit: + break + count += 1 + if count % frame_interval == 0: + image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png") + cv2.imwrite(image_path, frame) + image_paths.append(image_path) + video_frame_num += 1 + + return image_paths + + def update_gallery_on_upload( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]: + """ + Handle file uploads and update gallery. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + if not input_video and not input_images: + return None, None, None, None + + target_dir, image_paths = self.handle_uploads(input_video, input_images, s_time_interval) + return ( + None, + target_dir, + image_paths, + "Upload complete. Click 'Reconstruct' to begin 3D processing.", + ) + + def load_example_scene( + self, scene_name: str, examples_dir: str = "examples" + ) -> Tuple[Optional[str], Optional[str], Optional[List], str]: + """ + Load a scene from examples directory. + + Args: + scene_name: Name of the scene to load + examples_dir: Path to examples directory + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + from depth_anything_3.app.modules.utils import get_scene_info + + scenes = get_scene_info(examples_dir) + + # Find the selected scene + selected_scene = None + for scene in scenes: + if scene["name"] == scene_name: + selected_scene = scene + break + + if selected_scene is None: + return None, None, None, "Scene not found" + + # Use fixed directory name for examples (not timestamp-based) + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + input_images_dir = os.path.join(workspace_dir, "input_images") + if not os.path.exists(input_images_dir): + os.makedirs(input_images_dir) + + # Create a fixed folder name based on scene name + target_dir = os.path.join(input_images_dir, f"example_{scene_name}") + target_dir_images = os.path.join(target_dir, "images") + + # Check if already cached (GLB file exists) + glb_path = os.path.join(target_dir, "scene.glb") + is_cached = os.path.exists(glb_path) + + # Create directory if it doesn't exist + if not os.path.exists(target_dir): + os.makedirs(target_dir) + os.makedirs(target_dir_images) + + # Copy images if directory is new or empty + if not os.path.exists(target_dir_images) or len(os.listdir(target_dir_images)) == 0: + os.makedirs(target_dir_images, exist_ok=True) + image_paths = [] + for file_path in selected_scene["image_files"]: + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + else: + # Use existing images + image_paths = sorted( + [ + os.path.join(target_dir_images, f) + for f in os.listdir(target_dir_images) + if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif")) + ] + ) + + # Return cached GLB if available + if is_cached: + return ( + glb_path, # Return cached reconstruction + target_dir, # Set target directory + image_paths, # Set gallery + f"Loaded cached scene '{scene_name}' with {selected_scene['num_images']} images.", + ) + else: + return ( + None, # No cached reconstruction + target_dir, # Set target directory + image_paths, # Set gallery + ( + f"Loaded scene '{scene_name}' with {selected_scene['num_images']} images. " + "Click 'Reconstruct' to begin 3D processing." + ), + ) diff --git a/depth_anything_3/app/modules/model_inference.py b/depth_anything_3/app/modules/model_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..6ce77b4388b303f95c2f0a9da5c81b2dc55f38f6 --- /dev/null +++ b/depth_anything_3/app/modules/model_inference.py @@ -0,0 +1,260 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model inference module for Depth Anything 3 Gradio app. + +This module handles all model-related operations including inference, +data processing, and result preparation. +""" + +import glob +import os +from typing import Any, Dict, Optional, Tuple +import numpy as np +import torch + +from depth_anything_3.api import DepthAnything3 +from depth_anything_3.utils.memory import cleanup_cuda_memory +from depth_anything_3.utils.export.glb import export_to_glb +from depth_anything_3.utils.export.gs import export_to_gs_video + + +class ModelInference: + """ + Handles model inference and data processing for Depth Anything 3. + """ + + def __init__(self): + """Initialize the model inference handler.""" + self.model = None + + def initialize_model(self, device: str = "cuda") -> None: + """ + Initialize the DepthAnything3 model. + + Args: + device: Device to load the model on + """ + if self.model is None: + # Get model directory from environment variable or use default + model_dir = os.environ.get( + "DA3_MODEL_DIR", "/dev/shm/da3_models/DA3HF-VITG-METRIC_VITL" + ) + self.model = DepthAnything3.from_pretrained(model_dir) + self.model = self.model.to(device) + else: + self.model = self.model.to(device) + + self.model.eval() + + def run_inference( + self, + target_dir: str, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + show_camera: bool = True, + save_percentage: float = 30.0, + num_max_points: int = 1_000_000, + infer_gs: bool = False, + ref_view_strategy: str = "saddle_balanced", + gs_trj_mode: str = "extend", + gs_video_quality: str = "high", + ) -> Tuple[Any, Dict[int, Dict[str, Any]]]: + """ + Run DepthAnything3 model inference on images. + + Args: + target_dir: Directory containing images + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + show_camera: Whether to show camera in 3D view + save_percentage: Percentage of points to save (0-100) + num_max_points: Maximum number of points in point cloud + infer_gs: Whether to infer 3D Gaussian Splatting + ref_view_strategy: Reference view selection strategy + gs_trj_mode: Trajectory mode for 3DGS + gs_video_quality: Video quality for 3DGS + + Returns: + Tuple of (prediction, processed_data) + """ + print(f"Processing images from {target_dir}") + + # Device check + device = "cuda" if torch.cuda.is_available() else "cpu" + device = torch.device(device) + + # Initialize model if needed + self.initialize_model(device) + + # Get image paths + print("Loading images...") + image_folder_path = os.path.join(target_dir, "images") + all_image_paths = sorted(glob.glob(os.path.join(image_folder_path, "*"))) + + # Filter for image files + image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"] + all_image_paths = [ + path + for path in all_image_paths + if any(path.lower().endswith(ext) for ext in image_extensions) + ] + + print(f"Found {len(all_image_paths)} images") + print(f"All image paths: {all_image_paths}") + + # Use sorted image order (reference view will be selected automatically) + image_paths = all_image_paths + print(f"Reference view selection strategy: {ref_view_strategy}") + + if len(image_paths) == 0: + raise ValueError("No images found. Check your upload.") + + # Map UI options to actual method names + method_mapping = {"high_res": "lower_bound_resize", "low_res": "upper_bound_resize"} + actual_method = method_mapping.get(process_res_method, "upper_bound_crop") + + # Run model inference + print(f"Running inference with method: {actual_method}") + with torch.no_grad(): + prediction = self.model.inference( + image_paths, + export_dir=None, + process_res_method=actual_method, + infer_gs=infer_gs, + ref_view_strategy=ref_view_strategy, + ) + # num_max_points: int = 1_000_000, + export_to_glb( + prediction, + filter_black_bg=filter_black_bg, + filter_white_bg=filter_white_bg, + export_dir=target_dir, + show_cameras=show_camera, + conf_thresh_percentile=save_percentage, + num_max_points=int(num_max_points), + ) + + # export to gs video if needed + if infer_gs: + mode_mapping = {"extend": "extend", "smooth": "interpolate_smooth"} + print(f"GS mode: {gs_trj_mode}; Backend mode: {mode_mapping[gs_trj_mode]}") + export_to_gs_video( + prediction, + export_dir=target_dir, + chunk_size=4, + trj_mode=mode_mapping.get(gs_trj_mode, "extend"), + enable_tqdm=True, + vis_depth="hcat", + video_quality=gs_video_quality, + ) + + # Save predictions.npz for caching metric depth data + self._save_predictions_cache(target_dir, prediction) + + # Process results + processed_data = self._process_results(target_dir, prediction, image_paths) + + # Clean up using centralized memory utilities for consistency with backend + cleanup_cuda_memory() + + return prediction, processed_data + + def _save_predictions_cache(self, target_dir: str, prediction: Any) -> None: + """ + Save predictions data to predictions.npz for caching. + + Args: + target_dir: Directory to save the cache + prediction: Model prediction object + """ + try: + output_file = os.path.join(target_dir, "predictions.npz") + + # Build save dict with prediction data + save_dict = {} + + # Save processed images if available + if prediction.processed_images is not None: + save_dict["images"] = prediction.processed_images + + # Save depth data + if prediction.depth is not None: + save_dict["depths"] = np.round(prediction.depth, 6) + + # Save confidence if available + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + + # Save camera parameters + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + # Save to file + np.savez_compressed(output_file, **save_dict) + print(f"Saved predictions cache to: {output_file}") + + except Exception as e: + print(f"Warning: Failed to save predictions cache: {e}") + + def _process_results( + self, target_dir: str, prediction: Any, image_paths: list + ) -> Dict[int, Dict[str, Any]]: + """ + Process model results into structured data. + + Args: + target_dir: Directory containing results + prediction: Model prediction object + image_paths: List of input image paths + + Returns: + Dictionary containing processed data for each view + """ + processed_data = {} + + # Read generated depth visualization files + depth_vis_dir = os.path.join(target_dir, "depth_vis") + + if os.path.exists(depth_vis_dir): + depth_files = sorted(glob.glob(os.path.join(depth_vis_dir, "*.jpg"))) + for i, depth_file in enumerate(depth_files): + # Use processed images directly from API + processed_image = None + if prediction.processed_images is not None and i < len( + prediction.processed_images + ): + processed_image = prediction.processed_images[i] + + processed_data[i] = { + "depth_image": depth_file, + "image": processed_image, + "original_image_path": image_paths[i] if i < len(image_paths) else None, + "depth": prediction.depth[i] if i < len(prediction.depth) else None, + "intrinsics": ( + prediction.intrinsics[i] + if prediction.intrinsics is not None and i < len(prediction.intrinsics) + else None + ), + "mask": None, # No mask information available + } + + return processed_data + + # cleanup() removed: call cleanup_cuda_memory() directly where needed. diff --git a/depth_anything_3/app/modules/ui_components.py b/depth_anything_3/app/modules/ui_components.py new file mode 100644 index 0000000000000000000000000000000000000000..3a3762fff3f14284364dfe0b87a5eaf46ec8669d --- /dev/null +++ b/depth_anything_3/app/modules/ui_components.py @@ -0,0 +1,477 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +UI components module for Depth Anything 3 Gradio app. + +This module contains UI component definitions and layout functions. +""" + +import os +from typing import Any, Dict, List, Tuple +import gradio as gr + +from depth_anything_3.app.modules.utils import get_logo_base64, get_scene_info + + +class UIComponents: + """ + Handles UI component creation and layout for the Gradio app. + """ + + def __init__(self): + """Initialize the UI components handler.""" + + def create_upload_section(self) -> Tuple[gr.Video, gr.Slider, gr.File, gr.Gallery]: + """ + Create the upload section with video, images, and gallery components. + + Returns: + A tuple of Gradio components: (input_video, s_time_interval, input_images, image_gallery). + """ + input_video = gr.Video(label="Upload Video", interactive=True) + s_time_interval = gr.Slider( + minimum=0.1, + maximum=60, + value=10, + step=0.1, + label="Sampling FPS (Frames Per Second)", + interactive=True, + visible=True, + ) + input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True) + image_gallery = gr.Gallery( + label="Preview", + columns=4, + height="300px", + show_download_button=True, + object_fit="contain", + preview=True, + interactive=False, + ) + + return input_video, s_time_interval, input_images, image_gallery + + def create_3d_viewer_section(self) -> gr.Model3D: + """ + Create the 3D viewer component. + + Returns: + 3D model viewer component + """ + return gr.Model3D( + height=520, + zoom_speed=0.5, + pan_speed=0.5, + clear_color=[0.0, 0.0, 0.0, 0.0], + key="persistent_3d_viewer", + elem_id="reconstruction_3d_viewer", + ) + + def create_nvs_video(self) -> Tuple[gr.Video, gr.Markdown]: + """ + Create the 3DGS rendered video display component and info message. + + Returns: + Tuple of (video component, info message component) + """ + with gr.Column(): + gs_info = gr.Markdown( + ( + "‼️ **3D Gaussian Splatting rendering is currently DISABLED.**


" + "To render novel views from 3DGS, " + "enable **Infer 3D Gaussian Splatting** below.
" + "Next, in **Visualization Options**, " + "*optionally* configure the **rendering trajectory** (default: smooth) " + "and **video quality** (default: low), " + "then click **Reconstruct**." + ), + visible=True, + height=520, + ) + gs_video = gr.Video( + height=520, + label="3DGS Rendered NVS Video (depth shown for reference only)", + interactive=False, + visible=False, + ) + return gs_video, gs_info + + def create_depth_section(self) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image]: + """ + Create the depth visualization section. + + Returns: + A tuple of (prev_depth_btn, depth_view_selector, next_depth_btn, depth_map) + """ + with gr.Row(elem_classes=["navigation-row"]): + prev_depth_btn = gr.Button("◀ Previous", size="sm", scale=1) + depth_view_selector = gr.Dropdown( + choices=["View 1"], + value="View 1", + label="Select View", + scale=2, + interactive=True, + allow_custom_value=True, + ) + next_depth_btn = gr.Button("Next ▶", size="sm", scale=1) + depth_map = gr.Image( + type="numpy", + label="Colorized Depth Map", + format="png", + interactive=False, + ) + + return prev_depth_btn, depth_view_selector, next_depth_btn, depth_map + + def create_measure_section( + self, + ) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image, gr.Image, gr.Markdown]: + """ + Create the measurement section. + + Returns: + A tuple of (prev_measure_btn, measure_view_selector, next_measure_btn, measure_image, + measure_depth_image, measure_text) + """ + from depth_anything_3.app.css_and_html import MEASURE_INSTRUCTIONS_HTML + + gr.Markdown(MEASURE_INSTRUCTIONS_HTML) + with gr.Row(elem_classes=["navigation-row"]): + prev_measure_btn = gr.Button("◀ Previous", size="sm", scale=1) + measure_view_selector = gr.Dropdown( + choices=["View 1"], + value="View 1", + label="Select View", + scale=2, + interactive=True, + allow_custom_value=True, + ) + next_measure_btn = gr.Button("Next ▶", size="sm", scale=1) + with gr.Row(): + measure_image = gr.Image( + type="numpy", + show_label=False, + format="webp", + interactive=False, + sources=[], + label="RGB Image", + scale=1, + height=275, + ) + measure_depth_image = gr.Image( + type="numpy", + show_label=False, + format="webp", + interactive=False, + sources=[], + label="Depth Visualization (Right Half)", + scale=1, + height=275, + ) + gr.Markdown( + "**Note:** Images have been adjusted to model processing size. " + "Click two points on the RGB image to measure distance." + ) + measure_text = gr.Markdown("") + + return ( + prev_measure_btn, + measure_view_selector, + next_measure_btn, + measure_image, + measure_depth_image, + measure_text, + ) + + def create_inference_control_section(self) -> Tuple[gr.Dropdown, gr.Checkbox, gr.Dropdown]: + """ + Create the inference control section (before inference). + + Returns: + Tuple of (process_res_method_dropdown, infer_gs, ref_view_strategy) + """ + with gr.Row(): + process_res_method_dropdown = gr.Dropdown( + choices=["high_res", "low_res"], + value="low_res", + label="Image Processing Method", + info="low_res for much more images", + scale=1, + ) + # Modify line 220, add color class + infer_gs = gr.Checkbox( + label="Infer 3D Gaussian Splatting", + value=False, + info=( + 'Enable novel view rendering from 3DGS ( requires extra processing time)' + ), + scale=1, + ) + ref_view_strategy = gr.Dropdown( + choices=["saddle_balanced", "saddle_sim_range", "first", "middle"], + value="saddle_balanced", + label="Reference View Strategy", + info="Strategy for selecting reference view from multiple inputs", + scale=1, + ) + + return (process_res_method_dropdown, infer_gs, ref_view_strategy) + + def create_display_control_section( + self, + ) -> Tuple[ + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Slider, + gr.Slider, + gr.Dropdown, + gr.Dropdown, + gr.Button, + gr.ClearButton, + ]: + """ + Create the display control section (options for visualization). + + Returns: + Tuple of display control components including buttons + """ + with gr.Column(): + # 3DGS options at the top + with gr.Row(): + gs_trj_mode = gr.Dropdown( + choices=["smooth", "extend"], + value="smooth", + label=("Rendering trajectory for 3DGS viewpoints (requires n_views ≥ 2)"), + info=("'smooth' for view interpolation; 'extend' for longer trajectory"), + visible=False, # initially hidden + ) + gs_video_quality = gr.Dropdown( + choices=["low", "medium", "high"], + value="low", + label=("Video quality for 3DGS rendered outputs"), + info=("'low' for faster loading speed; 'high' for better visual quality"), + visible=False, # initially hidden + ) + + # Reconstruct and Clear buttons (before Visualization Options) + with gr.Row(): + submit_btn = gr.Button("Reconstruct", scale=1, variant="primary") + clear_btn = gr.ClearButton(scale=1) + + gr.Markdown("### Visualization Options: (Click Reconstruct to update)") + show_cam = gr.Checkbox(label="Show Camera", value=True) + filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False) + filter_white_bg = gr.Checkbox(label="Filter White Background", value=False) + save_percentage = gr.Slider( + minimum=0, + maximum=100, + value=10, + step=1, + label="Filter Percentage", + info="Confidence Threshold (%): Higher values filter more points.", + ) + num_max_points = gr.Slider( + minimum=1000, + maximum=100000, + value=1000, + step=1000, + label="Max Points (K points)", + info="Maximum number of points to export to GLB (in thousands)", + ) + + return ( + show_cam, + filter_black_bg, + filter_white_bg, + save_percentage, + num_max_points, + gs_trj_mode, + gs_video_quality, + submit_btn, + clear_btn, + ) + + def create_control_section( + self, + ) -> Tuple[ + gr.Button, + gr.ClearButton, + gr.Dropdown, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Dropdown, + gr.Checkbox, + gr.Textbox, + ]: + """ + Create the control section with buttons and options. + + Returns: + Tuple of control components + """ + with gr.Row(): + submit_btn = gr.Button("Reconstruct", scale=1, variant="primary") + clear_btn = gr.ClearButton( + scale=1, + ) + + with gr.Row(): + frame_filter = gr.Dropdown( + choices=["All"], value="All", label="Show Points from Frame" + ) + with gr.Column(): + gr.Markdown("### Visualization Option: (Click Reconstruct to update)") + show_cam = gr.Checkbox(label="Show Camera", value=True) + show_mesh = gr.Checkbox(label="Show Mesh", value=True) + filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False) + filter_white_bg = gr.Checkbox(label="Filter White Background", value=False) + gr.Markdown("### Reconstruction Options: (updated on next run)") + apply_mask_checkbox = gr.Checkbox( + label="Apply mask for predicted ambiguous depth classes & edges", + value=True, + ) + process_res_method_dropdown = gr.Dropdown( + choices=[ + "upper_bound_resize", + "upper_bound_crop", + "lower_bound_resize", + "lower_bound_crop", + ], + value="upper_bound_resize", + label="Image Processing Method", + info="Method for resizing input images", + ) + save_to_gallery_checkbox = gr.Checkbox( + label="Save to Gallery", + value=False, + info="Save current reconstruction results to gallery directory", + ) + gallery_name_input = gr.Textbox( + label="Gallery Name", + placeholder="Enter a name for the gallery folder", + value="", + info="Leave empty for auto-generated name with timestamp", + ) + + return ( + submit_btn, + clear_btn, + frame_filter, + show_cam, + show_mesh, + filter_black_bg, + filter_white_bg, + apply_mask_checkbox, + process_res_method_dropdown, + save_to_gallery_checkbox, + gallery_name_input, + ) + + def create_example_scenes_section(self) -> List[Dict[str, Any]]: + """ + Create the example scenes section. + + Returns: + List of scene information dictionaries + """ + # Get workspace directory from environment variable + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + examples_dir = os.path.join(workspace_dir, "examples") + + # Get scene information + scenes = get_scene_info(examples_dir) + + return scenes + + def create_example_scene_grid(self, scenes: List[Dict[str, Any]]) -> List[gr.Image]: + """ + Create the example scene grid. + + Args: + scenes: List of scene information dictionaries + + Returns: + List of scene image components + """ + scene_components = [] + + if scenes: + for i in range(0, len(scenes), 4): # Process 4 scenes per row + with gr.Row(): + for j in range(4): + scene_idx = i + j + if scene_idx < len(scenes): + scene = scenes[scene_idx] + with gr.Column(scale=1, elem_classes=["clickable-thumbnail"]): + # Clickable thumbnail + scene_img = gr.Image( + value=scene["thumbnail"], + height=150, + interactive=False, + show_label=False, + elem_id=f"scene_thumb_{scene['name']}", + sources=[], + ) + scene_components.append(scene_img) + + # Scene name and image count as text below thumbnail + gr.Markdown( + f"**{scene['name']}** \n {scene['num_images']} images", + elem_classes=["scene-info"], + ) + else: + # Empty column to maintain grid structure + with gr.Column(scale=1): + pass + + return scene_components + + def create_header_section(self) -> gr.HTML: + """ + Create the header section with logo and title. + + Returns: + Header HTML component + """ + from depth_anything_3.app.css_and_html import get_header_html + + return gr.HTML(get_header_html(get_logo_base64())) + + def create_description_section(self) -> gr.HTML: + """ + Create the description section. + + Returns: + Description HTML component + """ + from depth_anything_3.app.css_and_html import get_description_html + + return gr.HTML(get_description_html()) + + def create_acknowledgements_section(self) -> gr.HTML: + """ + Create the acknowledgements section. + + Returns: + Acknowledgements HTML component + """ + from depth_anything_3.app.css_and_html import get_acknowledgements_html + + return gr.HTML(get_acknowledgements_html()) diff --git a/depth_anything_3/app/modules/utils.py b/depth_anything_3/app/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..985abc145753c77143cd87d762ded76ac1dd0755 --- /dev/null +++ b/depth_anything_3/app/modules/utils.py @@ -0,0 +1,207 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utility functions for Depth Anything 3 Gradio app. + +This module contains helper functions for data processing, visualization, +and file operations. +""" + + +import json +import os +import shutil +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple +import numpy as np + +def create_depth_visualization(depth: np.ndarray) -> Optional[np.ndarray]: + """ + Create a colored depth visualization. + + Args: + depth: Depth array + + Returns: + Colored depth visualization or None + """ + if depth is None: + return None + + # Normalize depth to 0-1 range + depth_min = depth[depth > 0].min() if (depth > 0).any() else 0 + depth_max = depth.max() + + if depth_max <= depth_min: + return None + + # Normalize depth + depth_norm = (depth - depth_min) / (depth_max - depth_min) + depth_norm = np.clip(depth_norm, 0, 1) + + # Apply colormap (using matplotlib's viridis colormap) + import matplotlib.cm as cm + + # Convert to colored image + depth_colored = cm.viridis(depth_norm)[:, :, :3] # Remove alpha channel + depth_colored = (depth_colored * 255).astype(np.uint8) + + return depth_colored + + +def save_to_gallery_func( + target_dir: str, processed_data: Dict[int, Dict[str, Any]], gallery_name: Optional[str] = None +) -> Tuple[bool, str]: + """ + Save the current reconstruction results to the gallery directory. + + Args: + target_dir: Source directory containing reconstruction results + processed_data: Processed data dictionary + gallery_name: Name for the gallery folder + + Returns: + Tuple of (success, message) + """ + try: + # Get gallery directory from environment variable or use default + gallery_dir = os.environ.get( + "DA3_GALLERY_DIR", + "workspace/gallery", + ) + if not os.path.exists(gallery_dir): + os.makedirs(gallery_dir) + + # Use provided name or create a unique name + if gallery_name is None or gallery_name.strip() == "": + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + gallery_name = f"reconstruction_{timestamp}" + + gallery_path = os.path.join(gallery_dir, gallery_name) + + # Check if directory already exists + if os.path.exists(gallery_path): + return False, f"Save failed: folder '{gallery_name}' already exists" + + # Create the gallery directory + os.makedirs(gallery_path, exist_ok=True) + + # Copy GLB file + glb_source = os.path.join(target_dir, "scene.glb") + glb_dest = os.path.join(gallery_path, "scene.glb") + if os.path.exists(glb_source): + shutil.copy2(glb_source, glb_dest) + + # Copy depth visualization images + depth_vis_dir = os.path.join(target_dir, "depth_vis") + if os.path.exists(depth_vis_dir): + gallery_depth_vis = os.path.join(gallery_path, "depth_vis") + shutil.copytree(depth_vis_dir, gallery_depth_vis) + + # Copy original images + images_source = os.path.join(target_dir, "images") + if os.path.exists(images_source): + gallery_images = os.path.join(gallery_path, "images") + shutil.copytree(images_source, gallery_images) + + scene_preview_source = os.path.join(target_dir, "scene.jpg") + scene_preview_dest = os.path.join(gallery_path, "scene.jpg") + shutil.copy2(scene_preview_source, scene_preview_dest) + + # Save metadata + metadata = { + "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"), + "num_images": len(processed_data) if processed_data else 0, + "gallery_name": gallery_name, + } + + with open(os.path.join(gallery_path, "metadata.json"), "w") as f: + json.dump(metadata, f, indent=2) + + print(f"Saved reconstruction to gallery: {gallery_path}") + return True, f"Save successful: saved to {gallery_path}" + + except Exception as e: + print(f"Error saving to gallery: {e}") + return False, f"Save failed: {str(e)}" + + +def get_scene_info(examples_dir: str) -> List[Dict[str, Any]]: + """ + Get information about scenes in the examples directory. + + Args: + examples_dir: Path to examples directory + + Returns: + List of scene information dictionaries + """ + import glob + + scenes = [] + if not os.path.exists(examples_dir): + return scenes + + for scene_folder in sorted(os.listdir(examples_dir)): + scene_path = os.path.join(examples_dir, scene_folder) + if os.path.isdir(scene_path): + # Find all image files in the scene folder + image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tiff", "*.tif"] + image_files = [] + for ext in image_extensions: + image_files.extend(glob.glob(os.path.join(scene_path, ext))) + image_files.extend(glob.glob(os.path.join(scene_path, ext.upper()))) + + if image_files: + # Sort images and get the first one for thumbnail + image_files = sorted(image_files) + first_image = image_files[0] + num_images = len(image_files) + + scenes.append( + { + "name": scene_folder, + "path": scene_path, + "thumbnail": first_image, + "num_images": num_images, + "image_files": image_files, + } + ) + + return scenes + + +# NOTE: cleanup was moved to a single canonical helper in +# `depth_anything_3.utils.memory.cleanup_cuda_memory`. +# Callers should import and call that directly instead of using this module. + + +def get_logo_base64() -> Optional[str]: + """ + Convert WAI logo to base64 for embedding in HTML. + + Returns: + Base64 encoded logo string or None + """ + import base64 + + logo_path = "examples/WAI-Logo/wai_logo.png" + try: + with open(logo_path, "rb") as img_file: + img_data = img_file.read() + base64_str = base64.b64encode(img_data).decode() + return f"data:image/png;base64,{base64_str}" + except FileNotFoundError: + return None diff --git a/depth_anything_3/app/modules/visualization.py b/depth_anything_3/app/modules/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..ada49b25532a6ddbeb0e7f99856498e4ce3fb2ad --- /dev/null +++ b/depth_anything_3/app/modules/visualization.py @@ -0,0 +1,434 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Visualization module for Depth Anything 3 Gradio app. + +This module handles visualization updates, navigation, and measurement functionality. +""" + +import os +from typing import Any, Dict, List, Optional, Tuple +import cv2 +import gradio as gr +import numpy as np + + +class VisualizationHandler: + """ + Handles visualization updates and navigation for the Gradio app. + """ + + def __init__(self): + """Initialize the visualization handler.""" + + def update_view_selectors( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[gr.Dropdown, gr.Dropdown]: + """ + Update view selector dropdowns based on available views. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (depth_view_selector, measure_view_selector) + """ + if processed_data is None or len(processed_data) == 0: + choices = ["View 1"] + else: + num_views = len(processed_data) + choices = [f"View {i + 1}" for i in range(num_views)] + + return ( + gr.Dropdown(choices=choices, value=choices[0]), # depth_view_selector + gr.Dropdown(choices=choices, value=choices[0]), # measure_view_selector + ) + + def get_view_data_by_index( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[Dict[str, Any]]: + """ + Get view data by index, handling bounds. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to get + + Returns: + View data dictionary or None + """ + if processed_data is None or len(processed_data) == 0: + return None + + view_keys = list(processed_data.keys()) + if view_index < 0 or view_index >= len(view_keys): + view_index = 0 + + return processed_data[view_keys[view_index]] + + def update_depth_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[str]: + """ + Update depth view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Path to depth visualization image or None + """ + view_data = self.get_view_data_by_index(processed_data, view_index) + if view_data is None or view_data.get("depth_image") is None: + return None + + # Return the depth visualization image directly + return view_data["depth_image"] + + def navigate_depth_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector_value: str, + direction: int, + ) -> Tuple[str, Optional[str]]: + """ + Navigate depth view (direction: -1 for previous, +1 for next). + + Args: + processed_data: Processed data dictionary + current_selector_value: Current selector value + direction: Direction to navigate (-1 for previous, +1 for next) + + Returns: + Tuple of (new_selector_value, depth_vis) + """ + if processed_data is None or len(processed_data) == 0: + return "View 1", None + + # Parse current view number + try: + current_view = int(current_selector_value.split()[1]) - 1 + except: # noqa + current_view = 0 + + num_views = len(processed_data) + new_view = (current_view + direction) % num_views + + new_selector_value = f"View {new_view + 1}" + depth_vis = self.update_depth_view(processed_data, new_view) + + return new_selector_value, depth_vis + + def update_measure_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Update measure view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Tuple of (measure_image, depth_right_half, measure_points) + """ + view_data = self.get_view_data_by_index(processed_data, view_index) + if view_data is None: + return None, None, [] # image, depth_right_half, measure_points + + # Get the processed (resized) image + if "image" in view_data and view_data["image"] is not None: + image = view_data["image"].copy() + else: + return None, None, [] + + # Ensure image is in uint8 format + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255).astype(np.uint8) + else: + image = image.astype(np.uint8) + + # Extract right half of the depth visualization (pure depth part) + depth_image_path = view_data.get("depth_image", None) + depth_right_half = None + + if depth_image_path and os.path.exists(depth_image_path): + try: + # Load the combined depth visualization image + depth_combined = cv2.imread(depth_image_path) + depth_combined = cv2.cvtColor(depth_combined, cv2.COLOR_BGR2RGB) + if depth_combined is not None: + height, width = depth_combined.shape[:2] + # Extract right half (depth visualization part) + depth_right_half = depth_combined[:, width // 2 :] + except Exception as e: + print(f"Error extracting depth right half: {e}") + + return image, depth_right_half, [] + + def navigate_measure_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector_value: str, + direction: int, + ) -> Tuple[str, Optional[np.ndarray], Optional[str], List]: + """ + Navigate measure view (direction: -1 for previous, +1 for next). + + Args: + processed_data: Processed data dictionary + current_selector_value: Current selector value + direction: Direction to navigate (-1 for previous, +1 for next) + + Returns: + Tuple of (new_selector_value, measure_image, depth_image_path, measure_points) + """ + if processed_data is None or len(processed_data) == 0: + return "View 1", None, None, [] + + # Parse current view number + try: + current_view = int(current_selector_value.split()[1]) - 1 + except: # noqa + current_view = 0 + + num_views = len(processed_data) + new_view = (current_view + direction) % num_views + + new_selector_value = f"View {new_view + 1}" + measure_image, depth_right_half, measure_points = self.update_measure_view( + processed_data, new_view + ) + + return new_selector_value, measure_image, depth_right_half, measure_points + + def populate_visualization_tabs( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[Optional[str], Optional[np.ndarray], Optional[str], List]: + """ + Populate the depth and measure tabs with processed data. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (depth_vis, measure_img, depth_image_path, measure_points) + """ + if processed_data is None or len(processed_data) == 0: + return None, None, None, [] + + # Use update function to get depth visualization + depth_vis = self.update_depth_view(processed_data, 0) + measure_img, depth_right_half, _ = self.update_measure_view(processed_data, 0) + + return depth_vis, measure_img, depth_right_half, [] + + def reset_measure( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[Optional[np.ndarray], List, str]: + """ + Reset measure points. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (image, measure_points, text) + """ + if processed_data is None or len(processed_data) == 0: + return None, [], "" + + # Return the first view image + first_view = list(processed_data.values())[0] + return first_view["image"], [], "" + + def measure( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + measure_points: List, + current_view_selector: str, + event: gr.SelectData, + ) -> List: + """ + Handle measurement on images. + + Args: + processed_data: Processed data dictionary + measure_points: List of current measure points + current_view_selector: Current view selector value + event: Gradio select event + + Returns: + List of [image, depth_right_half, measure_points, text] + """ + try: + print(f"Measure function called with selector: {current_view_selector}") + + if processed_data is None or len(processed_data) == 0: + return [None, [], "No data available"] + + # Use the currently selected view instead of always using the first view + try: + current_view_index = int(current_view_selector.split()[1]) - 1 + except: # noqa + current_view_index = 0 + + print(f"Using view index: {current_view_index}") + + # Get view data safely + if current_view_index < 0 or current_view_index >= len(processed_data): + current_view_index = 0 + + view_keys = list(processed_data.keys()) + current_view = processed_data[view_keys[current_view_index]] + + if current_view is None: + return [None, [], "No view data available"] + + point2d = event.index[0], event.index[1] + print(f"Clicked point: {point2d}") + + measure_points.append(point2d) + + # Get image and depth visualization + image, depth_right_half, _ = self.update_measure_view( + processed_data, current_view_index + ) + if image is None: + return [None, [], "No image available"] + + image = image.copy() + + # Ensure image is in uint8 format for proper cv2 operations + try: + if image.dtype != np.uint8: + if image.max() <= 1.0: + # Image is in [0, 1] range, convert to [0, 255] + image = (image * 255).astype(np.uint8) + else: + # Image is already in [0, 255] range + image = image.astype(np.uint8) + except Exception as e: + print(f"Image conversion error: {e}") + return [None, [], f"Image conversion error: {e}"] + + # Draw circles for points + try: + for p in measure_points: + if 0 <= p[0] < image.shape[1] and 0 <= p[1] < image.shape[0]: + image = cv2.circle(image, p, radius=5, color=(255, 0, 0), thickness=2) + except Exception as e: + print(f"Drawing error: {e}") + return [None, [], f"Drawing error: {e}"] + + # Get depth information from processed_data + depth_text = "" + try: + for i, p in enumerate(measure_points): + if ( + current_view["depth"] is not None + and 0 <= p[1] < current_view["depth"].shape[0] + and 0 <= p[0] < current_view["depth"].shape[1] + ): + d = current_view["depth"][p[1], p[0]] + depth_text += f"- **P{i + 1} depth: {d:.2f}m**\n" + else: + depth_text += f"- **P{i + 1}: Click position ({p[0]}, {p[1]}) - No depth information**\n" # noqa: E501 + except Exception as e: + print(f"Depth text error: {e}") + depth_text = f"Error computing depth: {e}\n" + + if len(measure_points) == 2: + try: + point1, point2 = measure_points + # Draw line + if ( + 0 <= point1[0] < image.shape[1] + and 0 <= point1[1] < image.shape[0] + and 0 <= point2[0] < image.shape[1] + and 0 <= point2[1] < image.shape[0] + ): + image = cv2.line(image, point1, point2, color=(255, 0, 0), thickness=2) + + # Compute 3D distance using depth information and camera intrinsics + distance_text = "- **Distance: Unable to calculate 3D distance**" + if ( + current_view["depth"] is not None + and 0 <= point1[1] < current_view["depth"].shape[0] + and 0 <= point1[0] < current_view["depth"].shape[1] + and 0 <= point2[1] < current_view["depth"].shape[0] + and 0 <= point2[0] < current_view["depth"].shape[1] + ): + try: + # Get depth values at the two points + d1 = current_view["depth"][point1[1], point1[0]] + d2 = current_view["depth"][point2[1], point2[0]] + + # Convert 2D pixel coordinates to 3D world coordinates + if current_view["intrinsics"] is not None: + # Get camera intrinsics + K = current_view["intrinsics"] # 3x3 intrinsic matrix + fx, fy = K[0, 0], K[1, 1] # focal lengths + cx, cy = K[0, 2], K[1, 2] # principal point + + # Convert pixel coordinates to normalized camera coordinates + # Point 1: (u1, v1) -> (x1, y1, z1) + u1, v1 = point1[0], point1[1] + x1 = (u1 - cx) * d1 / fx + y1 = (v1 - cy) * d1 / fy + z1 = d1 + + # Point 2: (u2, v2) -> (x2, y2, z2) + u2, v2 = point2[0], point2[1] + x2 = (u2 - cx) * d2 / fx + y2 = (v2 - cy) * d2 / fy + z2 = d2 + + # Calculate 3D Euclidean distance + p1_3d = np.array([x1, y1, z1]) + p2_3d = np.array([x2, y2, z2]) + distance_3d = np.linalg.norm(p1_3d - p2_3d) + + distance_text = f"- **Distance: {distance_3d:.2f}m**" + else: + # Fallback to simplified calculation if no intrinsics + pixel_distance = np.sqrt( + (point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2 + ) + avg_depth = (d1 + d2) / 2 + scale_factor = avg_depth / 1000 # Rough scaling factor + estimated_3d_distance = pixel_distance * scale_factor + distance_text = f"- **Distance: {estimated_3d_distance:.2f}m (estimated, no intrinsics)**" # noqa: E501 + + except Exception as e: + print(f"Distance computation error: {e}") + distance_text = f"- **Distance computation error: {e}**" + + measure_points = [] + text = depth_text + distance_text + print(f"Measurement complete: {text}") + return [image, depth_right_half, measure_points, text] + except Exception as e: + print(f"Final measurement error: {e}") + return [None, [], f"Measurement error: {e}"] + else: + print(f"Single point measurement: {depth_text}") + return [image, depth_right_half, measure_points, depth_text] + + except Exception as e: + print(f"Overall measure function error: {e}") + return [None, [], f"Measure function error: {e}"] diff --git a/depth_anything_3/bench/__init__.py b/depth_anything_3/bench/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..664ee41e3fde08213d64dc81e543af3b0b0783ad --- /dev/null +++ b/depth_anything_3/bench/__init__.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Depth Anything 3 Benchmark Evaluation Module. + +This module provides tools for evaluating DepthAnything3 model on various benchmark datasets. +Currently supported datasets: +- DTU (3D Reconstruction) +- DTU-64 (Pose Evaluation Only) +- ETH3D (3D Reconstruction) +- 7Scenes (3D Reconstruction) +- ScanNet++ (3D Reconstruction) +- HiRoom (3D Reconstruction) + +Supported evaluation modes: +- pose: Camera pose estimation evaluation +- recon_unposed: 3D reconstruction with predicted poses +- recon_posed: 3D reconstruction with ground truth poses +""" + +from depth_anything_3.bench.registries import MV_REGISTRY, MONO_REGISTRY + + +def __getattr__(name): + """Lazy import to avoid circular import when running as __main__.""" + if name == "Evaluator": + from depth_anything_3.bench.evaluator import Evaluator + return Evaluator + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = ["Evaluator", "MV_REGISTRY", "MONO_REGISTRY"] + diff --git a/depth_anything_3/bench/configs/eval_bench.yaml b/depth_anything_3/bench/configs/eval_bench.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1d9924ddfeecb42a5dda47d19f5e86d2bedd8273 --- /dev/null +++ b/depth_anything_3/bench/configs/eval_bench.yaml @@ -0,0 +1,98 @@ +# DepthAnything3 Benchmark Evaluation Configuration +# +# This config can be loaded and overridden via command line. +# Example: python -m depth_anything_3.bench.evaluator --model /path/to/model --work_dir /path/to/workspace +# +# See depth_anything_3.cfg for config utility functions. + +# ============================================================================== +# Model Configuration +# ============================================================================== +model: + # Path to model checkpoint or HuggingFace model ID + path: depth-anything/DA3-GIANT + +# ============================================================================== +# Workspace Configuration +# ============================================================================== +workspace: + # Working directory for outputs (model results, metrics, etc.) + work_dir: ./workspace/evaluation + +# ============================================================================== +# Evaluation Configuration +# ============================================================================== +eval: + # Datasets to evaluate + # Options: dtu, dtu64, eth3d, 7scenes (sevenscenes), scannetpp, hiroom + datasets: + - eth3d + - 7scenes + - scannetpp + - hiroom + - dtu + - dtu64 + + # Evaluation modes + # Options: pose, recon_unposed, recon_posed, view_syn + modes: + - pose + - recon_unposed + - recon_posed + + # Reference view selection strategy for inference + # Options: first, saddle_balanced, auto, mid + ref_view_strategy: "first" + + # Specific scenes to evaluate (null = all scenes) + # Example: [courtyard, relief] for eth3d + scenes: null + + # Maximum number of frames per scene (for sampling) + # If a scene has more frames, randomly sample to this limit. + # Set to -1 to disable sampling. + max_frames: 100 + + # Only run evaluation (skip inference) + eval_only: false + + # Only print saved metrics (skip inference and evaluation) + print_only: false + +# ============================================================================== +# Inference Configuration +# ============================================================================== +inference: + # Number of parallel workers for TSDF fusion + num_fusion_workers: 4 + + # Enable debug mode with verbose output + debug: false + +# ============================================================================== +# Preset Configurations +# ============================================================================== +# These can be activated via command line: --preset full_eval + +presets: + # Full evaluation on all 6 datasets + full_eval: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu, dtu64] + modes: [pose, recon_unposed, recon_posed] + + # Pose-only evaluation + pose_only: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu64] + modes: [pose] + + # Reconstruction-only evaluation (5 datasets, excluding dtu64) + recon_only: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu] + modes: [recon_unposed, recon_posed] + + # Quick test (single scene per dataset) + quick_test: + datasets: [eth3d] + modes: [pose, recon_unposed] + scenes: [courtyard] + diff --git a/depth_anything_3/bench/dataset.py b/depth_anything_3/bench/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2f30a0f96d7414f0ec086e65f37d3fb1ef8ec31e --- /dev/null +++ b/depth_anything_3/bench/dataset.py @@ -0,0 +1,136 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Base dataset class for benchmark evaluation. + +All dataset implementations should inherit from this class and implement +the required abstract methods. +""" + +import os +import time +from abc import abstractmethod +from typing import Dict as TDict + +import numpy as np +import torch +from addict import Dict + +from depth_anything_3.bench.utils import compute_pose +from depth_anything_3.utils.geometry import as_homogeneous + + +def _wait_for_file_ready(path: str, timeout: float = 3.0, interval: float = 0.2) -> None: + """Wait until file size stabilizes for 2 consecutive checks.""" + last_size = -1 + stable_count = 0 + start = time.time() + while time.time() - start < timeout: + time.sleep(interval) + size = os.path.getsize(path) + if size == last_size and size > 0: + stable_count += 1 + if stable_count >= 2: # Need 2 consecutive stable checks + return + else: + stable_count = 0 + last_size = size + + +class Dataset: + """ + Base class for all benchmark datasets. + + Subclasses must implement: + - SCENES: List of scene identifiers + - data_root: Path to dataset root + - get_data(scene): Return scene data (images, intrinsics, extrinsics, etc.) + - eval3d(scene, fuse_path): Evaluate 3D reconstruction + - fuse3d(scene, result_path, fuse_path, mode): Fuse depth maps into point cloud + + Optional overrides: + - eval_pose(scene, result_path): Evaluate pose estimation (default provided) + """ + + # Subclasses should define these + SCENES: list = [] + data_root: str = "" + + def __init__(self): + pass + + def eval_pose(self, scene: str, result_path: str) -> TDict[str, float]: + """ + Evaluate camera pose estimation accuracy. + + Args: + scene: Scene identifier + result_path: Path to .npz file containing predicted extrinsics + + Returns: + Dict with pose metrics (auc30, auc15, auc05, auc03) + """ + _wait_for_file_ready(result_path) + pred = np.load(result_path) + gt = self.get_data(scene) + return compute_pose( + torch.from_numpy(as_homogeneous(pred["extrinsics"])), + torch.from_numpy(as_homogeneous(gt["extrinsics"])), + ) + + @abstractmethod + def get_data(self, scene: str) -> Dict: + """ + Get scene data including images, camera parameters, and auxiliary info. + + Args: + scene: Scene identifier + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - camera extrinsics (world-to-camera) + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict - auxiliary data (masks, GT paths, etc.) + """ + raise NotImplementedError + + @abstractmethod + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate 3D reconstruction quality against ground truth. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with reconstruction metrics (e.g., acc, comp, overall) + """ + raise NotImplementedError + + @abstractmethod + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depth maps into a single point cloud. + + Args: + scene: Scene identifier + result_path: Path to .npz file with predicted depths and poses + fuse_path: Output path for fused point cloud (.ply) + mode: Fusion mode ("recon_unposed" or "recon_posed") + """ + raise NotImplementedError + diff --git a/depth_anything_3/bench/datasets/__init__.py b/depth_anything_3/bench/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2e9f464300c07432fc515fd1236a9883f1d8cd69 --- /dev/null +++ b/depth_anything_3/bench/datasets/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Benchmark dataset implementations. + +Datasets are auto-registered via decorators when imported. +Add new dataset files here and they will be automatically discovered. +""" + diff --git a/depth_anything_3/bench/datasets/dtu.py b/depth_anything_3/bench/datasets/dtu.py new file mode 100644 index 0000000000000000000000000000000000000000..4d07c91577b3e29786471331c55f4609adb6c112 --- /dev/null +++ b/depth_anything_3/bench/datasets/dtu.py @@ -0,0 +1,681 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +DTU Benchmark dataset implementation. + +DTU is a multi-view stereo benchmark for 3D reconstruction evaluation. +Reference: https://roboimagedata.compute.dtu.dk/ + +Note: DepthAnything3 was never trained on any images from DTU. +""" + +import glob +import os +from typing import Dict as TDict, List + +import numpy as np +import open3d as o3d +import torch +import torch.nn.functional as F +from addict import Dict +from PIL import Image +from plyfile import PlyData +from scipy.io import loadmat +from sklearn import neighbors as skln +from tqdm import tqdm + +from depth_anything_3.bench.dataset import Dataset +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.utils.constants import ( + DTU_DIST_THRESH, + DTU_EVAL_DATA_ROOT, + DTU_MAX_POINTS, + DTU_NUM_CONSIST, + DTU_SCENES, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="dtu") +@MONO_REGISTRY.register(name="dtu") +class DTU(Dataset): + """ + DTU Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (accuracy, completeness, overall) + - Point cloud fusion from depth maps + + The dataset uses MVSNet evaluation protocol: + https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view + """ + + data_root = DTU_EVAL_DATA_ROOT + SCENES = DTU_SCENES + + # Evaluation/triangulation hyperparameters from constants + dist_thresh = DTU_DIST_THRESH + num_consist = DTU_NUM_CONSIST + + # ------------------------------ + # Public API + # ------------------------------ + + def read_cam_file(self, filename: str) -> tuple: + """ + Read DTU camera file containing extrinsics and intrinsics. + + Args: + filename: Path to camera text file + + Returns: + Tuple of (intrinsics [3,3], extrinsics [4,4]) + """ + with open(filename) as f: + lines = [line.rstrip() for line in f.readlines()] + extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4)) + intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3)) + return intrinsics, extrinsics + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics, and GT masks. + + Args: + scene: Scene identifier (e.g., "scan1") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] + - intrinsics: np.ndarray [N, 3, 3] + - aux.mask_files: List[str] - paths to depth masks + """ + rgb_folder = os.path.join(self.data_root, "Rectified", scene) + camera_folder = os.path.join(self.data_root, "Cameras") + + files = sorted(glob.glob(os.path.join(rgb_folder, "*.png"))) + # Reorder: place index 33 first (reference view convention) + files = [files[33]] + files[:33] + files[34:] + + out = Dict( + { + "image_files": files, + "extrinsics": [], + "intrinsics": [], + "aux": Dict({"mask_files": []}), + } + ) + + for rgb_file in files: + basename = os.path.basename(rgb_file) + file_idx = basename.split("_")[1] + cam_idx = depth_idx = int(file_idx) - 1 + + mask_file = self._depth_mask_path(scene, depth_idx) + proj_mat_filename = os.path.join(camera_folder, f"{cam_idx:0>8}_cam.txt") + + ixt, ext = self.read_cam_file(proj_mat_filename) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.mask_files.append(mask_file) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + return out + + def get_3dgtpath(self, scene: str) -> str: + """Get path to ground truth point cloud for a scene.""" + scene_id = int(scene[4:]) + return os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply") + + def eval3d(self, scene: str, fuse_path: str, use_gpu: bool = False) -> TDict[str, float]: + """ + Evaluate fused point cloud against DTU GT with ObsMask/Plane. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud + use_gpu: If True, use GPU-accelerated distance computation (faster but may have minor numerical differences) + + Returns: + Dict with metrics: {"comp": float, "acc": float, "overall": float} + """ + scene_id = int(scene[4:]) + gt_ply = os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply") + mask_file = os.path.join( + self.data_root, f"SampleSet/mvs_data/ObsMask/ObsMask{scene_id}_10.mat" + ) + plane_file = os.path.join( + self.data_root, f"SampleSet/mvs_data/ObsMask/Plane{scene_id}.mat" + ) + result = self._evaluate_reconstruction( + scene, fuse_path, gt_ply, mask_file, plane_file, use_gpu=use_gpu + ) + return {"comp": result[0], "acc": result[1], "overall": result[2]} + + def load_masks(self, mask_files: List[str]) -> np.ndarray: + """ + Load DTU depth validity masks. + + Args: + mask_files: List of paths to mask images + + Returns: + Boolean array [N, H, W] indicating valid depth regions + """ + masks = [] + for mask_file in mask_files: + mask = Image.open(mask_file) + mask = np.array(mask, dtype=np.float32) + masks.append(mask > 10) + return np.asarray(masks) + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud and save to PLY. + + Args: + scene: Scene identifier (e.g., "scan114") + result_path: Path to npz file containing predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + gt_data = self.get_data(scene) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + masks = self.load_masks(gt_data.aux.mask_files) + + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed(pred_data, gt_data, masks) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed(pred_data, gt_data, masks) + else: + raise ValueError(f"Invalid mode: {mode}") + + proj_mat = self._build_proj_mats(intrinsics, extrinsics) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + depths_t = torch.from_numpy(depths).to(device=device, dtype=dtype).unsqueeze(1) + proj_t = torch.from_numpy(proj_mat).to(device=device, dtype=dtype) + height, width = depths_t.shape[-2:] + + points: List[np.ndarray] = [] + for idx in range(len(gt_data.image_files)): + if mode == "recon_unposed": + # Simple unfiltered back-projection per frame + cur_p_pcd = self._generate_points_from_depth( + depths_t[idx : idx + 1], proj_t[idx : idx + 1] + ) + mask = (depths_t[idx : idx + 1] > 0.001).squeeze() + cur_p_pcd = cur_p_pcd[:, :, mask] + no_filter_pc = cur_p_pcd.squeeze(0).permute(1, 0).cpu().numpy() + points.append(no_filter_pc) + else: # recon_posed + final_pc = self._fuse_consistent_points(depths_t, proj_t, idx, height, width) + points.append(final_pc) + + # Concatenate and optionally downsample to hard cap + points_np = np.concatenate(points, axis=0) + points_np = self._cap_points(points_np, max_points=DTU_MAX_POINTS) + + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points_np) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Geometry helpers + # ------------------------------ + + def _generate_points_from_depth( + self, depth: torch.Tensor, proj: torch.Tensor + ) -> torch.Tensor: + """ + Back-project depth map into 3D world coordinates. + + Args: + depth: Depth tensor [B, 1, H, W] + proj: Projection matrix [B, 4, 4] = [[K@R, K@t], [0,0,0,1]] + + Returns: + Point cloud tensor [B, 3, H, W] + """ + batch, height, width = depth.shape[0], depth.shape[2], depth.shape[3] + inv_proj = torch.inverse(proj) + rot = inv_proj[:, :3, :3] + trans = inv_proj[:, :3, 3:4] + + y, x = torch.meshgrid( + [ + torch.arange(0, height, dtype=torch.float32, device=depth.device), + torch.arange(0, width, dtype=torch.float32, device=depth.device), + ], + indexing="ij", + ) + y, x = y.contiguous(), x.contiguous() + y, x = y.view(height * width), x.view(height * width) + xyz = torch.stack((x, y, torch.ones_like(x))) + xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1) + rot_xyz = torch.matmul(rot, xyz) + rot_depth_xyz = rot_xyz * depth.view(batch, 1, -1) + proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1) + return proj_xyz.view(batch, 3, height, width) + + def _homo_warping( + self, + src_fea: torch.Tensor, + src_proj: torch.Tensor, + ref_proj: torch.Tensor, + depth_values: torch.Tensor, + ) -> torch.Tensor: + """ + Homography warping for multi-view consistency checking. + + Args: + src_fea: Source features [B, C, H, W] + src_proj: Source projection [B, 4, 4] + ref_proj: Reference projection [B, 4, 4] + depth_values: Depth values [B, Ndepth] or [B, Ndepth, H, W] + + Returns: + Warped features [B, C, H, W] + """ + batch, channels = src_fea.shape[0], src_fea.shape[1] + height, width = src_fea.shape[2], src_fea.shape[3] + + with torch.no_grad(): + proj = torch.matmul(src_proj, torch.inverse(ref_proj)) + rot = proj[:, :3, :3] + trans = proj[:, :3, 3:4] + + y, x = torch.meshgrid( + [ + torch.arange(0, height, dtype=torch.float32, device=src_fea.device), + torch.arange(0, width, dtype=torch.float32, device=src_fea.device), + ], + indexing="ij", + ) + y, x = y.contiguous(), x.contiguous() + y, x = y.view(height * width), x.view(height * width) + xyz = torch.stack((x, y, torch.ones_like(x))) + xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1) + rot_xyz = torch.matmul(rot, xyz) + + rot_depth_xyz = rot_xyz.unsqueeze(2) * depth_values.view(-1, 1, 1, height * width) + proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1) + proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :] + proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1 + proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1 + grid = torch.stack((proj_x_normalized, proj_y_normalized), dim=3) + + warped_src_fea = F.grid_sample( + src_fea, + grid.view(batch, height, width, 2), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + return warped_src_fea.view(batch, channels, height, width) + + def _filter_depth( + self, + ref_depth: torch.Tensor, + src_depths: torch.Tensor, + ref_proj: torch.Tensor, + src_projs: torch.Tensor, + ) -> tuple: + """ + Compute geometric consistency between reference and source depths. + + Args: + ref_depth: Reference depth [1, 1, H, W] + src_depths: Source depths [B, 1, H, W] + ref_proj: Reference projection [1, 4, 4] + src_projs: Source projections [B, 4, 4] + + Returns: + Tuple of (ref_pc, aligned_pcs, dist) + """ + ref_pc = self._generate_points_from_depth(ref_depth, ref_proj) + src_pcs = self._generate_points_from_depth(src_depths, src_projs) + aligned_pcs = self._homo_warping(src_pcs, src_projs, ref_proj, ref_depth) + x_2 = (ref_pc[:, 0] - aligned_pcs[:, 0]) ** 2 + y_2 = (ref_pc[:, 1] - aligned_pcs[:, 1]) ** 2 + z_2 = (ref_pc[:, 2] - aligned_pcs[:, 2]) ** 2 + dist = torch.sqrt(x_2 + y_2 + z_2).unsqueeze(1) + return ref_pc, aligned_pcs, dist + + def _extract_points( + self, pc: torch.Tensor, mask: torch.Tensor, rgb: np.ndarray = None + ) -> np.ndarray: + """Extract masked points from a dense grid.""" + pc = pc.cpu().numpy() + mask = mask.cpu().numpy().reshape(-1) + pc = pc.reshape(-1, 3) + points = pc[np.where(mask)] + if rgb is not None: + rgb = rgb.reshape(-1, 3) + colors = rgb[np.where(mask)] + return np.concatenate([points, colors], axis=1) + return points + + # ------------------------------ + # 3D Reconstruction Evaluation + # ------------------------------ + + def _evaluate_reconstruction( + self, + scanid: str, + pred_ply: str, + gt_ply: str, + mask_file: str, + plane_file: str, + down_dense: float = 0.2, + patch: int = 60, + max_dist: int = 20, + use_gpu: bool = False, + ) -> tuple: + """ + Compute accuracy, completeness, and overall metrics for one scan. + + Args: + scanid: Scan identifier + pred_ply: Predicted point cloud path or array + gt_ply: Ground truth point cloud path or array + mask_file: ObsMask file path + plane_file: Plane file path + down_dense: Downsample density (min distance between points) + patch: Patch size for boundary + max_dist: Outlier threshold in mm + use_gpu: If True, use GPU-accelerated distance computation + + Returns: + Tuple of (mean_d2s, mean_s2d, overall) + """ + thresh = down_dense + + # Load and downsample predicted point cloud + data_pcd = self._read_ply(pred_ply) if isinstance(pred_ply, str) else pred_ply + # Use fixed seed for reproducibility + shuffle_rng = np.random.default_rng(seed=42) + shuffle_rng.shuffle(data_pcd, axis=0) + + # Downsample point cloud + nn_engine = skln.NearestNeighbors( + n_neighbors=1, radius=thresh, algorithm="kd_tree", n_jobs=-1 + ) + nn_engine.fit(data_pcd) + rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False) + mask = np.ones(data_pcd.shape[0], dtype=np.bool_) + for curr, idxs in enumerate(rnn_idxs): + if mask[curr]: + mask[idxs] = 0 + mask[curr] = 1 + data_down = data_pcd[mask] + + # Restrict to observed volume (ObsMask) + obs_mask_file = loadmat(mask_file) + ObsMask, BB, Res = (obs_mask_file[attr] for attr in ["ObsMask", "BB", "Res"]) + BB = BB.astype(np.float32) + + inbound = ((data_down >= BB[:1] - patch) & (data_down < BB[1:] + patch * 2)).sum( + axis=-1 + ) == 3 + data_in = data_down[inbound] + + data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32) + grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum( + axis=-1 + ) == 3 + data_grid_in = data_grid[grid_inbound] + in_obs = ObsMask[data_grid_in[:, 0], data_grid_in[:, 1], data_grid_in[:, 2]].astype( + np.bool_ + ) + data_in_obs = data_in[grid_inbound][in_obs] + + # Compute accuracy (pred -> GT) and completeness (GT -> pred) + stl = self._read_ply(gt_ply) if isinstance(gt_ply, str) else gt_ply + + if use_gpu and torch.cuda.is_available(): + # GPU-accelerated distance computation + mean_d2s = self._knn_dist_gpu(data_in_obs, stl, max_dist) + else: + # CPU version (original, for exact reproduction) + nn_engine.fit(stl) + dist_d2s, _ = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True) + mean_d2s = dist_d2s[dist_d2s < max_dist].mean() + + ground_plane = loadmat(plane_file)["P"] + stl_hom = np.concatenate([stl, np.ones_like(stl[:, :1])], -1) + above = (ground_plane.reshape((1, 4)) * stl_hom).sum(-1) > 0 + stl_above = stl[above] + + if use_gpu and torch.cuda.is_available(): + # GPU-accelerated distance computation + mean_s2d = self._knn_dist_gpu(stl_above, data_in, max_dist) + else: + # CPU version (original, for exact reproduction) + nn_engine.fit(data_in) + dist_s2d, _ = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True) + mean_s2d = dist_s2d[dist_s2d < max_dist].mean() + + overall = (mean_d2s + mean_s2d) / 2 + return mean_d2s, mean_s2d, overall + + def _knn_dist_gpu( + self, + query: np.ndarray, + target: np.ndarray, + max_dist: float, + batch_size: int = 8192, + target_batch_size: int = 50000, + ) -> float: + """ + GPU-accelerated nearest neighbor distance computation. + + Args: + query: Query points [N, 3] + target: Target points [M, 3] + max_dist: Outlier threshold + batch_size: Batch size for query to avoid OOM (tuned for 16GB GPU) + target_batch_size: Batch size for target to avoid OOM + + Returns: + Mean distance (excluding outliers) + """ + device = torch.device("cuda") + + all_min_dists = [] + n_query_batches = (len(query) + batch_size - 1) // batch_size + n_target_batches = (len(target) + target_batch_size - 1) // target_batch_size + + # Pre-load target batches to GPU to avoid repeated transfers + # Memory: ~50000 pts * 3 coords * 4 bytes * n_batches + target_batches = [] + for j in range(0, len(target), target_batch_size): + target_batch = target[j : j + target_batch_size] + target_t = torch.from_numpy(target_batch).float().to(device) + target_batches.append(target_t) + + with tqdm(total=n_query_batches, desc=" GPU KNN", leave=False, ncols=100) as pbar: + for i in range(0, len(query), batch_size): + batch = query[i : i + batch_size] + query_t = torch.from_numpy(batch).float().to(device) + + # Compute distances to all target batches + # Memory peak: query_batch × target_batch_size × 4 bytes + # = 8192 × 50000 × 4 = ~1.6 GB per cdist call + batch_min_dists = [] + for target_t in target_batches: + dists = torch.cdist(query_t, target_t) + batch_min_dists.append(dists.min(dim=1).values) + del dists # Free immediately + + # Get minimum distance across all target batches + min_dists = torch.stack(batch_min_dists, dim=1).min(dim=1).values + all_min_dists.append(min_dists.cpu().numpy()) + + del query_t, min_dists, batch_min_dists + pbar.update(1) + + # Clean up target batches + for target_t in target_batches: + del target_t + torch.cuda.empty_cache() + + all_min_dists = np.concatenate(all_min_dists) + return all_min_dists[all_min_dists < max_dist].mean() + + def _read_ply(self, file: str) -> np.ndarray: + """Read point cloud from PLY file.""" + data = PlyData.read(file) + vertex = data["vertex"] + return np.stack([vertex["x"], vertex["y"], vertex["z"]], axis=-1) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _depth_mask_path(self, scene: str, depth_idx: int) -> str: + """Get path to depth mask for a scene and frame.""" + return os.path.join( + self.data_root, "depth_raw", "Depths", scene, f"depth_visual_{depth_idx:04d}.png" + ) + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, masks: np.ndarray + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Applies Umeyama scale, rescales intrinsics if depth resolution differs, + and zeroes invalid-mask depths (nearest interpolation as in paper). + """ + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + ransac=True, + return_aligned=True, + random_state=42, + ) + depths = pred_data.depth * scale + intrinsics = pred_data.intrinsics.copy() + + if depths.shape[-2:] != masks.shape[-2:]: + # When resizing depths to mask size, adjust intrinsics accordingly + sx = masks.shape[-1] / depths.shape[-1] + sy = masks.shape[-2] / depths.shape[-2] + intrinsics[:, 0:1] *= sx + intrinsics[:, 1:2] *= sy + depths = F.interpolate( + torch.from_numpy(depths)[None].float(), + size=(masks.shape[-2], masks.shape[-1]), + mode="nearest", + )[0].numpy() + depths[masks == False] = 0.0 # noqa: E712 + + return depths, intrinsics, extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, masks: np.ndarray + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + + Uses GT intrinsics/extrinsics but aligns scale via Umeyama. + Same mask order as other datasets: mask BEFORE scale. + """ + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + ransac=True, + return_aligned=True, + random_state=42, + ) + depths = pred_data.depth.copy() + intrinsics = gt_data.intrinsics.copy() + extrinsics = gt_data.extrinsics.copy() + + if depths.shape[-2:] != masks.shape[-2:]: + depths = F.interpolate( + torch.from_numpy(depths)[None].float(), + size=(masks.shape[-2], masks.shape[-1]), + mode="nearest", + )[0].numpy() + + # Mask BEFORE scale (same as other datasets) + depths[masks == False] = 0.0 # noqa: E712 + depths = depths * scale + + return depths, intrinsics, extrinsics + + def _build_proj_mats( + self, intrinsics: np.ndarray, extrinsics: np.ndarray + ) -> np.ndarray: + """Compute per-view 4x4 projection matrices from K and [R|t].""" + proj_mat_list = [] + for i in range(len(intrinsics)): + proj_mat = np.eye(4, dtype=np.float32) + proj_mat[:3, :4] = np.dot(intrinsics[i], extrinsics[i][:3]) + proj_mat_list.append(proj_mat) + return np.stack(proj_mat_list, axis=0) + + def _fuse_consistent_points( + self, + depths_t: torch.Tensor, + proj_t: torch.Tensor, + idx: int, + H: int, + W: int, + ) -> np.ndarray: + """Fuse points consistent across multiple source views for a reference index.""" + device, dtype = depths_t.device, depths_t.dtype + pc_buff = torch.zeros((3, H, W), device=device, dtype=dtype) + val_cnt = torch.zeros((1, H, W), device=device, dtype=dtype) + + j = 0 + batch_size = 20 + tot_frame = depths_t.shape[0] + while True: + ref_pc, pcs, dist = self._filter_depth( + ref_depth=depths_t[idx : idx + 1], + src_depths=depths_t[j : min(j + batch_size, tot_frame)], + ref_proj=proj_t[idx : idx + 1], + src_projs=proj_t[j : min(j + batch_size, tot_frame)], + ) + masks = (dist < self.dist_thresh).float() + masked_pc = pcs * masks + pc_buff += masked_pc.sum(dim=0, keepdim=False) + val_cnt += masks.sum(dim=0, keepdim=False) + j += batch_size + if j >= tot_frame: + break + + final_mask = (val_cnt >= self.num_consist).squeeze(0) + avg_points = torch.div(pc_buff, val_cnt).permute(1, 2, 0) + final_pc = self._extract_points(avg_points, final_mask) + return final_pc + + def _cap_points(self, points: np.ndarray, max_points: int) -> np.ndarray: + """Downsample points if exceeding max count.""" + if len(points) <= max_points: + return points + # Use fixed seed for reproducibility + rng = np.random.default_rng(seed=42) + random_idx = rng.choice(len(points), max_points, replace=False) + return points[random_idx] + diff --git a/depth_anything_3/bench/datasets/dtu64.py b/depth_anything_3/bench/datasets/dtu64.py new file mode 100644 index 0000000000000000000000000000000000000000..36ebb63223e1140dff9ff8caaee26cfcefe1200f --- /dev/null +++ b/depth_anything_3/bench/datasets/dtu64.py @@ -0,0 +1,182 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +DTU-64 Dataset implementation for POSE EVALUATION ONLY. + +This is a subset of DTU with 64 images per scene, specifically designed for +camera pose estimation evaluation. It does NOT support 3D reconstruction. + +Note: GT depth loading is not implemented as it's not needed for pose evaluation. +""" + +import glob +import os +from typing import Dict as TDict + +import numpy as np +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.utils.constants import ( + DTU64_CAMERA_ROOT, + DTU64_EVAL_DATA_ROOT, + DTU64_SCENES, +) + + +@MV_REGISTRY.register(name="dtu64") +@MONO_REGISTRY.register(name="dtu64") +class DTU64(Dataset): + """ + DTU-64 Dataset wrapper for DepthAnything3 POSE EVALUATION ONLY. + + This dataset is a subset of DTU with 64 images per scene. + It is specifically designed for camera pose estimation evaluation + and does NOT support 3D reconstruction evaluation. + + Dataset structure: + DTU/scans/ + ├── {scene}/ + │ └── image/ # RGB images (64 per scene) + └── Cameras/ + └── {idx}_cam.txt # Camera parameters + + Supported modes: + - pose: Camera pose estimation evaluation + + NOT supported: + - recon_unposed: 3D reconstruction (no GT depth available) + - recon_posed: 3D reconstruction (no GT depth available) + """ + + data_root = DTU64_EVAL_DATA_ROOT + camera_root = DTU64_CAMERA_ROOT + SCENES = DTU64_SCENES + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Camera file parsing + # ------------------------------ + + def read_cam_file(self, filename: str) -> tuple: + """ + Read DTU camera file containing extrinsics and intrinsics. + + Args: + filename: Path to camera text file + + Returns: + Tuple of (intrinsics [3,3], extrinsics [4,4]) + """ + with open(filename) as f: + lines = [line.rstrip() for line in f.readlines()] + # extrinsics: line [1,5), 4x4 matrix + extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4)) + # intrinsics: line [7-10), 3x3 matrix + intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3)) + return intrinsics, extrinsics + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "scan105") + + Returns: + Dict with: + - image_files: List[str] - paths to images (64 per scene) + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict (empty for this dataset) + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + rgb_folder = os.path.join(self.data_root, scene, "image") + + # Get all PNG files sorted + files = sorted(glob.glob(os.path.join(rgb_folder, "*.png"))) + + # Reorder: place index 33 first (reference view convention) + if len(files) > 33: + files = [files[33]] + files[:33] + files[34:] + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({}), + }) + + for rgb_file in files: + basename = os.path.basename(rgb_file) + # File naming: "00000033.png" -> cam_idx = 33 + file_idx = basename.split(".")[0] + cam_idx = int(file_idx) + + # Camera file path + cam_file = os.path.join(self.camera_root, f"{cam_idx:0>8}_cam.txt") + + if not os.path.exists(cam_file): + print(f"[DTU-64] Warning: Camera file not found: {cam_file}") + continue + + intrinsics, extrinsics = self.read_cam_file(cam_file) + + out.image_files.append(rgb_file) + out.extrinsics.append(extrinsics) + out.intrinsics.append(intrinsics) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[DTU-64] {scene}: {len(out.image_files)} images (pose evaluation only)") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + NOT SUPPORTED for DTU-64. + + DTU-64 is only for pose evaluation, not 3D reconstruction. + """ + raise NotImplementedError( + "DTU-64 dataset is for POSE EVALUATION ONLY. " + "3D reconstruction evaluation is not supported. " + "Use the standard 'dtu' dataset for 3D reconstruction evaluation." + ) + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + NOT SUPPORTED for DTU-64. + + DTU-64 is only for pose evaluation, not 3D reconstruction. + """ + raise NotImplementedError( + "DTU-64 dataset is for POSE EVALUATION ONLY. " + "3D reconstruction (fuse3d) is not supported. " + "Use the standard 'dtu' dataset for 3D reconstruction." + ) + diff --git a/depth_anything_3/bench/datasets/eth3d.py b/depth_anything_3/bench/datasets/eth3d.py new file mode 100644 index 0000000000000000000000000000000000000000..b39069517cca9dadafd54d881e8c119c944af526 --- /dev/null +++ b/depth_anything_3/bench/datasets/eth3d.py @@ -0,0 +1,594 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +ETH3D Benchmark dataset implementation. + +ETH3D is a multi-view stereo benchmark with high-resolution images and +accurate ground truth geometry from laser scanning. +Reference: https://www.eth3d.net/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import glob +import os +from typing import Dict as TDict, List, Optional + +import cv2 +import numpy as np +import open3d as o3d +import torch +import torch.nn.functional as F +from addict import Dict +from PIL import Image + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + quat2rotmat, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + ETH3D_DOWN_SAMPLE, + ETH3D_EVAL_DATA_ROOT, + ETH3D_EVAL_THRESHOLD, + ETH3D_FILTER_KEYS, + ETH3D_MAX_DEPTH, + ETH3D_SAMPLING_NUMBER, + ETH3D_SCENES, + ETH3D_SDF_TRUNC, + ETH3D_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="eth3d") +@MONO_REGISTRY.register(name="eth3d") +class ETH3D(Dataset): + """ + ETH3D Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + eth3d/multiview/ + ├── scene_name/ + │ ├── images/ # RGB images + │ ├── dslr_calibration_jpg/ + │ │ ├── cameras.txt # Camera intrinsics + │ │ └── images.txt # Camera poses + │ ├── combined_mesh.ply # Ground truth mesh + │ └── ground_truth_depth/ # GT depth maps (optional) + """ + + data_root = ETH3D_EVAL_DATA_ROOT + SCENES = ETH3D_SCENES + + # Evaluation hyperparameters from constants + max_depth = ETH3D_MAX_DEPTH + sampling_number = ETH3D_SAMPLING_NUMBER + voxel_length = ETH3D_VOXEL_LENGTH + sdf_trunc = ETH3D_SDF_TRUNC + eval_threshold = ETH3D_EVAL_THRESHOLD + down_sample = ETH3D_DOWN_SAMPLE + + def __init__(self): + super().__init__() + # Pre-load scene data for efficiency + self._scene_cache = {} + + # ------------------------------ + # Camera file parsing + # ------------------------------ + + def _parse_cameras_txt(self, filepath: str) -> dict: + """ + Parse COLMAP-style cameras.txt file. + + Returns: + Dict mapping camera_id to intrinsic parameters + """ + camera_dict = {} + with open(filepath) as f: + lines = f.readlines() + for line in lines[3:]: # Skip header + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 8: + continue + cam_id = parts[0] + # Format: ID, MODEL, WIDTH, HEIGHT, fx, fy, cx, cy, [distortion params...] + camera_dict[cam_id] = { + "width": float(parts[2]), + "height": float(parts[3]), + "fx": float(parts[4]), + "fy": float(parts[5]), + "cx": float(parts[6]), + "cy": float(parts[7]), + } + return camera_dict + + def _parse_images_txt(self, filepath: str) -> dict: + """ + Parse COLMAP-style images.txt file. + + Returns: + Dict mapping image path to pose parameters + """ + pose_dict = {} + with open(filepath) as f: + lines = f.readlines() + for idx, line in enumerate(lines[4:]): # Skip header + line = line.strip() + if not line or line.startswith("#"): + continue + # Every other line contains pose info + if idx % 2 == 0: + parts = line.split() + if len(parts) < 10: + continue + # Format: IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME + image_id = parts[0] + qw, qx, qy, qz = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]) + tx, ty, tz = float(parts[5]), float(parts[6]), float(parts[7]) + camera_id = parts[8] + name = parts[9] + pose_dict[name] = { + "image_id": image_id, + "quat": [qw, qx, qy, qz], + "trans": [tx, ty, tz], + "camera_id": camera_id, + } + return pose_dict + + def _should_filter_image(self, scene: str, image_name: str) -> bool: + """Check if image should be filtered out based on known problematic views.""" + filter_keys = ETH3D_FILTER_KEYS.get(scene, []) + for key in filter_keys: + if image_name.endswith(key): + return True + return False + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "courtyard") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path + """ + # Check cache + if scene in self._scene_cache: + return self._scene_cache[scene] + + scene_dir = os.path.join(self.data_root, scene) + + # Parse camera files + cameras_file = os.path.join(scene_dir, "dslr_calibration_jpg", "cameras.txt") + images_file = os.path.join(scene_dir, "dslr_calibration_jpg", "images.txt") + camera_dict = self._parse_cameras_txt(cameras_file) + pose_dict = self._parse_images_txt(images_file) + + # Ground truth mesh path + gt_mesh_path = os.path.join(scene_dir, "combined_mesh.ply") + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "heights": [], + "widths": [], + }), + }) + + # Process each image (preserve original order from images.txt) + filtered_count = 0 + for image_name, pose_info in pose_dict.items(): + # Filter problematic views + if self._should_filter_image(scene, image_name): + filtered_count += 1 + continue + + image_path = os.path.join(scene_dir, "images", image_name) + if not os.path.exists(image_path): + continue + + cam_info = camera_dict.get(pose_info["camera_id"]) + if cam_info is None: + continue + + # Build intrinsics matrix + ixt = np.array([ + [cam_info["fx"], 0, cam_info["cx"]], + [0, cam_info["fy"], cam_info["cy"]], + [0, 0, 1], + ], dtype=np.float32) + + # Build extrinsics matrix (world-to-camera) + # COLMAP format: world point -> camera point + rot = quat2rotmat(pose_info["quat"]) + ext = np.eye(4, dtype=np.float32) + ext[:3, :3] = rot + ext[:3, 3] = pose_info["trans"] + + out.image_files.append(image_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.heights.append(cam_info["height"]) + out.aux.widths.append(cam_info["width"]) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + # Print scene info + total_images = len(pose_dict) + used_images = len(out.image_files) + print(f"[ETH3D] {scene}: {used_images}/{total_images} images " + f"(filtered {filtered_count}, missing {total_images - used_images - filtered_count})") + + if used_images < 3: + print(f"[ETH3D] ⚠️ WARNING: {scene} has only {used_images} images - evaluation may fail!") + + # Cache result + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against ETH3D ground truth mesh. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load and sample ground truth mesh + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """ + Load saved GT meta (extrinsics, intrinsics, image_files) for fusion. + + This is needed when frames are sampled, so fuse3d uses the correct + (sampled) GT instead of full dataset GT. + + Args: + result_path: Path to npz file (used to derive gt_meta.npz path) + + Returns: + Dict with GT data, or None if gt_meta.npz doesn't exist + """ + # gt_meta.npz is in the same exports/ directory as results.npz + export_dir = os.path.dirname(result_path) # exports/mini_npz/ + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": data["image_files"] if "image_files" in data else None, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Pipeline: + 1. Load original images (keep original size) + 2. Resize depth to original image size (nearest interpolation) + 3. Adjust intrinsics to original image size + 4. Apply scale alignment and mask invalid depths + 5. TSDF fusion + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + else: + gt_data = self.get_data(scene) + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load original images (keep original size) + images = [] + orig_sizes = [] # (H, W) for each image + for img_path in gt_data.image_files: + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + # Prepare depths, intrinsics, extrinsics with resize to original size + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + images = np.stack(images, axis=0) + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Pipeline: + 1. Umeyama scale alignment + 2. Load GT mask for each frame + 3. Resize depth to original image size (nearest) + 4. Apply GT mask BEFORE scale + 5. Apply scale + 6. Adjust intrinsics to original image size + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + # Get model output size + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + # Process each frame + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + image_name = os.path.basename(gt_data.image_files[i]) + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask (apply BEFORE scale) + gt_zero_mask = None + if scene is not None: + gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w)) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio # fx, 0, cx + ixt[1, :] *= h_ratio # 0, fy, cy + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + + Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama. + Depth is resized to original image size. + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + # Process each frame + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + image_name = os.path.basename(gt_data.image_files[i]) + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask (apply BEFORE scale) + gt_zero_mask = None + if scene is not None: + gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w)) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics (already at original image size) + return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy() + + def _load_gt_mask(self, scene: str, image_name: str, shape: tuple) -> np.ndarray: + """ + Load GT mask for masking invalid regions. + + GT mask marks occluded or invalid regions that should be excluded + from depth fusion and evaluation. + + Args: + scene: Scene identifier + image_name: Image filename (e.g., "DSC_0307.JPG") + shape: (height, width) of the image + + Returns: + Boolean mask where True = valid region to keep + """ + h, w = shape + + # GT mask file path + gt_mask_path = os.path.join( + self.data_root, scene, "masks_for_images", "dslr_images", + image_name.replace(".JPG", ".png") + ) + + # GT depth file path (used to determine valid depth regions) + gt_depth_path = os.path.join( + self.data_root, scene, "ground_truth_depth", "dslr_images", image_name + ) + + # Load GT depth + if os.path.exists(gt_depth_path): + gt_depth = np.fromfile(gt_depth_path, dtype=np.float32).reshape(h, w) + else: + gt_depth = np.ones((h, w), dtype=np.float32) + + # Load GT mask + if os.path.exists(gt_mask_path): + gt_mask = cv2.imread(gt_mask_path, cv2.IMREAD_GRAYSCALE) + gt_mask = np.asarray(gt_mask) + else: + gt_mask = np.zeros((h, w), dtype=np.uint8) + + # Compute zero_mask + # gt_mask == 1 means occluded/invalid region + invalid_mask_from_gt = gt_mask == 1 + gt_depth_copy = gt_depth.copy() + gt_depth_copy[gt_mask == 1] = 0 + + invalid_mask_from_gt_depth = np.logical_or(gt_depth_copy == 0, gt_depth_copy == np.inf) + + # zero_mask: valid region that should be kept + zero_mask = np.logical_and( + np.logical_not(invalid_mask_from_gt), + np.logical_not(invalid_mask_from_gt_depth) + ) + + return zero_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """ + Mask invalid depth values by setting them to 0. + + Logic: + 1. Apply GT mask (if provided) - marks occluded/invalid regions + 2. Mask pred invalid values (nan, inf) + + Args: + depth: Depth map to mask + gt_zero_mask: Optional GT mask (True = valid region) + + Returns: + Masked depth map with invalid regions set to 0 + """ + depth = depth.copy() + + # Apply GT mask first (before scale) + if gt_zero_mask is not None: + # Also mask out invalid pred depth + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + # Fallback: only mask pred invalid values + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/depth_anything_3/bench/datasets/hiroom.py b/depth_anything_3/bench/datasets/hiroom.py new file mode 100644 index 0000000000000000000000000000000000000000..45619f76a4d7409bf5e9674815a10d0bd02085f5 --- /dev/null +++ b/depth_anything_3/bench/datasets/hiroom.py @@ -0,0 +1,440 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +HiRoom Dataset implementation. + +HiRoom is an indoor RGB-D dataset containing ground truth camera poses, +depth maps, and fused point clouds. + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict, List + +import cv2 +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + HIROOM_DOWN_SAMPLE, + HIROOM_EVAL_DATA_ROOT, + HIROOM_EVAL_THRESHOLD, + HIROOM_GT_ROOT_PATH, + HIROOM_MAX_DEPTH, + HIROOM_SAMPLING_NUMBER, + HIROOM_SCENE_LIST_PATH, + HIROOM_SDF_TRUNC, + HIROOM_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +def _load_scene_list() -> List[str]: + """Load scene list from file.""" + if os.path.exists(HIROOM_SCENE_LIST_PATH): + with open(HIROOM_SCENE_LIST_PATH, "r") as f: + return f.read().splitlines() + return [] + + +@MV_REGISTRY.register(name="hiroom") +@MONO_REGISTRY.register(name="hiroom") +class HiRoomDataset(Dataset): + """ + HiRoom Dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + HiRoom/ + ├── {scene_path}/ + │ ├── image/ # RGB images + │ ├── depth/ # GT depth maps + │ ├── pose/ # Camera poses (.npy) + │ ├── cam_K.npy # Camera intrinsics + │ └── aliasing_mask/ # Aliasing masks + + fused_pcd/ + └── {scene_name}.ply # Ground truth fused point cloud + """ + + data_root = HIROOM_EVAL_DATA_ROOT + gt_root_path = HIROOM_GT_ROOT_PATH + SCENES = _load_scene_list() + + # Evaluation hyperparameters from constants + max_depth = HIROOM_MAX_DEPTH + sampling_number = HIROOM_SAMPLING_NUMBER + voxel_length = HIROOM_VOXEL_LENGTH + sdf_trunc = HIROOM_SDF_TRUNC + eval_threshold = HIROOM_EVAL_THRESHOLD + down_sample = HIROOM_DOWN_SAMPLE + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene path (e.g., "xxx/yyy/zzz") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_pcd_path, gt_depth_files, aliasing_mask_files + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + scene_dir = os.path.join(self.data_root, scene) + image_dir = os.path.join(scene_dir, "image") + + # Get scene name for GT point cloud + scene_name = "-".join(scene.split("/")[-3:]) + gt_pcd_path = os.path.join(self.gt_root_path, f"{scene_name}.ply") + + # Load shared camera intrinsics + intrin_path = os.path.join(scene_dir, "cam_K.npy") + ixt_shared = np.load(intrin_path).astype(np.float32) + + # Get all image names sorted + image_names = sorted(os.listdir(image_dir)) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_pcd_path": gt_pcd_path, + "gt_depth_files": [], + "aliasing_mask_files": [], + }), + }) + + for img_name in image_names: + img_path = os.path.join(image_dir, img_name) + frame_name = img_name.split(".")[0] + + # Depth and pose paths + depth_path = os.path.join(scene_dir, "depth", f"{frame_name}.png") + pose_path = os.path.join(scene_dir, "pose", f"{frame_name}.npy") + aliasing_mask_path = os.path.join(scene_dir, "aliasing_mask", f"{frame_name}.png") + + if not os.path.exists(pose_path): + continue + + # Load extrinsics (world-to-camera) + ext = np.load(pose_path).astype(np.float32) + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt_shared.copy()) + out.aux.gt_depth_files.append(depth_path) + out.aux.aliasing_mask_files.append(aliasing_mask_path) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[HiRoom] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against HiRoom ground truth point cloud. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_pcd_path = gt_data.aux.gt_pcd_path + + # Load ground truth point cloud + gt_pcd = o3d.io.read_point_cloud(gt_pcd_path) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """Load saved GT meta for fusion.""" + export_dir = os.path.dirname(result_path) + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + image_files = list(data["image_files"]) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Get full GT data + full_gt_data = self.get_data(scene) + + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + image_indices = [ + full_gt_data.image_files.index(f) + for f in gt_data.image_files + if f in full_gt_data.image_files + ] + else: + gt_data = full_gt_data + image_indices = list(range(len(full_gt_data.image_files))) + + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load images + images = [] + orig_sizes = [] + for img_idx in image_indices: + img_path = full_gt_data.image_files[img_idx] + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + images = np.stack(images, axis=0) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, orig_sizes: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_unposed mode.""" + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + img_idx = image_indices[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask + gt_zero_mask = self._load_gt_mask( + full_gt_data.aux.gt_depth_files[img_idx], + full_gt_data.aux.aliasing_mask_files[img_idx], + ) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, orig_sizes: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_posed mode.""" + # Scale alignment + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + img_idx = image_indices[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask + gt_zero_mask = self._load_gt_mask( + full_gt_data.aux.gt_depth_files[img_idx], + full_gt_data.aux.aliasing_mask_files[img_idx], + ) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics + gt_intrinsics = np.stack([full_gt_data.intrinsics[idx] for idx in image_indices]) + gt_extrinsics = np.stack([full_gt_data.extrinsics[idx] for idx in image_indices]) + + return np.stack(depths_out), gt_intrinsics, gt_extrinsics + + def _load_gt_mask(self, gt_depth_path: str, aliasing_mask_path: str) -> np.ndarray: + """ + Load GT depth and aliasing mask to create valid mask. + + For HiRoom: + - GT depth is stored as 16-bit PNG, scaled to 100m range + - Aliasing mask marks regions to exclude + + Returns: + Boolean mask where True = valid region to keep + """ + # Load GT depth + if os.path.exists(gt_depth_path): + gt_depth = cv2.imread(gt_depth_path, -1) / 65535.0 * 100.0 + else: + return None + + # Load aliasing mask + aliasing_mask = None + if os.path.exists(aliasing_mask_path): + aliasing_mask = cv2.imread(aliasing_mask_path, -1) > 0 + + # Valid mask: depth > 0 and not in aliasing region + valid_mask = gt_depth > 0 + if aliasing_mask is not None: + valid_mask = np.logical_and(valid_mask, np.logical_not(aliasing_mask)) + + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """Mask invalid depth values by setting them to 0.""" + depth = depth.copy() + + if gt_zero_mask is not None: + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/depth_anything_3/bench/datasets/scannetpp.py b/depth_anything_3/bench/datasets/scannetpp.py new file mode 100644 index 0000000000000000000000000000000000000000..f841ad11c701a93792ef8814f8093c8a60110d74 --- /dev/null +++ b/depth_anything_3/bench/datasets/scannetpp.py @@ -0,0 +1,591 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +ScanNet++ Benchmark dataset implementation. + +ScanNet++ is a high-quality indoor RGB-D dataset with iPhone and DSLR images, +ground truth camera poses from COLMAP, and high-resolution 3D meshes. +Reference: https://kaldir.vc.in.tum.de/scannetpp/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict + +import cv2 +import imageio +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + fuse_depth_to_tsdf, + nn_correspondance, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + SCANNETPP_DOWN_SAMPLE, + SCANNETPP_EVAL_DATA_ROOT, + SCANNETPP_EVAL_THRESHOLD, + SCANNETPP_INPUT_H, + SCANNETPP_INPUT_W, + SCANNETPP_MAX_DEPTH, + SCANNETPP_SAMPLING_NUMBER, + SCANNETPP_SCENES, + SCANNETPP_SDF_TRUNC, + SCANNETPP_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama +from depth_anything_3.utils.read_write_model import read_model + + +@MV_REGISTRY.register(name="scannetpp") +@MONO_REGISTRY.register(name="scannetpp") +class ScanNetPP(Dataset): + """ + ScanNet++ Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + scannetpp/data/ + ├── {scene_id}/ + │ ├── merge_dslr_iphone/ + │ │ ├── colmap/sparse_render_rgb/ # COLMAP reconstruction + │ │ ├── images/ # RGB images + │ │ └── render_depth/ # GT depth maps + │ └── scans/ + │ └── mesh_aligned_0.05.ply # Ground truth mesh + """ + + data_root = SCANNETPP_EVAL_DATA_ROOT + SCENES = SCANNETPP_SCENES + + # Input resolution after undistortion and resize + input_h = SCANNETPP_INPUT_H + input_w = SCANNETPP_INPUT_W + + # Evaluation hyperparameters from constants + max_depth = SCANNETPP_MAX_DEPTH + sampling_number = SCANNETPP_SAMPLING_NUMBER + voxel_length = SCANNETPP_VOXEL_LENGTH + sdf_trunc = SCANNETPP_SDF_TRUNC + eval_threshold = SCANNETPP_EVAL_THRESHOLD + down_sample = SCANNETPP_DOWN_SAMPLE + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Only uses iPhone images (not DSLR). + + Args: + scene: Scene identifier (e.g., "09c1414f1b") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path, dist, roi, cam_hw, etc. + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + input_path = os.path.join(self.data_root, scene, "merge_dslr_iphone") + colmap_path = os.path.join(input_path, "colmap/sparse_render_rgb") + image_path = os.path.join(input_path, "images") + depth_path_dir = os.path.join(input_path, "render_depth") + + # Read COLMAP model + cams, images, points3d = read_model(colmap_path) + + # Map image names to IDs + name2id = {image.name: k for k, image in images.items()} + names = sorted([image.name for k, image in images.items()]) + # Only use iPhone images + names = [name for name in names if "iphone" in name] + + gt_mesh_path = os.path.join( + input_path.replace("merge_dslr_iphone", "scans"), "mesh_aligned_0.05.ply" + ) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "dist_list": [], + "roi_list": [], + "cam_hw_list": [], + "ixt_raw_list": [], + "gt_depth_files": [], + }), + }) + + for name in names: + image = images[name2id[name]] + img_path = os.path.join(image_path, name) + + if not os.path.exists(img_path): + continue + + # Build extrinsics (world-to-camera) + ext = np.eye(4, dtype=np.float32) + ext[:3, :3] = image.qvec2rotmat() + ext[:3, 3] = image.tvec + + # Get camera parameters + cam_id = image.camera_id + camera = cams[cam_id] + cam_height, cam_width = camera.height, camera.width + + # Build intrinsics + ixt = np.eye(3, dtype=np.float32) + ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2] = camera.params[:4] + ixt[:2, 2] -= 0.5 # COLMAP convention adjustment + ixt_raw = ixt.copy() + + # Handle distortion (OPENCV model) + dist = np.zeros(5, dtype=np.float32) + roi = (0, 0, cam_width, cam_height) + if camera.model == "OPENCV": + dist[:4] = camera.params[4:] + ixt, roi = cv2.getOptimalNewCameraMatrix( + ixt, dist, (cam_width, cam_height), 1, (cam_width, cam_height) + ) + + # Depth file path + frame_name = os.path.basename(name)[:-4] # Remove .jpg + depth_file = os.path.join(depth_path_dir, f"{frame_name}.png") + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.dist_list.append(dist) + out.aux.roi_list.append(roi) + out.aux.cam_hw_list.append((cam_height, cam_width)) + out.aux.ixt_raw_list.append(ixt_raw) + out.aux.gt_depth_files.append(depth_file) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[ScanNet++] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def load_image(self, img_path: str, idx: int, aux: Dict) -> np.ndarray: + """ + Load and preprocess image with undistortion and cropping. + + Args: + img_path: Path to image file + idx: Index of the image in the dataset + aux: Auxiliary data from get_data + + Returns: + Preprocessed RGB image + """ + image = imageio.imread(img_path).astype(np.uint8) + ixt_raw = aux.ixt_raw_list[idx] + ixt = aux.intrinsics[idx] if hasattr(aux, 'intrinsics') else None + dist = aux.dist_list[idx] + roi = aux.roi_list[idx] + + # Undistort using raw intrinsics + # Use the stored intrinsics from get_data for newCameraMatrix + stored_ixt = self._scene_cache.get(aux.scene, {}).get('intrinsics', [None])[idx] if hasattr(aux, 'scene') else None + if stored_ixt is None: + # Recompute optimal camera matrix for undistortion + cam_h, cam_w = aux.cam_hw_list[idx] + ixt_for_undistort = ixt_raw.copy() + ixt_for_undistort, _ = cv2.getOptimalNewCameraMatrix( + ixt_raw, dist, (cam_w, cam_h), 1, (cam_w, cam_h) + ) + else: + ixt_for_undistort = stored_ixt + + image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt_for_undistort) + + # Crop to ROI + x, y, w, h = roi + image = image[y:y+h, x:x+w] + + # Resize to target resolution + image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA) + + return image + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against ScanNet++ ground truth mesh. + + Uses AABB cropping to only evaluate points within GT bounding box. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load ground truth mesh and sample points + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Crop prediction to GT bounding box (with 0.1m margin) + aabb = gt_pcd.get_axis_aligned_bounding_box() + points = np.asarray(pred_pcd.points) + inside_mask = ( + (points[:, 0] >= aabb.min_bound[0] - 0.1) & + (points[:, 0] <= aabb.max_bound[0] + 0.1) & + (points[:, 1] >= aabb.min_bound[1] - 0.1) & + (points[:, 1] <= aabb.max_bound[1] + 0.1) & + (points[:, 2] >= aabb.min_bound[2] - 0.1) & + (points[:, 2] <= aabb.max_bound[2] + 0.1) + ) + pred_pcd = pred_pcd.select_by_index(inside_mask.nonzero()[0]) + + # Downsample + if self.down_sample > 0: + pred_pcd = pred_pcd.voxel_down_sample(self.down_sample) + gt_pcd = gt_pcd.voxel_down_sample(self.down_sample) + + verts_pred = np.asarray(pred_pcd.points) + verts_gt = np.asarray(gt_pcd.points) + + if len(verts_pred) == 0 or len(verts_gt) == 0: + return { + "acc": float("inf"), + "comp": float("inf"), + "overall": float("inf"), + "precision": 0.0, + "recall": 0.0, + "fscore": 0.0, + } + + # Compute distances + dist_pred_to_gt = nn_correspondance(verts_gt, verts_pred) + dist_gt_to_pred = nn_correspondance(verts_pred, verts_gt) + + # Compute metrics + accuracy = float(np.mean(dist_pred_to_gt)) + completeness = float(np.mean(dist_gt_to_pred)) + overall = (accuracy + completeness) / 2 + + precision = float(np.mean((dist_pred_to_gt < self.eval_threshold).astype(float))) + recall = float(np.mean((dist_gt_to_pred < self.eval_threshold).astype(float))) + + if precision + recall > 0: + fscore = 2 * precision * recall / (precision + recall) + else: + fscore = 0.0 + + return { + "acc": accuracy, + "comp": completeness, + "overall": overall, + "precision": precision, + "recall": recall, + "fscore": fscore, + } + + def _load_gt_meta(self, result_path: str) -> Dict: + """Load saved GT meta for fusion.""" + export_dir = os.path.dirname(result_path) + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + image_files = list(data["image_files"]) + + # Reconstruct aux data from image files + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Get GT data + full_gt_data = self.get_data(scene) + + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + # Need to rebuild aux from full GT data based on image indices + image_indices = [ + full_gt_data.image_files.index(f) + for f in gt_data.image_files + if f in full_gt_data.image_files + ] + else: + gt_data = full_gt_data + image_indices = list(range(len(full_gt_data.image_files))) + + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load and preprocess images + images = [] + for idx, img_idx in enumerate(image_indices): + img_path = full_gt_data.image_files[img_idx] + image = imageio.imread(img_path).astype(np.uint8) + + # Undistort and crop + ixt_raw = full_gt_data.aux.ixt_raw_list[img_idx] + ixt = full_gt_data.intrinsics[img_idx] + dist = full_gt_data.aux.dist_list[img_idx] + roi = full_gt_data.aux.roi_list[img_idx] + + image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt) + x, y, w, h = roi + image = image[y:y+h, x:x+w] + image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA) + + images.append(image) + + images = np.stack(images, axis=0) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, full_gt_data, image_indices, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, full_gt_data, image_indices, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_unposed mode.""" + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + img_idx = image_indices[i] + + # Get original image size (after undistort+crop, before resize to input_h/w) + orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx] + + # Step 1: nearest resize to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Step 2: linear resize to target resolution + depth = cv2.resize( + depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to target resolution + h_ratio = self.input_h / model_h + w_ratio = self.input_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_posed mode.""" + # Scale alignment + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + depths_out = [] + intrinsics_out = [] + extrinsics_out = [] + + for i in range(len(pred_data.depth)): + img_idx = image_indices[i] + + # Get original image size (after undistort+crop, before resize to input_h/w) + orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx] + + # Step 1: nearest resize to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Step 2: linear resize to target resolution + depth = cv2.resize( + depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Get GT intrinsics and scale to target resolution + ixt = full_gt_data.intrinsics[img_idx].copy() + cam_h, cam_w = full_gt_data.aux.cam_hw_list[img_idx] + ixt[:2, 2] += 0.5 # Undo COLMAP convention + ixt[0, :] *= self.input_w / cam_w + ixt[1, :] *= self.input_h / cam_h + intrinsics_out.append(ixt) + + extrinsics_out.append(full_gt_data.extrinsics[img_idx]) + + return np.stack(depths_out), np.stack(intrinsics_out), np.stack(extrinsics_out) + + def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray: + """ + Load GT depth and create valid mask. + + For ScanNet++, GT depth is stored as 16-bit PNG in millimeters. + + Returns: + Boolean mask where True = valid region to keep + """ + if not os.path.exists(gt_depth_path): + return None + + gt_depth = imageio.imread(gt_depth_path) / 1000.0 # mm to meters + + # Resize to target resolution + gt_depth = cv2.resize( + gt_depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Valid mask: depth > 0 and not inf + valid_mask = np.logical_and(gt_depth > 0, gt_depth != np.inf) + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """Mask invalid depth values by setting them to 0.""" + depth = depth.copy() + + if gt_zero_mask is not None: + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/depth_anything_3/bench/datasets/sevenscenes.py b/depth_anything_3/bench/datasets/sevenscenes.py new file mode 100644 index 0000000000000000000000000000000000000000..cce82b385d4460f948d240f3c0991b7570d2177a --- /dev/null +++ b/depth_anything_3/bench/datasets/sevenscenes.py @@ -0,0 +1,449 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +7Scenes Benchmark dataset implementation. + +7Scenes is an indoor RGB-D dataset with ground truth camera poses and 3D meshes. +Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict + +import cv2 +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + SEVENSCENES_CX, + SEVENSCENES_CY, + SEVENSCENES_DOWN_SAMPLE, + SEVENSCENES_EVAL_DATA_ROOT, + SEVENSCENES_EVAL_THRESHOLD, + SEVENSCENES_FX, + SEVENSCENES_FY, + SEVENSCENES_MAX_DEPTH, + SEVENSCENES_SAMPLING_NUMBER, + SEVENSCENES_SCENES, + SEVENSCENES_SDF_TRUNC, + SEVENSCENES_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="7scenes") +@MONO_REGISTRY.register(name="7scenes") +class SevenScenes(Dataset): + """ + 7Scenes Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + 7scenes/ + ├── 7Scenes/ + │ ├── {scene}/ + │ │ └── seq-01/ (or seq-02 for stairs) + │ │ ├── frame-XXXXXX.color.png + │ │ ├── frame-XXXXXX.depth.png + │ │ └── frame-XXXXXX.pose.txt + │ └── meshes/ + │ └── {scene}.ply # Ground truth mesh + """ + + data_root = SEVENSCENES_EVAL_DATA_ROOT + SCENES = SEVENSCENES_SCENES + + # Evaluation hyperparameters from constants + max_depth = SEVENSCENES_MAX_DEPTH + sampling_number = SEVENSCENES_SAMPLING_NUMBER + voxel_length = SEVENSCENES_VOXEL_LENGTH + sdf_trunc = SEVENSCENES_SDF_TRUNC + eval_threshold = SEVENSCENES_EVAL_THRESHOLD + down_sample = SEVENSCENES_DOWN_SAMPLE + + # Fixed camera intrinsics for all 7Scenes images + fx = SEVENSCENES_FX + fy = SEVENSCENES_FY + cx = SEVENSCENES_CX + cy = SEVENSCENES_CY + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "chess") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path, gt_depth_files + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + # Different sequence for stairs scene + if scene == "stairs": + data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-02") + n_imgs = 500 + else: + data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-01") + n_imgs = 1000 + + gt_mesh_path = os.path.join(self.data_root, "7Scenes", "meshes", f"{scene}.ply") + + # Fixed intrinsics for all images + ixt = np.array([ + [self.fx, 0, self.cx], + [0, self.fy, self.cy], + [0, 0, 1], + ], dtype=np.float32) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "gt_depth_files": [], + }), + }) + + for i in range(0, n_imgs, 1): + img_path = os.path.join(data_folder, f"frame-{i:06d}.color.png") + pose_path = os.path.join(data_folder, f"frame-{i:06d}.pose.txt") + depth_path = os.path.join(data_folder, f"frame-{i:06d}.depth.png") + + if not os.path.exists(img_path) or not os.path.exists(pose_path): + continue + + # Load camera-to-world pose and convert to world-to-camera (extrinsic) + c2w = np.loadtxt(pose_path) + ext = np.linalg.inv(c2w).astype(np.float32) + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt.copy()) + out.aux.gt_depth_files.append(depth_path) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[7Scenes] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against 7Scenes ground truth mesh. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load and sample ground truth mesh + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """ + Load saved GT meta (extrinsics, intrinsics, image_files) for fusion. + + This is needed when frames are sampled, so fuse3d uses the correct + (sampled) GT instead of full dataset GT. + + Args: + result_path: Path to npz file (used to derive gt_meta.npz path) + + Returns: + Dict with GT data, or None if gt_meta.npz doesn't exist + """ + export_dir = os.path.dirname(result_path) # exports/mini_npz/ + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + # Build aux with gt_depth_files derived from image_files + image_files = list(data["image_files"]) + gt_depth_files = [ + img_path.replace("color", "depth").replace(".color.", ".depth.") + for img_path in image_files + ] + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + "aux": Dict({"gt_depth_files": gt_depth_files}), + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + else: + gt_data = self.get_data(scene) + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load original images (keep original size) + images = [] + orig_sizes = [] + for img_path in gt_data.image_files: + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + images = np.stack(images, axis=0) + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Similar to ETH3D but uses GT depth for masking instead of separate mask files. + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama. + """ + # Scale alignment with fixed random_state + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics + return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy() + + def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray: + """ + Load GT depth and create valid mask. + + For 7Scenes, GT depth is stored as 16-bit PNG in millimeters. + Value 65535 indicates invalid depth. + + Returns: + Boolean mask where True = valid region to keep + """ + if not os.path.exists(gt_depth_path): + return None + + gt_depth = cv2.imread(gt_depth_path, -1) + if gt_depth is None: + return None + + # 65535 is invalid depth marker in 7Scenes + gt_depth[gt_depth == 65535] = 0 + # Convert to meters + gt_depth = gt_depth / 1000.0 + + # Valid mask: depth > 0 + valid_mask = gt_depth > 0 + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """ + Mask invalid depth values by setting them to 0. + + Args: + depth: Depth map to mask + gt_zero_mask: Optional GT mask (True = valid region) + + Returns: + Masked depth map with invalid regions set to 0 + """ + depth = depth.copy() + + if gt_zero_mask is not None: + # Also mask out invalid pred depth + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + # Fallback: only mask pred invalid values + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + + diff --git a/depth_anything_3/bench/evaluator.py b/depth_anything_3/bench/evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..5b93b98b1017d25060f0541c1201bfaea0806b14 --- /dev/null +++ b/depth_anything_3/bench/evaluator.py @@ -0,0 +1,752 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Main Evaluator class for DepthAnything3 benchmark evaluation. + +Supports multiple datasets and evaluation modes: +- pose: Camera pose estimation (AUC metrics) +- recon_unposed: 3D reconstruction with predicted poses +- recon_posed: 3D reconstruction with GT poses +- view_syn: Novel view synthesis (TODO) +""" + +import json +import os +import random +from typing import Dict as TDict, Iterable, List + +import numpy as np +import torch +from addict import Dict +from tqdm import tqdm + +from depth_anything_3.bench.print_metrics import MetricsPrinter +from depth_anything_3.utils.parallel_utils import parallel_execution +from depth_anything_3.bench.registries import MV_REGISTRY +from depth_anything_3.utils.constants import EVAL_REF_VIEW_STRATEGY + + +class Evaluator: + """ + Main evaluation orchestrator for DepthAnything3 benchmarks. + + Usage: + evaluator = Evaluator( + work_dir="./eval_workspace", + datas=["dtu"], + modes=["pose", "recon_unposed", "recon_posed"], + ) + api = DepthAnything3.from_pretrained("...") + evaluator.infer(api) + metrics = evaluator.eval() + evaluator.print_metrics() + """ + + VALID_MODES = {"pose", "recon_unposed", "recon_posed", "view_syn"} + + def __init__( + self, + work_dir: str = "./eval_workspace", + datas: List[str] = ("dtu",), + modes: List[str] = ("recon_unposed",), + ref_view_strategy: str = EVAL_REF_VIEW_STRATEGY, + scenes: List[str] = None, + debug: bool = False, + num_fusion_workers: int = 4, + max_frames: int = 100, + gpu_id: int = 0, + total_gpus: int = 1, + ): + """ + Initialize the evaluator. + + Args: + work_dir: Base directory for model outputs and metric files + datas: List of dataset names (must be registered in MV_REGISTRY) + modes: List of evaluation modes to run + ref_view_strategy: Reference view selection strategy for inference + ("first", "saddle_balanced", etc.) + scenes: Specific scenes to evaluate (None = all scenes) + debug: Enable verbose debug output + num_fusion_workers: Number of parallel workers for TSDF fusion (default: 4) + max_frames: Maximum number of frames per scene (default: 100). + If a scene has more frames, randomly sample to this limit. + Set to -1 to disable sampling. + gpu_id: GPU index for multi-GPU (0-indexed) + total_gpus: Total number of GPUs for task distribution + """ + self.work_dir = work_dir + self.datas = list(datas) + self.modes = set(modes) + self.ref_view_strategy = ref_view_strategy + self.scenes_filter = scenes + self.debug = debug + self.num_fusion_workers = num_fusion_workers + self.max_frames = max_frames + self.gpu_id = gpu_id + self.total_gpus = total_gpus + + # Validate modes + unknown = self.modes - self.VALID_MODES + if unknown: + raise ValueError(f"Unknown modes: {unknown}. Valid: {sorted(self.VALID_MODES)}") + + os.makedirs(self.work_dir, exist_ok=True) + + # Initialize datasets + self.datasets = Dict() + for data in self.datas: + if not MV_REGISTRY.has(data): + available = list(MV_REGISTRY.all().keys()) + raise ValueError(f"Dataset '{data}' not found. Available: {available}") + self.datasets[data] = MV_REGISTRY.get(data)() + + # Initialize metrics printer + self._printer = MetricsPrinter() + + # -------------------- Public APIs -------------------- # + + def all(self, api) -> TDict[str, dict]: + """ + Run complete evaluation pipeline: inference + evaluation. + + Args: + api: DepthAnything3 API instance + + Returns: + Combined metrics dictionary + """ + self.infer(api) + return self.eval() + + def _get_scenes(self, dataset) -> List[str]: + """Get list of scenes to evaluate, optionally filtered.""" + all_scenes = dataset.SCENES + if self.scenes_filter: + scenes = [s for s in all_scenes if s in self.scenes_filter] + if self.debug: + print(f"[DEBUG] Filtered scenes: {scenes} (from {len(all_scenes)} total)") + return scenes + return all_scenes + + def infer(self, api, model_path: str = None) -> None: + """ + Run inference according to requested modes. + + - Unposed export if 'pose' or 'recon_unposed' is in modes + - Posed export if 'recon_posed' or 'view_syn' is in modes + + Multi-GPU: Use --gpu_id and --total_gpus to distribute tasks. + Example: Launch 4 processes with gpu_id=0,1,2,3 and total_gpus=4 + + Args: + api: DepthAnything3 API instance + model_path: Model path (unused, kept for API compatibility) + """ + need_unposed = {"pose", "recon_unposed"} & self.modes + need_posed = {"recon_posed", "view_syn"} & self.modes + export_format = "mini_npz-glb" if self.debug else "mini_npz" + + # Collect all tasks + all_tasks = [] + for data in self.datas: + dataset = self.datasets[data] + for scene in self._get_scenes(dataset): + all_tasks.append((data, scene)) + + # Distribute tasks across GPUs + if self.total_gpus > 1: + tasks = [t for i, t in enumerate(all_tasks) if i % self.total_gpus == self.gpu_id] + print(f"[INFO] GPU {self.gpu_id}/{self.total_gpus}: {len(tasks)}/{len(all_tasks)} tasks") + else: + tasks = all_tasks + print(f"[INFO] Total inference tasks: {len(tasks)}") + + for data, scene in tqdm(tasks, desc=f"Inference (GPU {self.gpu_id})"): + dataset = self.datasets[data] + scene_data = dataset.get_data(scene) + scene_data = self._sample_frames(scene_data, scene) + + if need_unposed: + export_dir = self._export_dir(data, scene, posed=False) + api.inference( + scene_data.image_files, + export_dir=export_dir, + export_format=export_format, + ref_view_strategy=self.ref_view_strategy, + ) + self._save_gt_meta(export_dir, scene_data) + + if need_posed: + export_dir = self._export_dir(data, scene, posed=True) + api.inference( + scene_data.image_files, + scene_data.extrinsics, + scene_data.intrinsics, + export_dir=export_dir, + export_format=export_format, + ref_view_strategy=self.ref_view_strategy, + ) + self._save_gt_meta(export_dir, scene_data) + + def eval(self) -> TDict[str, dict]: + """ + Evaluate for all configured modes and write JSON files. + + Evaluation order by mode (all datasets per mode): + 1. pose - all datasets + 2. recon_unposed - all datasets + 3. recon_posed - all datasets + + Returns: + Summary mapping: {"_": metrics_dict} + """ + summary: TDict[str, dict] = {} + + # Evaluate by mode (all datasets per mode) + if "pose" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating POSE for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_pose(): + summary[f"{data}_pose"] = result + + if "recon_unposed" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating RECON_UNPOSED for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_reconstruction("recon_unposed"): + summary[f"{data}_recon_unposed"] = result + + if "recon_posed" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating RECON_POSED for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_reconstruction("recon_posed"): + summary[f"{data}_recon_posed"] = result + + if "view_syn" in self.modes: + # TODO: Add view synthesis metrics here when available + pass + + return summary + + def print_metrics(self, metrics: TDict[str, dict] = None) -> None: + """ + Print evaluation metrics in a beautiful tabular format. + + Args: + metrics: Metrics dictionary. If None, loads from saved JSON files. + """ + if metrics is None: + metrics = self._load_metrics() + + self._printer.print_results(metrics) + + # -------------------- Evaluation Methods -------------------- # + + def _eval_pose(self) -> Iterable[tuple]: + """Compute pose-estimation metrics for each dataset and scene.""" + os.makedirs(self._metric_dir, exist_ok=True) + + for data in tqdm(self.datas, desc="Datasets (pose eval)"): + dataset = self.datasets[data] + dataset_results = Dict() + scenes = self._get_scenes(dataset) + + for scene in tqdm(scenes, desc=f"{data} scenes", leave=False): + export_dir = self._export_dir(data, scene, posed=False) + result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + + # Check if result file exists and is valid + if not os.path.exists(result_path): + print(f"\n[ERROR] Result file not found: {result_path}") + print(f"[ERROR] CWD: {os.getcwd()}") + print(f"[ERROR] Please run inference first (remove --eval_only)") + continue + + try: + # Use saved GT meta (handles frame sampling correctly) + gt_meta = self._load_gt_meta(export_dir) + if gt_meta is not None: + result = self._compute_pose_with_gt(result_path, gt_meta) + else: + # Fallback to dataset GT (no sampling was done) + result = dataset.eval_pose(scene, result_path) + dataset_results[scene] = self._to_float_dict(result) + except Exception as e: + print(f"\n[ERROR] Failed to evaluate pose for {data}/{scene}: {e}") + print(f"[ERROR] File path: {os.path.abspath(result_path)}") + if self.debug: + import traceback + traceback.print_exc() + continue + + if not dataset_results: + print(f"[WARNING] No valid results for {data}") + continue + + dataset_results["mean"] = self._mean_of_dicts(dataset_results.values()) + out_path = os.path.join(self._metric_dir, f"{data}_pose.json") + self._dump_json(out_path, dataset_results) + yield data, dataset_results + + def _eval_reconstruction(self, mode: str) -> Iterable[tuple]: + """ + Compute reconstruction metrics for each dataset and scene. + + Args: + mode: "recon_unposed" or "recon_posed" + """ + assert mode in {"recon_unposed", "recon_posed"} + os.makedirs(self._metric_dir, exist_ok=True) + + posed_flag = mode == "recon_posed" + + # Filter out datasets that don't support reconstruction (e.g., dtu64) + recon_datas = [d for d in self.datas if d != "dtu64"] + + for data in tqdm(recon_datas, desc=f"Datasets ({mode} eval)"): + dataset = self.datasets[data] + dataset_results = Dict() + scenes = self._get_scenes(dataset) + + # Prepare paths for all scenes + scene_list = [] + result_paths = [] + fuse_paths = [] + for scene in scenes: + export_dir = self._export_dir(data, scene, posed=posed_flag) + result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + fuse_path = os.path.join(export_dir, "exports", "fuse", "pcd.ply") + scene_list.append(scene) + result_paths.append(result_path) + fuse_paths.append(fuse_path) + + # Parallel fusion (default 4 workers) + # DTU uses CUDA operations in fusion, which doesn't work well with ThreadPool + use_sequential = (data == "dtu") + parallel_execution( + scene_list, + result_paths, + fuse_paths, + action=lambda s, rp, fp: dataset.fuse3d(s, rp, fp, mode), + num_processes=self.num_fusion_workers, + print_progress=True, + desc=f"{data} fusion", + sequential=use_sequential, + ) + + # Sequential evaluation (fast, no need to parallelize) + for scene, fuse_path in zip(scene_list, fuse_paths): + # DTU supports CPU-based evaluation + if data == "dtu" and hasattr(dataset, "eval3d"): + result = dataset.eval3d(scene, fuse_path) + else: + result = dataset.eval3d(scene, fuse_path) + dataset_results[scene] = self._to_float_dict(result) + print(f" {mode} | {data} | {scene}: {result}") + + dataset_results["mean"] = self._mean_of_dicts(dataset_results.values()) + out_path = os.path.join(self._metric_dir, f"{data}_{mode}.json") + self._dump_json(out_path, dataset_results) + yield data, dataset_results + + # -------------------- Helpers -------------------- # + + def _save_gt_meta(self, export_dir: str, scene_data: Dict) -> None: + """ + Save GT extrinsics/intrinsics/image_files for evaluation. + + This is needed when frames are sampled, so eval_pose and fuse3d can use + the correct (sampled) GT instead of full dataset GT. + + Args: + export_dir: Export directory for the scene + scene_data: Sampled scene data + """ + meta_path = os.path.join(export_dir, "exports", "gt_meta.npz") + os.makedirs(os.path.dirname(meta_path), exist_ok=True) + np.savez_compressed( + meta_path, + extrinsics=scene_data.extrinsics, + intrinsics=scene_data.intrinsics, + image_files=np.array(scene_data.image_files, dtype=object), + ) + + def _load_gt_meta(self, export_dir: str) -> Dict: + """ + Load saved GT extrinsics/intrinsics for evaluation. + + Returns: + Dict with extrinsics and intrinsics, or None if not found + """ + meta_path = os.path.join(export_dir, "exports", "gt_meta.npz") + if os.path.exists(meta_path): + data = np.load(meta_path) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + }) + return None + + def _compute_pose_with_gt(self, result_path: str, gt_meta: Dict) -> TDict[str, float]: + """ + Compute pose metrics using saved GT meta (handles frame sampling). + + Args: + result_path: Path to npz with predicted extrinsics + gt_meta: Dict with GT extrinsics from saved meta + + Returns: + Dict with pose metrics + """ + from depth_anything_3.bench.dataset import _wait_for_file_ready + from depth_anything_3.bench.utils import compute_pose + from depth_anything_3.utils.geometry import as_homogeneous + + _wait_for_file_ready(result_path) + pred = np.load(result_path) + return compute_pose( + torch.from_numpy(as_homogeneous(pred["extrinsics"])), + torch.from_numpy(as_homogeneous(gt_meta["extrinsics"])), + ) + + def _sample_frames(self, scene_data: Dict, scene: str) -> Dict: + """ + Sample frames if scene has more than max_frames. + + Uses fixed random seed (42) for reproducibility. + + Args: + scene_data: Scene data dict with image_files, extrinsics, intrinsics, aux + scene: Scene name (for logging) + + Returns: + Sampled scene_data if num_frames > max_frames, otherwise original + """ + if self.max_frames <= 0: + return scene_data + + num_frames = len(scene_data.image_files) + if num_frames <= self.max_frames: + return scene_data + + # Sample with fixed seed for reproducibility + random.seed(42) + indices = list(range(num_frames)) + random.shuffle(indices) + sampled_indices = sorted(indices[:self.max_frames]) + + print(f" [Sampling] {scene}: {num_frames} -> {self.max_frames} frames") + + # Create new scene_data with sampled frames + sampled = Dict() + sampled.image_files = [scene_data.image_files[i] for i in sampled_indices] + sampled.extrinsics = scene_data.extrinsics[sampled_indices] + sampled.intrinsics = scene_data.intrinsics[sampled_indices] + + # Copy aux data, sampling lists if needed + sampled.aux = Dict() + for key, val in scene_data.aux.items(): + if isinstance(val, list) and len(val) == num_frames: + sampled.aux[key] = [val[i] for i in sampled_indices] + elif isinstance(val, np.ndarray) and len(val) == num_frames: + sampled.aux[key] = val[sampled_indices] + else: + sampled.aux[key] = val + + return sampled + + @property + def _metric_dir(self) -> str: + """Directory for storing metric JSON files.""" + return os.path.join(self.work_dir, "metric_results") + + def _export_dir(self, data: str, scene: str, posed: bool) -> str: + """ + Get export directory path. + + Structure: .../model_results/{data}/{scene}/{posed|unposed} + """ + suffix = "posed" if posed else "unposed" + export_dir = os.path.join(self.work_dir, "model_results", data, scene, suffix) + os.makedirs(export_dir, exist_ok=True) + return export_dir + + @staticmethod + def _to_float_dict(d: TDict[str, float]) -> dict: + """Convert numpy scalars to plain Python floats for JSON safety.""" + return {k: float(v) for k, v in d.items()} + + @staticmethod + def _mean_of_dicts(dicts: Iterable[dict]) -> dict: + """Compute elementwise mean across a list of homogeneous metric dicts.""" + dicts = list(dicts) + if not dicts: + return {} + keys = dicts[0].keys() + return {k: float(np.mean([d[k] for d in dicts]).item()) for k in keys} + + @staticmethod + def _dump_json(path: str, obj: dict, indent: int = 4) -> None: + """Write JSON with UTF-8 and pretty indentation.""" + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(obj, f, indent=indent, ensure_ascii=False) + + def _load_metrics(self) -> TDict[str, dict]: + """Load evaluation metrics from JSON files.""" + metrics = {} + metric_dir = self._metric_dir + + if not os.path.exists(metric_dir): + return metrics + + for filename in os.listdir(metric_dir): + if filename.endswith(".json"): + filepath = os.path.join(metric_dir, filename) + try: + with open(filepath, encoding="utf-8") as f: + data = json.load(f) + key = filename[:-5] # Remove .json extension + metrics[key] = data + except Exception as e: + print(f"Warning: Failed to read metrics file: {filename} - {e}") + + return metrics + + +# -------------------- CLI Entry Point -------------------- # + + +if __name__ == "__main__": + import sys + from omegaconf import OmegaConf + from depth_anything_3.cfg import load_config + + # Get default config path (relative to this file) + _default_config = os.path.join( + os.path.dirname(__file__), "configs", "eval_bench.yaml" + ) + + # Check for help flag first (we need to handle this before OmegaConf) + if "--help" in sys.argv or "-h" in sys.argv: + pass # Will handle after config loading + + # Set up argv for OmegaConf processing + argv = sys.argv[1:] + + # Check if user provides custom config + config_path = _default_config + if "--config" in argv: + config_idx = argv.index("--config") + if config_idx + 1 < len(argv): + config_path = argv[config_idx + 1] + # Remove --config and its value + argv = argv[:config_idx] + argv[config_idx + 2:] + + # Print help if requested + if "--help" in sys.argv or "-h" in sys.argv: + print(""" +DepthAnything3 Benchmark Evaluation + +Usage: + python -m depth_anything_3.bench.evaluator [OPTIONS] [KEY=VALUE ...] + +Configuration: + --config PATH Config YAML file (default: bench/configs/eval_bench.yaml) + +Config Overrides (using dotlist notation): + model.path=VALUE Model path or HuggingFace ID + workspace.work_dir=VALUE Working directory for outputs + eval.datasets=[dataset1,dataset2] Datasets to evaluate (eth3d,7scenes,scannetpp,hiroom,dtu,dtu64) + eval.modes=[mode1,mode2] Evaluation modes (pose,recon_unposed,recon_posed) + eval.scenes=[scene1,scene2] Specific scenes to evaluate (null=all) + eval.max_frames=VALUE Max frames per scene (-1=no limit, default: 100) + eval.ref_view_strategy=VALUE Reference view strategy (default: first) + eval.eval_only=VALUE Only run evaluation (skip inference) (true/false) + eval.print_only=VALUE Only print saved metrics (true/false) + inference.num_fusion_workers=VALUE Number of parallel workers (default: 4) + inference.debug=VALUE Enable debug mode (true/false) + +Special Flags: + --help, -h Show this help message + +Multi-GPU: + Use CUDA_VISIBLE_DEVICES to specify GPUs (auto-detected and distributed) + +Examples: + # Use default config + python -m depth_anything_3.bench.evaluator + + # Override model path + python -m depth_anything_3.bench.evaluator model.path=depth-anything/DA3-LARGE + + # Evaluate specific datasets and modes + python -m depth_anything_3.bench.evaluator \\ + eval.datasets=[eth3d,hiroom] \\ + eval.modes=[pose] + + # Use custom config with overrides + python -m depth_anything_3.bench.evaluator \\ + --config my_config.yaml \\ + model.path=/path/to/model \\ + eval.max_frames=50 + + # Multi-GPU inference (auto-distributed) + CUDA_VISIBLE_DEVICES=0,1,2,3 python -m depth_anything_3.bench.evaluator + + # Debug specific scenes + python -m depth_anything_3.bench.evaluator \\ + eval.datasets=[eth3d] \\ + eval.scenes=[courtyard] \\ + inference.debug=true + + # Only evaluate (skip inference) + python -m depth_anything_3.bench.evaluator eval.eval_only=true + + # Only print saved metrics + python -m depth_anything_3.bench.evaluator eval.print_only=true + + """) + sys.exit(0) + + # Load config with CLI overrides using OmegaConf dotlist + # Example: python evaluator.py model.path=/path/to/model eval.datasets=[eth3d,dtu] + config = load_config(config_path, argv=argv) + + # Extract config values + work_dir = config.workspace.work_dir + model_path = config.model.path + datasets = config.eval.datasets + modes = config.eval.modes + ref_view_strategy = config.eval.ref_view_strategy + scenes = config.eval.scenes + max_frames = config.eval.max_frames + eval_only = config.eval.eval_only + print_only = config.eval.print_only + debug = config.inference.debug + num_fusion_workers = config.inference.num_fusion_workers + + # GPU settings: parse from CLI dotlist args (gpu_id=X total_gpus=Y) + # These are passed by the main process when spawning workers + gpu_id = 0 + total_gpus = 1 + for arg in argv: + if arg.startswith("gpu_id="): + gpu_id = int(arg.split("=")[1]) + elif arg.startswith("total_gpus="): + total_gpus = int(arg.split("=")[1]) + + # Override dataset scenes if specified + if scenes: + print(f"[INFO] Running on specific scenes: {scenes}") + + evaluator = Evaluator( + work_dir=work_dir, + datas=datasets, + modes=modes, + ref_view_strategy=ref_view_strategy, + scenes=scenes, + debug=debug, + num_fusion_workers=num_fusion_workers, + max_frames=max_frames, + gpu_id=gpu_id, + total_gpus=total_gpus, + ) + + if print_only: + evaluator.print_metrics() + elif eval_only: + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + else: + # Parse CUDA_VISIBLE_DEVICES to get GPU list + # If not set, use all available GPUs + cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if cuda_devices is not None and cuda_devices.strip(): + gpu_list = [g.strip() for g in cuda_devices.split(",") if g.strip()] + else: + # CUDA_VISIBLE_DEVICES not set, use all available GPUs + num_available = torch.cuda.device_count() + gpu_list = [str(i) for i in range(num_available)] if num_available > 0 else ["0"] + + # Auto multi-GPU: if multiple GPUs and not a worker process + is_worker = os.environ.get("_DA3_WORKER") == "1" + + if len(gpu_list) > 1 and not is_worker: + # Launch worker processes + import subprocess + + num_gpus = len(gpu_list) + print(f"[INFO] Detected {num_gpus} GPUs: {gpu_list}") + print(f"[INFO] Launching {num_gpus} workers...") + + # Build base command + base_cmd = [sys.executable, "-m", "depth_anything_3.bench.evaluator"] + # Pass config via dotlist instead of CLI args + if config_path != _default_config: + base_cmd += ["--config", config_path] + base_cmd += [f"model.path={model_path}"] + base_cmd += [f"workspace.work_dir={work_dir}"] + base_cmd += [f"eval.datasets=[{','.join(datasets)}]"] + base_cmd += [f"eval.modes=[{','.join(modes)}]"] + if scenes: + base_cmd += [f"eval.scenes=[{','.join(scenes)}]"] + base_cmd += [f"eval.max_frames={max_frames}"] + base_cmd += [f"eval.ref_view_strategy={ref_view_strategy}"] + base_cmd += [f"inference.debug={str(debug).lower()}"] + base_cmd += [f"inference.num_fusion_workers={num_fusion_workers}"] + + # Launch workers + processes = [] + for idx, gpu_id in enumerate(gpu_list): + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = gpu_id + env["_DA3_WORKER"] = "1" # Mark as worker process + + cmd = base_cmd.copy() + # GPU-specific worker config + cmd += [f"gpu_id={idx}", f"total_gpus={num_gpus}"] + + print(f"[INFO] Starting worker {idx} on GPU {gpu_id}") + p = subprocess.Popen(cmd, env=env) + processes.append(p) + + # Wait for all workers + for p in processes: + p.wait() + + print(f"[INFO] All {num_gpus} workers completed") + + # Run evaluation after all inference is done + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + else: + # Single GPU or worker process + from depth_anything_3.api import DepthAnything3 + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + api = DepthAnything3.from_pretrained(model_path) + api = api.to(device) + + evaluator.infer(api, model_path=model_path) + + # Only run eval if single GPU mode (workers don't eval) + if not is_worker: + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + diff --git a/depth_anything_3/bench/print_metrics.py b/depth_anything_3/bench/print_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..9964f3bf74fdf91a15ea0ccedb30995ae79b7281 --- /dev/null +++ b/depth_anything_3/bench/print_metrics.py @@ -0,0 +1,618 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Beautiful metrics printing utilities for benchmark evaluation. + +Provides colorized, well-formatted tabular output for evaluation results. +Supports highlighting best/worst values and grouping by dataset/mode. +""" + +import argparse +import json +import os +import re +from typing import Dict as TDict, List, Optional + + +# ANSI color codes for terminal output +class Colors: + """ANSI escape codes for terminal colors.""" + + RESET = "\033[0m" + BOLD = "\033[1m" + RED = "\033[31m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + BLUE = "\033[34m" + MAGENTA = "\033[35m" + CYAN = "\033[36m" + WHITE = "\033[37m" + + # Bold variants + BOLD_RED = "\033[1;31m" + BOLD_GREEN = "\033[1;32m" + BOLD_YELLOW = "\033[1;33m" + BOLD_BLUE = "\033[1;34m" + BOLD_MAGENTA = "\033[1;35m" + BOLD_CYAN = "\033[1;36m" + + # Background + BG_DARK = "\033[48;5;236m" + + +def strip_ansi(text: str) -> str: + """Remove ANSI escape sequences from string for length calculation.""" + ansi_escape = re.compile(r"\x1b\[[0-9;]*m") + return ansi_escape.sub("", text) + + +def colorize_value( + value: str, + is_best: bool = False, + is_worst: bool = False, + lower_is_better: bool = False, +) -> str: + """ + Apply color to a metric value based on whether it's best/worst. + + Args: + value: String representation of the value + is_best: Whether this is the best value in its column + is_worst: Whether this is the worst value in its column + lower_is_better: If True, lower values are better (e.g., error metrics) + + Returns: + Colorized string + """ + if lower_is_better: + # For metrics like error/distance, lower is better + if is_best: + return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}" + elif is_worst: + return f"{Colors.BOLD_RED}{value}{Colors.RESET}" + else: + # For metrics like accuracy/AUC, higher is better + if is_best: + return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}" + elif is_worst: + return f"{Colors.BOLD_RED}{value}{Colors.RESET}" + return value + + +class MetricsPrinter: + """ + Beautiful tabular metrics printer with color support. + + Features: + - Colorized best/worst values + - Grouped by dataset and evaluation mode + - Automatic column width calculation + - Support for multiple input directories comparison + """ + + # Metrics where lower values are better + LOWER_IS_BETTER = {"comp", "acc", "overall", "error", "loss", "rmse", "mae"} + + def __init__(self, use_color: bool = True): + """ + Initialize the printer. + + Args: + use_color: Whether to use ANSI colors in output + """ + self.use_color = use_color + + def print_results(self, metrics: TDict[str, dict], summary_only: bool = True) -> None: + """ + Print evaluation metrics in a beautiful tabular format. + + Args: + metrics: Dictionary mapping "dataset_mode" to metric results + summary_only: If True, only print summary table. If False, print per-dataset details too. + """ + if not metrics: + print(f"\n{Colors.BOLD_RED}❌ No evaluation metrics found{Colors.RESET}") + return + + if not summary_only: + self._print_header() + grouped = self._group_by_dataset(metrics) + + for dataset, modes_data in grouped.items(): + self._print_dataset_section(dataset, modes_data) + + # Print summary table with average metrics across datasets + self._print_summary(metrics) + + self._print_footer() + + def print_comparison( + self, + metrics_list: List[TDict[str, dict]], + labels: List[str], + ) -> None: + """ + Print comparison table for multiple evaluation runs. + + Args: + metrics_list: List of metrics dictionaries + labels: Labels for each metrics dictionary + """ + if not metrics_list or not all(metrics_list): + print(f"\n{Colors.BOLD_RED}❌ No metrics to compare{Colors.RESET}") + return + + # Collect all datasets and modes + all_keys = set() + for metrics in metrics_list: + all_keys.update(metrics.keys()) + + self._print_header("COMPARISON") + + for key in sorted(all_keys): + parts = key.rsplit("_", 1) + if len(parts) == 2: + dataset, mode = parts[0], parts[1] + else: + dataset, mode = key, "unknown" + + print(f"\n{Colors.BOLD_CYAN}📊 {dataset.upper()} - {mode.upper()}{Colors.RESET}") + print("-" * 100) + + # Collect metrics from all runs + all_metric_names = set() + for metrics in metrics_list: + if key in metrics and "mean" in metrics[key]: + all_metric_names.update(metrics[key]["mean"].keys()) + + if not all_metric_names: + continue + + # Build comparison table + metric_width = max(15, max(len(m) for m in all_metric_names) + 2) + label_width = max(15, max(len(l) for l in labels) + 2) + + # Header + header = f"{'Metric':<{metric_width}}" + for label in labels: + header += f"{label:<{label_width}}" + print(header) + print("-" * len(strip_ansi(header))) + + # Collect values for highlighting + for metric_name in sorted(all_metric_names): + values = [] + for metrics in metrics_list: + if key in metrics and "mean" in metrics[key]: + val = metrics[key]["mean"].get(metric_name) + values.append(val if val is not None else float("nan")) + else: + values.append(float("nan")) + + # Find best/worst + valid_values = [v for v in values if not (v != v)] # Filter NaN + if valid_values: + lower_better = any( + lb in metric_name.lower() for lb in self.LOWER_IS_BETTER + ) + best_val = min(valid_values) if lower_better else max(valid_values) + worst_val = max(valid_values) if lower_better else min(valid_values) + else: + best_val = worst_val = None + + # Print row + row = f"{metric_name:<{metric_width}}" + for val in values: + if val != val: # NaN check + val_str = "N/A" + else: + val_str = f"{val:.4f}" + if self.use_color and len(valid_values) > 1: + lower_better = any( + lb in metric_name.lower() for lb in self.LOWER_IS_BETTER + ) + is_best = abs(val - best_val) < 1e-8 if best_val else False + is_worst = abs(val - worst_val) < 1e-8 if worst_val else False + val_str_padded = f"{val_str:<{label_width}}" + val_str = colorize_value( + val_str_padded, is_best, is_worst, lower_better + ) + row += val_str + continue + row += f"{val_str:<{label_width}}" + print(row) + + self._print_footer() + + def _print_header(self, title: str = "EVALUATION RESULTS") -> None: + """Print report header.""" + width = 100 + print() + print("=" * width) + print(f"{Colors.BOLD_CYAN}📊 DEPTH ANYTHING 3 {title}{Colors.RESET}") + print("=" * width) + + def _print_footer(self) -> None: + """Print report footer.""" + width = 100 + print() + print("=" * width) + print(f"{Colors.BOLD_GREEN}✅ Evaluation Complete{Colors.RESET}") + print("=" * width) + print() + + def _group_by_dataset(self, metrics: TDict[str, dict]) -> TDict[str, dict]: + """Group metrics by dataset.""" + grouped = {} + for key, data in metrics.items(): + if not isinstance(data, dict) or "mean" not in data: + continue + # Parse key format: "dataset_mode" (e.g., "dtu_recon_unposed") + parts = key.split("_", 1) + if len(parts) == 2: + dataset, mode = parts + if dataset not in grouped: + grouped[dataset] = {} + grouped[dataset][mode] = data + return grouped + + def _print_dataset_section(self, dataset: str, modes_data: TDict[str, dict]) -> None: + """Print metrics section for a single dataset.""" + print(f"\n{Colors.BOLD_MAGENTA}🔍 {dataset.upper()}{Colors.RESET}") + print("-" * 100) + + # Collect all unique metrics across all modes + all_metrics = set() + for mode_data in modes_data.values(): + all_metrics.update(mode_data["mean"].keys()) + all_metrics = sorted(list(all_metrics)) + + if not all_metrics: + print(" No metrics available") + return + + # Calculate column widths + metric_width = max(18, max(len(m) for m in all_metrics) + 2) + mode_width = 18 + modes = list(modes_data.keys()) + + # Print header + header = f"{'Metric':<{metric_width}}" + for mode in modes: + header += f"{mode.upper():<{mode_width}}" + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(header)) + + # Print each metric row + for metric in all_metrics: + row = f"{metric:<{metric_width}}" + + # Collect values for this metric across modes + values = [] + for mode in modes: + if metric in modes_data[mode]["mean"]: + values.append(modes_data[mode]["mean"][metric]) + else: + values.append(None) + + # Find best/worst values + valid_values = [v for v in values if v is not None] + if valid_values: + lower_better = any(lb in metric.lower() for lb in self.LOWER_IS_BETTER) + best_val = min(valid_values) if lower_better else max(valid_values) + worst_val = max(valid_values) if lower_better else min(valid_values) + else: + best_val = worst_val = None + + # Format each value + for val in values: + if val is None: + row += f"{'N/A':<{mode_width}}" + else: + val_str = f"{val:.4f}" + if self.use_color and len(valid_values) > 1: + is_best = abs(val - best_val) < 1e-8 if best_val else False + is_worst = abs(val - worst_val) < 1e-8 if worst_val else False + lower_better = any( + lb in metric.lower() for lb in self.LOWER_IS_BETTER + ) + # Pad before colorizing to maintain alignment + val_str_padded = f"{val_str:<{mode_width}}" + row += colorize_value( + val_str_padded, is_best, is_worst, lower_better + ) + else: + row += f"{val_str:<{mode_width}}" + print(row) + + # Show scene counts + scene_info = [] + for mode, mode_data in modes_data.items(): + scene_count = len([k for k in mode_data.keys() if k != "mean"]) + scene_info.append(f"{mode}: {scene_count} scenes") + print(f"\n{Colors.CYAN}📈 {' | '.join(scene_info)}{Colors.RESET}") + + def _print_summary(self, metrics: TDict[str, dict]) -> None: + """ + Print summary table with key metrics across all datasets. + + Format: One row per metric, datasets as columns. + Order: HiRoom, ETH3D, DTU, 7Scenes, ScanNet++, (DTU-64 for pose only) + """ + print(f"\n{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}") + print(f"{Colors.BOLD_CYAN}📊 SUMMARY{Colors.RESET}") + print(f"{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}") + + # Dataset display order and names + DATASET_ORDER = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp", "dtu64"] + DATASET_DISPLAY = { + "hiroom": "HiRoom", + "eth3d": "ETH3D", + "dtu": "DTU", + "7scenes": "7Scenes", + "scannetpp": "ScanNet++", + "dtu64": "DTU-64", + } + + # Collect all metrics into a structured dict + # metric_data[dataset][mode] = {"Auc_3": x, "Auc_30": x, "fscore": x, "overall": x} + metric_data = {} + for key, data in metrics.items(): + if not isinstance(data, dict) or "mean" not in data: + continue + parts = key.split("_", 1) + if len(parts) != 2: + continue + dataset, mode = parts + dataset_lower = dataset.lower() + if dataset_lower not in metric_data: + metric_data[dataset_lower] = {} + metric_data[dataset_lower][mode] = data["mean"] + + col_width = 12 + + def fmt_val(val): + """Format value or return N/A.""" + if val is None: + return "N/A" + return f"{val:.4f}" + + def get_metric(dataset, mode, metric_name): + """Get metric value or None.""" + if dataset not in metric_data: + return None + if mode not in metric_data[dataset]: + return None + return metric_data[dataset][mode].get(metric_name) + + # ============ POSE METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🎯 POSE ESTIMATION{Colors.RESET}") + + # Pose: show all datasets except DTU (keep DTU-64 only) + # Order: HiRoom, ETH3D, DTU-64, 7Scenes, ScanNet++ + pose_datasets = ["hiroom", "eth3d", "dtu64", "7scenes", "scannetpp"] + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg':<{col_width}}" + for ds in pose_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # Helper to get metric with fallback names + def get_pose_metric(dataset, metric_name): + """Get pose metric with fallback for different naming conventions.""" + # Try different naming conventions + names = { + "Auc3": ["Auc_3", "auc03", "auc_3", "AUC_3", "Auc3", "auc3"], + "Auc30": ["Auc_30", "auc30", "auc_30", "AUC_30", "Auc30"], + } + for name in names.get(metric_name, [metric_name]): + val = get_metric(dataset, "pose", name) + if val is not None: + return val + return None + + # Auc3 row + values = [] + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc3") + if val is not None: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Auc3':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc3") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Auc30 row + values = [] + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc30") + if val is not None: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Auc30':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc30") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # ============ RECON_UNPOSED METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_UNPOSED (Pred Pose){Colors.RESET}") + + # For recon, exclude dtu64 from columns + recon_datasets = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp"] + avg_datasets = ["hiroom", "eth3d", "7scenes", "scannetpp"] # Exclude DTU from avg + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg*':<{col_width}}" + for ds in recon_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # F-score row (only metric for avg) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "fscore") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "fscore") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Overall row (avg over 4 datasets excluding DTU) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "overall") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "overall") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # ============ RECON_POSED METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_POSED (GT Pose){Colors.RESET}") + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg*':<{col_width}}" + for ds in recon_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # F-score row (only metric for avg) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "fscore") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "fscore") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Overall row (avg over 4 datasets excluding DTU) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "overall") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "overall") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + print(f"\n{Colors.CYAN}* Avg F-score / Overall = average over HiRoom, ETH3D, 7Scenes, ScanNet++ (4 datasets){Colors.RESET}") + + +def load_metrics_from_dir(metric_dir: str) -> TDict[str, dict]: + """ + Load all metrics JSON files from a directory. + + Args: + metric_dir: Path to directory containing metric JSON files + + Returns: + Dictionary mapping filename (without .json) to metric data + """ + metrics = {} + if not os.path.exists(metric_dir): + return metrics + + for filename in os.listdir(metric_dir): + if filename.endswith(".json"): + filepath = os.path.join(metric_dir, filename) + try: + with open(filepath, encoding="utf-8") as f: + content = f.read() + # Handle trailing commas in JSON + content = re.sub(r",\s*([\]\}])", r"\1", content) + data = json.loads(content) + key = filename[:-5] + metrics[key] = data + except Exception as e: + print(f"Warning: Failed to load {filename}: {e}") + + return metrics + + +def main(): + """Command-line interface for metrics printing.""" + parser = argparse.ArgumentParser( + description="Print DepthAnything3 benchmark evaluation metrics." + ) + parser.add_argument( + "--input_dir", + type=str, + default="./eval_workspace/metric_results", + help="Directory containing metric JSON files (comma-separated for comparison)", + ) + parser.add_argument( + "--no_color", + action="store_true", + help="Disable colored output", + ) + parser.add_argument( + "--key", + type=str, + default=None, + help="Specific metric key to highlight", + ) + args = parser.parse_args() + + # Support multiple directories for comparison + input_dirs = [d.strip() for d in args.input_dir.split(",") if d.strip()] + + printer = MetricsPrinter(use_color=not args.no_color) + + if len(input_dirs) == 1: + # Single directory - simple print + metrics = load_metrics_from_dir(input_dirs[0]) + printer.print_results(metrics) + else: + # Multiple directories - comparison mode + metrics_list = [] + labels = [] + for d in input_dirs: + metrics = load_metrics_from_dir(d) + if metrics: + metrics_list.append(metrics) + labels.append(os.path.basename(d.rstrip("/"))) + + if metrics_list: + printer.print_comparison(metrics_list, labels) + else: + print("No metrics found in specified directories") + + +if __name__ == "__main__": + main() + diff --git a/depth_anything_3/bench/registries.py b/depth_anything_3/bench/registries.py new file mode 100644 index 0000000000000000000000000000000000000000..0744661e3dcf7664fc70df5956add84f472109c4 --- /dev/null +++ b/depth_anything_3/bench/registries.py @@ -0,0 +1,85 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Auto-loading registry system for benchmark datasets. + +This module provides registry classes that automatically discover and import +dataset implementations from the datasets subpackage on first access. +""" + +import importlib +import pkgutil +import threading + +from depth_anything_3.utils.registry import Registry + +__all__ = ["METRIC_REGISTRY", "MONO_REGISTRY", "MV_REGISTRY", "NVS_REGISTRY"] + +# ---- Lazy import: Only scan and import all datasets submodules on first registry access ---- +_loaded = False +_lock = threading.Lock() + + +def _import_all_datasets_once(): + """ + Scan and import all .py submodules under depth_anything_3.bench.datasets + (skip files/packages starting with underscore), to trigger @REGISTRY.register(...) in each module. + """ + global _loaded + if _loaded: + return + + with _lock: + if _loaded: + return + + pkg_name = "depth_anything_3.bench.datasets" + pkg = importlib.import_module(pkg_name) + pkg_paths = list(getattr(pkg, "__path__", [])) + + for finder, name, ispkg in pkgutil.walk_packages(pkg_paths, prefix=pkg_name + "."): + base = name.rsplit(".", 1)[-1] + if base.startswith("_"): + continue + try: + importlib.import_module(name) + except Exception as e: + print(f"[datasets auto-import] Failed to import {name}: {e}") + + _loaded = True + + +class AutoRegistry(Registry): + """Registry that ensures all datasets are auto-discovered and imported on first use.""" + + def get(self, name): + _import_all_datasets_once() + return super().get(name) + + def all(self): + _import_all_datasets_once() + return super().all() + + def has(self, name): + _import_all_datasets_once() + return name in self._map + + +# Four auto-lazy registry instances for different evaluation types +METRIC_REGISTRY = AutoRegistry() # For metric depth evaluation +MONO_REGISTRY = AutoRegistry() # For monocular depth evaluation +MV_REGISTRY = AutoRegistry() # For multi-view evaluation +NVS_REGISTRY = AutoRegistry() # For novel view synthesis evaluation + diff --git a/depth_anything_3/bench/utils.py b/depth_anything_3/bench/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ca654773122c247e15d4a99a0beb5d0e80ac8e20 --- /dev/null +++ b/depth_anything_3/bench/utils.py @@ -0,0 +1,525 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utility functions for benchmark evaluation. + +Contains: +- Pose evaluation metrics (AUC) and helper functions +- 3D reconstruction evaluation metrics (Acc/Comp/F-score) +- Geometry utilities (quaternion conversion, etc.) +""" + +from typing import Dict as TDict, Optional, Tuple, Union + +import numpy as np +import open3d as o3d +import torch +from addict import Dict +from scipy.spatial import KDTree + +from depth_anything_3.utils.geometry import mat_to_quat + + +# ============================================================================= +# Geometry Utilities +# ============================================================================= + + +def quat2rotmat(qvec: list) -> np.ndarray: + """ + Convert quaternion (WXYZ order) to rotation matrix. + + Args: + qvec: Quaternion as [w, x, y, z] + + Returns: + 3x3 rotation matrix + """ + rotmat = np.array( + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ] + ) + rotmat = rotmat.reshape(3, 3) + return rotmat + + +# ============================================================================= +# 3D Reconstruction Evaluation +# ============================================================================= + + +def nn_correspondance(verts1: np.ndarray, verts2: np.ndarray) -> np.ndarray: + """ + Compute nearest neighbor distances from verts2 to verts1 using KDTree. + + Args: + verts1: Reference point cloud [N, 3] + verts2: Query point cloud [M, 3] + + Returns: + Distance array [M,] - distance from each point in verts2 to nearest in verts1 + """ + if len(verts1) == 0 or len(verts2) == 0: + return np.array([]) + + kdtree = KDTree(verts1) + distances, _ = kdtree.query(verts2) + return distances.reshape(-1) + + +def evaluate_3d_reconstruction( + pcd_pred: Union[o3d.geometry.PointCloud, np.ndarray], + pcd_trgt: Union[o3d.geometry.PointCloud, np.ndarray], + threshold: float = 0.05, + down_sample: Optional[float] = None, +) -> TDict[str, float]: + """ + Evaluate 3D reconstruction quality using standard metrics. + + This function computes: + - Accuracy: Mean distance from predicted points to GT surface + - Completeness: Mean distance from GT points to predicted surface + - Overall: Average of accuracy and completeness + - Precision: Fraction of predicted points within threshold of GT + - Recall: Fraction of GT points within threshold of prediction + - F-score: Harmonic mean of precision and recall + + Args: + pcd_pred: Predicted point cloud (Open3D or numpy array) + pcd_trgt: Ground truth point cloud (Open3D or numpy array) + threshold: Distance threshold for precision/recall (meters) + down_sample: Voxel size for downsampling (None to skip) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + # Convert to Open3D if needed + if isinstance(pcd_pred, np.ndarray): + pcd_pred_o3d = o3d.geometry.PointCloud() + pcd_pred_o3d.points = o3d.utility.Vector3dVector(pcd_pred) + pcd_pred = pcd_pred_o3d + if isinstance(pcd_trgt, np.ndarray): + pcd_trgt_o3d = o3d.geometry.PointCloud() + pcd_trgt_o3d.points = o3d.utility.Vector3dVector(pcd_trgt) + pcd_trgt = pcd_trgt_o3d + + # Downsample if requested + if down_sample is not None and down_sample > 0: + pcd_pred = pcd_pred.voxel_down_sample(down_sample) + pcd_trgt = pcd_trgt.voxel_down_sample(down_sample) + + verts_pred = np.asarray(pcd_pred.points) + verts_trgt = np.asarray(pcd_trgt.points) + + # Handle empty point clouds + if len(verts_pred) == 0 or len(verts_trgt) == 0: + return { + "acc": float("inf"), + "comp": float("inf"), + "overall": float("inf"), + "precision": 0.0, + "recall": 0.0, + "fscore": 0.0, + } + + # Compute distances + dist_pred_to_gt = nn_correspondance(verts_trgt, verts_pred) # Accuracy + dist_gt_to_pred = nn_correspondance(verts_pred, verts_trgt) # Completeness + + # Compute metrics + accuracy = float(np.mean(dist_pred_to_gt)) + completeness = float(np.mean(dist_gt_to_pred)) + overall = (accuracy + completeness) / 2 + + precision = float(np.mean((dist_pred_to_gt < threshold).astype(float))) + recall = float(np.mean((dist_gt_to_pred < threshold).astype(float))) + + if precision + recall > 0: + fscore = 2 * precision * recall / (precision + recall) + else: + fscore = 0.0 + + return { + "acc": accuracy, + "comp": completeness, + "overall": overall, + "precision": precision, + "recall": recall, + "fscore": fscore, + } + + +def create_tsdf_volume( + voxel_length: float = 4.0 / 512.0, + sdf_trunc: float = 0.04, + color_type: str = "RGB8", +) -> o3d.pipelines.integration.ScalableTSDFVolume: + """ + Create a scalable TSDF volume for depth fusion. + + Args: + voxel_length: Size of each voxel + sdf_trunc: Truncation distance for SDF + color_type: Color integration type ("RGB8" or "Gray32") + + Returns: + Initialized ScalableTSDFVolume + """ + if color_type == "RGB8": + color_enum = o3d.pipelines.integration.TSDFVolumeColorType.RGB8 + else: + color_enum = o3d.pipelines.integration.TSDFVolumeColorType.Gray32 + + volume = o3d.pipelines.integration.ScalableTSDFVolume( + voxel_length=voxel_length, + sdf_trunc=sdf_trunc, + color_type=color_enum, + ) + return volume + + +def fuse_depth_to_tsdf( + volume: o3d.pipelines.integration.ScalableTSDFVolume, + depths: np.ndarray, + images: np.ndarray, + intrinsics: np.ndarray, + extrinsics: np.ndarray, + max_depth: float = 10.0, +) -> o3d.geometry.TriangleMesh: + """ + Fuse multiple depth maps into TSDF volume and extract mesh. + + Args: + volume: TSDF volume to integrate into + depths: Depth maps [N, H, W] + images: RGB images [N, H, W, 3] + intrinsics: Camera intrinsics [N, 3, 3] + extrinsics: Camera extrinsics (world-to-camera) [N, 4, 4] + max_depth: Maximum depth for truncation + + Returns: + Extracted triangle mesh + """ + for i in range(len(depths)): + depth = depths[i] + image = images[i] + ixt = intrinsics[i] + ext = extrinsics[i] + + h, w = depth.shape[:2] + + # Create RGBD image + depth_o3d = o3d.geometry.Image(depth.astype(np.float32)) + color_o3d = o3d.geometry.Image(image.astype(np.uint8)) + rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth( + color_o3d, + depth_o3d, + depth_trunc=max_depth, + convert_rgb_to_intensity=False, + depth_scale=1.0, + ) + + # Create camera intrinsics + ixt_o3d = o3d.camera.PinholeCameraIntrinsic( + w, h, ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2] + ) + + # Integrate into volume + volume.integrate(rgbd, ixt_o3d, ext) + + # Extract mesh + mesh = volume.extract_triangle_mesh() + return mesh + + +def sample_points_from_mesh( + mesh: o3d.geometry.TriangleMesh, + num_points: int = 1000000, +) -> o3d.geometry.PointCloud: + """ + Uniformly sample points from a triangle mesh. + + Args: + mesh: Input triangle mesh + num_points: Number of points to sample + + Returns: + Sampled point cloud + """ + try: + pcd = mesh.sample_points_uniformly(number_of_points=num_points) + # Clamp colors to valid range [0, 1] for Open3D PLY export + if pcd.has_colors(): + colors = np.asarray(pcd.colors) + colors = np.clip(colors, 0.0, 1.0) + pcd.colors = o3d.utility.Vector3dVector(colors) + except Exception: + # Fallback: create random points if mesh is invalid (with fixed seed for reproducibility) + rng = np.random.default_rng(seed=42) + points = rng.uniform(-1, 1, size=(num_points, 3)) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + return pcd + + +# ============================================================================= +# Pose Evaluation +# ============================================================================= + + +def build_pair_index(N: int, B: int = 1): + """ + Build indices for all possible pairs of frames. + + Args: + N: Number of frames + B: Batch size + + Returns: + i1, i2: Indices for all possible pairs + """ + i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1) + i1, i2 = ((i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_]) + return i1, i2 + + +def compute_pose(pred_se3: torch.Tensor, gt_se3: torch.Tensor) -> Dict: + """ + Compute pose estimation metrics between predicted and ground truth trajectories. + + Args: + pred_se3: Predicted SE(3) transformations [N, 4, 4] + gt_se3: Ground truth SE(3) transformations [N, 4, 4] + + Returns: + Dict with AUC metrics at different thresholds (auc30, auc15, auc05, auc03) + """ + pred_se3 = align_to_first_camera(pred_se3) + gt_se3 = align_to_first_camera(gt_se3) + + rel_rangle_deg, rel_tangle_deg = se3_to_relative_pose_error(pred_se3, gt_se3, len(pred_se3)) + rError = rel_rangle_deg.cpu().numpy() + tError = rel_tangle_deg.cpu().numpy() + + output = Dict() + output.auc30, _ = calculate_auc_np(rError, tError, max_threshold=30) + output.auc15, _ = calculate_auc_np(rError, tError, max_threshold=15) + output.auc05, _ = calculate_auc_np(rError, tError, max_threshold=5) + output.auc03, _ = calculate_auc_np(rError, tError, max_threshold=3) + return output + + +def align_to_first_camera(camera_poses: torch.Tensor) -> torch.Tensor: + """ + Align all camera poses to the first camera's coordinate frame. + + Args: + camera_poses: Camera poses as SE3 transformations [N, 4, 4] + + Returns: + Aligned camera poses [N, 4, 4] + """ + first_cam_extrinsic_inv = closed_form_inverse_se3(camera_poses[0][None]) + aligned_poses = torch.matmul(camera_poses, first_cam_extrinsic_inv) + return aligned_poses + + +def rotation_angle( + rot_gt: torch.Tensor, rot_pred: torch.Tensor, batch_size: int = None, eps: float = 1e-15 +) -> torch.Tensor: + """ + Calculate rotation angle error between ground truth and predicted rotations. + + Args: + rot_gt: Ground truth rotation matrices + rot_pred: Predicted rotation matrices + batch_size: Batch size for reshaping the result + eps: Small value to avoid numerical issues + + Returns: + Rotation angle error in degrees + """ + q_pred = mat_to_quat(rot_pred) + q_gt = mat_to_quat(rot_gt) + + loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps) + err_q = torch.arccos(1 - 2 * loss_q) + + rel_rangle_deg = err_q * 180 / np.pi + + if batch_size is not None: + rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1) + + return rel_rangle_deg + + +def translation_angle( + tvec_gt: torch.Tensor, + tvec_pred: torch.Tensor, + batch_size: int = None, + ambiguity: bool = True, +) -> torch.Tensor: + """ + Calculate translation angle error between ground truth and predicted translations. + + Args: + tvec_gt: Ground truth translation vectors + tvec_pred: Predicted translation vectors + batch_size: Batch size for reshaping the result + ambiguity: Whether to handle direction ambiguity + + Returns: + Translation angle error in degrees + """ + rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred) + rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi + + if ambiguity: + rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs()) + + if batch_size is not None: + rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1) + + return rel_tangle_deg + + +def compare_translation_by_angle( + t_gt: torch.Tensor, t: torch.Tensor, eps: float = 1e-15, default_err: float = 1e6 +) -> torch.Tensor: + """ + Normalize the translation vectors and compute the angle between them. + + Args: + t_gt: Ground truth translation vectors + t: Predicted translation vectors + eps: Small value to avoid division by zero + default_err: Default error value for invalid cases + + Returns: + Angular error between translation vectors in radians + """ + t_norm = torch.norm(t, dim=1, keepdim=True) + t = t / (t_norm + eps) + + t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True) + t_gt = t_gt / (t_gt_norm + eps) + + loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps) + err_t = torch.acos(torch.sqrt(1 - loss_t)) + + err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err + return err_t + + +def calculate_auc_np( + r_error: np.ndarray, t_error: np.ndarray, max_threshold: int = 30 +) -> tuple: + """ + Calculate the Area Under the Curve (AUC) for the given error arrays. + + Args: + r_error: Rotation error values in degrees + t_error: Translation error values in degrees + max_threshold: Maximum threshold value for binning + + Returns: + Tuple of (AUC value, normalized histogram) + """ + error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1) + max_errors = np.max(error_matrix, axis=1) + bins = np.arange(max_threshold + 1) + histogram, _ = np.histogram(max_errors, bins=bins) + num_pairs = float(len(max_errors)) + normalized_histogram = histogram.astype(float) / num_pairs + return np.mean(np.cumsum(normalized_histogram)), normalized_histogram + + +def se3_to_relative_pose_error( + pred_se3: torch.Tensor, gt_se3: torch.Tensor, num_frames: int +) -> tuple: + """ + Compute rotation and translation errors between predicted and ground truth poses. + + Args: + pred_se3: Predicted SE(3) transformations + gt_se3: Ground truth SE(3) transformations + num_frames: Number of frames + + Returns: + Tuple of (rotation angle errors, translation angle errors) in degrees + """ + pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames) + + # Compute relative camera poses between pairs using closed-form inverse + relative_pose_gt = closed_form_inverse_se3(gt_se3[pair_idx_i1]).bmm(gt_se3[pair_idx_i2]) + relative_pose_pred = closed_form_inverse_se3(pred_se3[pair_idx_i1]).bmm(pred_se3[pair_idx_i2]) + + # Compute the difference in rotation and translation + rel_rangle_deg = rotation_angle(relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3]) + rel_tangle_deg = translation_angle(relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3]) + + return rel_rangle_deg, rel_tangle_deg + + +def closed_form_inverse_se3( + se3: torch.Tensor, R: torch.Tensor = None, T: torch.Tensor = None +) -> torch.Tensor: + """ + Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch. + + Uses closed-form solution instead of torch.inverse() for numerical stability. + + Args: + se3: Nx4x4 or Nx3x4 tensor of SE3 matrices + R: Optional Nx3x3 rotation matrices + T: Optional Nx3x1 translation vectors + + Returns: + Inverted SE3 matrices with same shape as input + """ + is_numpy = isinstance(se3, np.ndarray) + + if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4): + raise ValueError(f"se3 must be of shape (N,4,4), got {se3.shape}.") + + if R is None: + R = se3[:, :3, :3] + if T is None: + T = se3[:, :3, 3:] + + if is_numpy: + R_transposed = np.transpose(R, (0, 2, 1)) + top_right = -np.matmul(R_transposed, T) + inverted_matrix = np.tile(np.eye(4), (len(R), 1, 1)) + else: + R_transposed = R.transpose(1, 2) + top_right = -torch.bmm(R_transposed, T) + inverted_matrix = torch.eye(4, 4)[None].repeat(len(R), 1, 1) + inverted_matrix = inverted_matrix.to(R.dtype).to(R.device) + + inverted_matrix[:, :3, :3] = R_transposed + inverted_matrix[:, :3, 3:] = top_right + + return inverted_matrix + diff --git a/depth_anything_3/cfg.py b/depth_anything_3/cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..4607ff8c0983b67e6ccd2d80b78a163c4e487250 --- /dev/null +++ b/depth_anything_3/cfg.py @@ -0,0 +1,144 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Configuration utility functions +""" + +import importlib +from pathlib import Path +from typing import Any, Callable, List, Union +from omegaconf import DictConfig, ListConfig, OmegaConf + +try: + OmegaConf.register_new_resolver("eval", eval) +except Exception as e: + # if eval is not available, we can just pass + print(f"Error registering eval resolver: {e}") + + +def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]: + """ + Load a configuration. Will resolve inheritance. + Supports both file paths and module paths (e.g., depth_anything_3.configs.giant). + """ + # Check if path is a module path (contains dots but no slashes and doesn't end with .yaml) + if "." in path and "/" not in path and not path.endswith(".yaml"): + # It's a module path, load from package resources + path_parts = path.split(".")[1:] + config_path = Path(__file__).resolve().parent + for part in path_parts: + config_path = config_path.joinpath(part) + config_path = config_path.with_suffix(".yaml") + config = OmegaConf.load(str(config_path)) + else: + # It's a file path (absolute, relative, or with .yaml extension) + config = OmegaConf.load(path) + + if argv is not None: + config_argv = OmegaConf.from_dotlist(argv) + config = OmegaConf.merge(config, config_argv) + config = resolve_recursive(config, resolve_inheritance) + return config + + +def resolve_recursive( + config: Any, + resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]], +) -> Any: + config = resolver(config) + if isinstance(config, DictConfig): + for k in config.keys(): + v = config.get(k) + if isinstance(v, (DictConfig, ListConfig)): + config[k] = resolve_recursive(v, resolver) + if isinstance(config, ListConfig): + for i in range(len(config)): + v = config.get(i) + if isinstance(v, (DictConfig, ListConfig)): + config[i] = resolve_recursive(v, resolver) + return config + + +def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any: + """ + Recursively resolve inheritance if the config contains: + __inherit__: path/to/parent.yaml or a ListConfig of such paths. + """ + if isinstance(config, DictConfig): + inherit = config.pop("__inherit__", None) + + if inherit: + inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit] + + parent_config = None + for parent_path in inherit_list: + assert isinstance(parent_path, str) + parent_config = ( + load_config(parent_path) + if parent_config is None + else OmegaConf.merge(parent_config, load_config(parent_path)) + ) + + if len(config.keys()) > 0: + config = OmegaConf.merge(parent_config, config) + else: + config = parent_config + return config + + +def import_item(path: str, name: str) -> Any: + """ + Import a python item. Example: import_item("path.to.file", "MyClass") -> MyClass + """ + return getattr(importlib.import_module(path), name) + + +def create_object(config: DictConfig) -> Any: + """ + Create an object from config. + The config is expected to contains the following: + __object__: + path: path.to.module + name: MyClass + args: as_config | as_params (default to as_config) + """ + config = DictConfig(config) + item = import_item( + path=config.__object__.path, + name=config.__object__.name, + ) + args = config.__object__.get("args", "as_config") + if args == "as_config": + return item(config) + if args == "as_params": + config = OmegaConf.to_object(config) + config.pop("__object__") + return item(**config) + raise NotImplementedError(f"Unknown args type: {args}") + + +def create_dataset(path: str, *args, **kwargs) -> Any: + """ + Create a dataset. Requires the file to contain a "create_dataset" function. + """ + return import_item(path, "create_dataset")(*args, **kwargs) + + +def to_dict_recursive(config_obj): + if isinstance(config_obj, DictConfig): + return {k: to_dict_recursive(v) for k, v in config_obj.items()} + elif isinstance(config_obj, ListConfig): + return [to_dict_recursive(item) for item in config_obj] + return config_obj diff --git a/depth_anything_3/cli.py b/depth_anything_3/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..316d949015f33872276fb0f6a04e2712cc15ea6f --- /dev/null +++ b/depth_anything_3/cli.py @@ -0,0 +1,824 @@ +# flake8: noqa: E402 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Refactored Depth Anything 3 CLI +Clean, modular command-line interface +""" + +from __future__ import annotations + +import os +import typer + +from depth_anything_3.services import start_server +from depth_anything_3.services.gallery import gallery as gallery_main +from depth_anything_3.services.inference_service import run_inference +from depth_anything_3.services.input_handlers import ( + ColmapHandler, + ImageHandler, + ImagesHandler, + InputHandler, + VideoHandler, + parse_export_feat, +) +from depth_anything_3.utils.constants import ( + DEFAULT_EXPORT_DIR, + DEFAULT_GALLERY_DIR, + DEFAULT_GRADIO_DIR, + DEFAULT_MODEL, +) + +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + +app = typer.Typer(help="Depth Anything 3 - Video depth estimation CLI", add_completion=False) + + +# ============================================================================ +# Input type detection utilities +# ============================================================================ + +# Supported file extensions +IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif"} +VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm", ".m4v"} + + +def detect_input_type(input_path: str) -> str: + """ + Detect input type from path. + + Returns: + - "image": Single image file + - "images": Directory containing images + - "video": Video file + - "colmap": COLMAP directory structure + - "unknown": Cannot determine type + """ + if not os.path.exists(input_path): + return "unknown" + + # Check if it's a file + if os.path.isfile(input_path): + ext = os.path.splitext(input_path)[1].lower() + if ext in IMAGE_EXTENSIONS: + return "image" + elif ext in VIDEO_EXTENSIONS: + return "video" + return "unknown" + + # Check if it's a directory + if os.path.isdir(input_path): + # Check for COLMAP structure + images_dir = os.path.join(input_path, "images") + sparse_dir = os.path.join(input_path, "sparse") + + if os.path.isdir(images_dir) and os.path.isdir(sparse_dir): + return "colmap" + + # Check if directory contains image files + for item in os.listdir(input_path): + item_path = os.path.join(input_path, item) + if os.path.isfile(item_path): + ext = os.path.splitext(item)[1].lower() + if ext in IMAGE_EXTENSIONS: + return "images" + + return "unknown" + + return "unknown" + + +# ============================================================================ +# Common parameters and configuration +# ============================================================================ + +# ============================================================================ +# Inference commands +# ============================================================================ + + +@app.command() +def auto( + input_path: str = typer.Argument( + ..., help="Path to input (image, directory, video, or COLMAP)" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS]Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Video-specific options + fps: float = typer.Option(1.0, help="[Video] Sampling FPS for frame extraction"), + # COLMAP-specific options + sparse_subdir: str = typer.Option( + "", help="[COLMAP] Sparse reconstruction subdirectory (e.g., '0' for sparse/0/)" + ), + align_to_input_ext_scale: bool = typer.Option( + True, help="[COLMAP] Align prediction to input extrinsics scale" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """ + Automatically detect input type and run appropriate processing. + + Supports: + - Single image file (.jpg, .png, etc.) + - Directory of images + - Video file (.mp4, .avi, etc.) + - COLMAP directory (with 'images' and 'sparse' subdirectories) + """ + # Detect input type + input_type = detect_input_type(input_path) + + if input_type == "unknown": + typer.echo(f"❌ Error: Cannot determine input type for: {input_path}", err=True) + typer.echo("Supported inputs:", err=True) + typer.echo(" - Single image file (.jpg, .png, etc.)", err=True) + typer.echo(" - Directory containing images", err=True) + typer.echo(" - Video file (.mp4, .avi, etc.)", err=True) + typer.echo(" - COLMAP directory (with 'images/' and 'sparse/' subdirectories)", err=True) + raise typer.Exit(1) + + # Display detected type + typer.echo(f"🔍 Detected input type: {input_type.upper()}") + typer.echo(f"📁 Input path: {input_path}") + typer.echo() + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Route to appropriate handler + if input_type == "image": + typer.echo("Processing single image...") + # Process input + image_files = ImageHandler.process(input_path) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "images": + typer.echo("Processing directory of images...") + # Process input - use default extensions + image_files = ImagesHandler.process(input_path, "png,jpg,jpeg") + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "video": + typer.echo(f"Processing video with FPS={fps}...") + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Process input + image_files = VideoHandler.process(input_path, export_dir, fps) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "colmap": + typer.echo( + f"Processing COLMAP directory (sparse subdirectory: '{sparse_subdir or 'default'}')..." + ) + # Process input + image_files, extrinsics, intrinsics = ColmapHandler.process(input_path, sparse_subdir) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + typer.echo() + typer.echo("✅ Processing completed successfully!") + + +@app.command() +def image( + image_path: str = typer.Argument(..., help="Path to input image file"), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run camera pose and depth estimation on a single image.""" + # Process input + image_files = ImageHandler.process(image_path) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def images( + images_dir: str = typer.Argument(..., help="Path to directory containing input images"), + image_extensions: str = typer.Option( + "png,jpg,jpeg", help="Comma-separated image file extensions to process" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run camera pose and depth estimation on a directory of images.""" + # Process input + image_files = ImagesHandler.process(images_dir, image_extensions) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def colmap( + colmap_dir: str = typer.Argument( + ..., help="Path to COLMAP directory containing 'images' and 'sparse' subdirectories" + ), + sparse_subdir: str = typer.Option( + "", help="Sparse reconstruction subdirectory (e.g., '0' for sparse/0/, empty for sparse/)" + ), + align_to_input_ext_scale: bool = typer.Option( + True, help="Align prediction to input extrinsics scale" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run pose conditioned depth estimation on COLMAP data.""" + # Process input + image_files, extrinsics, intrinsics = ColmapHandler.process(colmap_dir, sparse_subdir) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def video( + video_path: str = typer.Argument(..., help="Path to input video file"), + fps: float = typer.Option(1.0, help="Sampling FPS for frame extraction"), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run depth estimation on video by extracting frames and processing them.""" + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Process input + image_files = VideoHandler.process(video_path, export_dir, fps) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +# ============================================================================ +# Service management commands +# ============================================================================ + + +@app.command() +def backend( + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + device: str = typer.Option("cuda", help="Device to use"), + host: str = typer.Option("127.0.0.1", help="Host to bind to"), + port: int = typer.Option(8008, help="Port to bind to"), + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path (optional)"), + api_key: str = typer.Option( + None, + help="Require this API key (X-API-Key header) on inference requests. Falls back " + "to the DA3_BACKEND_API_KEY env var, or an auto-generated key when binding to a " + "non-loopback host.", + ), + allow_unauthenticated: bool = typer.Option( + False, + help="Skip authentication entirely on a non-loopback host, instead of using an " + "auto-generated API key.", + ), +): + """Start model backend service with integrated gallery.""" + typer.echo("=" * 60) + typer.echo("🚀 Starting Depth Anything 3 Backend Server") + typer.echo("=" * 60) + typer.echo(f"Model directory: {model_dir}") + typer.echo(f"Device: {device}") + + # The gallery directory is also where /inference writes exports, so make sure it + # exists up front rather than treating a fresh checkout as "no gallery configured". + if gallery_dir: + os.makedirs(gallery_dir, exist_ok=True) + typer.echo(f"Gallery directory: {gallery_dir}") + else: + gallery_dir = None + + typer.echo() + typer.echo("📡 Server URLs (Ctrl/CMD+Click to open):") + typer.echo(f" 🏠 Home: http://{host}:{port}") + typer.echo(f" 📊 Dashboard: http://{host}:{port}/dashboard") + typer.echo(f" 📈 API Status: http://{host}:{port}/status") + + if gallery_dir: + typer.echo(f" 🎨 Gallery: http://{host}:{port}/gallery/") + + typer.echo("=" * 60) + + try: + start_server( + model_dir, + device, + host, + port, + gallery_dir, + api_key=api_key, + allow_unauthenticated=allow_unauthenticated, + ) + except KeyboardInterrupt: + typer.echo("\n👋 Backend server stopped.") + except Exception as e: + typer.echo(f"❌ Failed to start backend: {e}") + raise typer.Exit(1) + + +# ============================================================================ +# Application launch commands +# ============================================================================ + + +@app.command() +def gradio( + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + workspace_dir: str = typer.Option(DEFAULT_GRADIO_DIR, help="Workspace directory path"), + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path"), + host: str = typer.Option("127.0.0.1", help="Host address to bind to"), + port: int = typer.Option(7860, help="Port number to bind to"), + share: bool = typer.Option(False, help="Create a public link for the app"), + debug: bool = typer.Option(False, help="Enable debug mode"), + cache_examples: bool = typer.Option( + False, help="Pre-cache all example scenes at startup for faster loading" + ), + cache_gs_tag: str = typer.Option( + "", + help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.", + ), +): + """Launch Depth Anything 3 Gradio interactive web application""" + from depth_anything_3.app.gradio_app import DepthAnything3App + + # Create necessary directories + os.makedirs(workspace_dir, exist_ok=True) + os.makedirs(gallery_dir, exist_ok=True) + + typer.echo("Launching Depth Anything 3 Gradio application...") + typer.echo(f"Model directory: {model_dir}") + typer.echo(f"Workspace directory: {workspace_dir}") + typer.echo(f"Gallery directory: {gallery_dir}") + typer.echo(f"Host: {host}") + typer.echo(f"Port: {port}") + typer.echo(f"Share: {share}") + typer.echo(f"Debug mode: {debug}") + typer.echo(f"Cache examples: {cache_examples}") + if cache_examples: + if cache_gs_tag: + typer.echo( + f"Cache GS Tag: '{cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)" + ) + else: + typer.echo(f"Cache GS Tag: None (all scenes will use low-res only)") + + try: + # Initialize and launch application + app = DepthAnything3App( + model_dir=model_dir, workspace_dir=workspace_dir, gallery_dir=gallery_dir + ) + + # Pre-cache examples if requested + if cache_examples: + typer.echo("\n" + "=" * 60) + typer.echo("Pre-caching mode enabled") + if cache_gs_tag: + typer.echo(f"Scenes containing '{cache_gs_tag}' will use HIGH-RES + 3DGS") + typer.echo(f"Other scenes will use LOW-RES only") + else: + typer.echo(f"All scenes will use LOW-RES only") + typer.echo("=" * 60) + app.cache_examples( + show_cam=True, + filter_black_bg=False, + filter_white_bg=False, + save_percentage=20.0, + num_max_points=1000, + cache_gs_tag=cache_gs_tag, + gs_trj_mode="smooth", + gs_video_quality="low", + ) + + # Prepare launch arguments + launch_kwargs = {"share": share, "debug": debug} + + app.launch(host=host, port=port, **launch_kwargs) + + except KeyboardInterrupt: + typer.echo("\nGradio application stopped.") + except Exception as e: + typer.echo(f"Failed to launch Gradio application: {e}") + raise typer.Exit(1) + + +@app.command() +def gallery( + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery root directory"), + host: str = typer.Option("127.0.0.1", help="Host address to bind to"), + port: int = typer.Option(8007, help="Port number to bind to"), + open_browser: bool = typer.Option(False, help="Open browser after launch"), +): + """Launch Depth Anything 3 Gallery server""" + + # Validate gallery directory + if not os.path.exists(gallery_dir): + raise typer.BadParameter(f"Gallery directory not found: {gallery_dir}") + + typer.echo("Launching Depth Anything 3 Gallery server...") + typer.echo(f"Gallery directory: {gallery_dir}") + typer.echo(f"Host: {host}") + typer.echo(f"Port: {port}") + typer.echo(f"Auto-open browser: {open_browser}") + + try: + # Set command line arguments + import sys + + sys.argv = ["gallery", "--dir", gallery_dir, "--host", host, "--port", str(port)] + if open_browser: + sys.argv.append("--open") + + # Launch gallery server + gallery_main() + + except KeyboardInterrupt: + typer.echo("\nGallery server stopped.") + except Exception as e: + typer.echo(f"Failed to launch Gallery server: {e}") + raise typer.Exit(1) + + +if __name__ == "__main__": + app() diff --git a/depth_anything_3/configs/da3-base.yaml b/depth_anything_3/configs/da3-base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c52a7e5018388a174841469f9a94dc995e14f220 --- /dev/null +++ b/depth_anything_3/configs/da3-base.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitb + out_layers: [5, 7, 9, 11] + alt_start: 4 + qknorm_start: 4 + rope_start: 4 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 1536 + output_dim: 2 + features: &head_features 128 + out_channels: &head_out_channels [96, 192, 384, 768] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 768 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 1536 diff --git a/depth_anything_3/configs/da3-giant.yaml b/depth_anything_3/configs/da3-giant.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f5a75c043353aa3e5b7c5368e4b26416c2b0b8b0 --- /dev/null +++ b/depth_anything_3/configs/da3-giant.yaml @@ -0,0 +1,71 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitg + out_layers: [19, 27, 33, 39] + alt_start: 13 + qknorm_start: 13 + rope_start: 13 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 3072 + output_dim: 2 + features: &head_features 256 + out_channels: &head_out_channels [256, 512, 1024, 1024] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 1536 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 3072 + + +gs_head: + __object__: + path: depth_anything_3.model.gsdpt + name: GSDPT + args: as_params + + dim_in: *head_dim_in + output_dim: 38 # should align with gs_adapter's setting, for gs params + features: *head_features + out_channels: *head_out_channels + + +gs_adapter: + __object__: + path: depth_anything_3.model.gs_adapter + name: GaussianAdapter + args: as_params + + sh_degree: 2 + pred_color: false # predict SH coefficient if false + pred_offset_depth: true + pred_offset_xy: true + gaussian_scale_min: 1e-5 + gaussian_scale_max: 30.0 diff --git a/depth_anything_3/configs/da3-large.yaml b/depth_anything_3/configs/da3-large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4fa367c9d9b46eb9a62aef7041f68c709eb4c6e3 --- /dev/null +++ b/depth_anything_3/configs/da3-large.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [11, 15, 19, 23] + alt_start: 8 + qknorm_start: 8 + rope_start: 8 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 2048 + output_dim: 2 + features: &head_features 256 + out_channels: &head_out_channels [256, 512, 1024, 1024] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 1024 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 2048 diff --git a/depth_anything_3/configs/da3-small.yaml b/depth_anything_3/configs/da3-small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..10887437697fc9f2614c03add73fe9858b309d91 --- /dev/null +++ b/depth_anything_3/configs/da3-small.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vits + out_layers: [5, 7, 9, 11] + alt_start: 4 + qknorm_start: 4 + rope_start: 4 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 768 + output_dim: 2 + features: &head_features 64 + out_channels: &head_out_channels [48, 96, 192, 384] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 384 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 768 diff --git a/depth_anything_3/configs/da3metric-large.yaml b/depth_anything_3/configs/da3metric-large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..124635cfd952c25c8857ee3da63ce0444c4377f2 --- /dev/null +++ b/depth_anything_3/configs/da3metric-large.yaml @@ -0,0 +1,28 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [4, 11, 17, 23] + alt_start: -1 # -1 means disable + qknorm_start: -1 + rope_start: -1 + cat_token: False + +head: + __object__: + path: depth_anything_3.model.dpt + name: DPT + args: as_params + + dim_in: 1024 + output_dim: 1 + features: 256 + out_channels: [256, 512, 1024, 1024] diff --git a/depth_anything_3/configs/da3mono-large.yaml b/depth_anything_3/configs/da3mono-large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..124635cfd952c25c8857ee3da63ce0444c4377f2 --- /dev/null +++ b/depth_anything_3/configs/da3mono-large.yaml @@ -0,0 +1,28 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [4, 11, 17, 23] + alt_start: -1 # -1 means disable + qknorm_start: -1 + rope_start: -1 + cat_token: False + +head: + __object__: + path: depth_anything_3.model.dpt + name: DPT + args: as_params + + dim_in: 1024 + output_dim: 1 + features: 256 + out_channels: [256, 512, 1024, 1024] diff --git a/depth_anything_3/configs/da3nested-giant-large.yaml b/depth_anything_3/configs/da3nested-giant-large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..595c122b1dc976ecfec58b133b2b60d8f724618c --- /dev/null +++ b/depth_anything_3/configs/da3nested-giant-large.yaml @@ -0,0 +1,10 @@ +__object__: + path: depth_anything_3.model.da3 + name: NestedDepthAnything3Net + args: as_params + +anyview: + __inherit__: depth_anything_3.configs.da3-giant + +metric: + __inherit__: depth_anything_3.configs.da3metric-large diff --git a/depth_anything_3/model/__init__.py b/depth_anything_3/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..57a2a45132eeae8d58a11a26036c54feef9cfe16 --- /dev/null +++ b/depth_anything_3/model/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from depth_anything_3.model.da3 import DepthAnything3Net, NestedDepthAnything3Net + +__export__ = [ + NestedDepthAnything3Net, + DepthAnything3Net, +] diff --git a/depth_anything_3/model/cam_dec.py b/depth_anything_3/model/cam_dec.py new file mode 100644 index 0000000000000000000000000000000000000000..3353b403683bf556b3823081863573dc7f5f719e --- /dev/null +++ b/depth_anything_3/model/cam_dec.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn + + +class CameraDec(nn.Module): + def __init__(self, dim_in=1536): + super().__init__() + output_dim = dim_in + self.backbone = nn.Sequential( + nn.Linear(output_dim, output_dim), + nn.ReLU(), + nn.Linear(output_dim, output_dim), + nn.ReLU(), + ) + self.fc_t = nn.Linear(output_dim, 3) + self.fc_qvec = nn.Linear(output_dim, 4) + self.fc_fov = nn.Sequential(nn.Linear(output_dim, 2), nn.ReLU()) + + def forward(self, feat, camera_encoding=None, *args, **kwargs): + B, N = feat.shape[:2] + feat = feat.reshape(B * N, -1) + feat = self.backbone(feat) + out_t = self.fc_t(feat.float()).reshape(B, N, 3) + if camera_encoding is None: + out_qvec = self.fc_qvec(feat.float()).reshape(B, N, 4) + out_fov = self.fc_fov(feat.float()).reshape(B, N, 2) + else: + out_qvec = camera_encoding[..., 3:7] + out_fov = camera_encoding[..., -2:] + pose_enc = torch.cat([out_t, out_qvec, out_fov], dim=-1) + return pose_enc diff --git a/depth_anything_3/model/cam_enc.py b/depth_anything_3/model/cam_enc.py new file mode 100644 index 0000000000000000000000000000000000000000..bf28e701442fa73d89e54b409800908c138a93d8 --- /dev/null +++ b/depth_anything_3/model/cam_enc.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch.nn as nn + +from depth_anything_3.model.utils.attention import Mlp +from depth_anything_3.model.utils.block import Block +from depth_anything_3.model.utils.transform import extri_intri_to_pose_encoding +from depth_anything_3.utils.geometry import affine_inverse + + +class CameraEnc(nn.Module): + """ + CameraHead predicts camera parameters from token representations using iterative refinement. + + It applies a series of transformer blocks (the "trunk") to dedicated camera tokens. + """ + + def __init__( + self, + dim_out: int = 1024, + dim_in: int = 9, + trunk_depth: int = 4, + target_dim: int = 9, + num_heads: int = 16, + mlp_ratio: int = 4, + init_values: float = 0.01, + **kwargs, + ): + super().__init__() + self.target_dim = target_dim + self.trunk_depth = trunk_depth + self.trunk = nn.Sequential( + *[ + Block( + dim=dim_out, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + init_values=init_values, + ) + for _ in range(trunk_depth) + ] + ) + self.token_norm = nn.LayerNorm(dim_out) + self.trunk_norm = nn.LayerNorm(dim_out) + self.pose_branch = Mlp( + in_features=dim_in, + hidden_features=dim_out // 2, + out_features=dim_out, + drop=0, + ) + + def forward( + self, + ext, + ixt, + image_size, + ) -> tuple: + c2ws = affine_inverse(ext) + pose_encoding = extri_intri_to_pose_encoding( + c2ws, + ixt, + image_size, + ) + pose_tokens = self.pose_branch(pose_encoding) + pose_tokens = self.token_norm(pose_tokens) + pose_tokens = self.trunk(pose_tokens) + pose_tokens = self.trunk_norm(pose_tokens) + return pose_tokens diff --git a/depth_anything_3/model/da3.py b/depth_anything_3/model/da3.py new file mode 100644 index 0000000000000000000000000000000000000000..f9e5bd683a8a697647da30e4d93a0ead03e9545d --- /dev/null +++ b/depth_anything_3/model/da3.py @@ -0,0 +1,442 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import torch +import torch.nn as nn +from addict import Dict +from omegaconf import DictConfig, OmegaConf + +from depth_anything_3.cfg import create_object +from depth_anything_3.model.utils.transform import pose_encoding_to_extri_intri +from depth_anything_3.utils.alignment import ( + apply_metric_scaling, + compute_alignment_mask, + compute_sky_mask, + least_squares_scale_scalar, + sample_tensor_for_quantile, + set_sky_regions_to_max_depth, +) +from depth_anything_3.utils.geometry import affine_inverse, as_homogeneous, map_pdf_to_opacity +from depth_anything_3.utils.ray_utils import get_extrinsic_from_camray + + +def _wrap_cfg(cfg_obj): + return OmegaConf.create(cfg_obj) + + +class DepthAnything3Net(nn.Module): + """ + Depth Anything 3 network for depth estimation and camera pose estimation. + + This network consists of: + - Backbone: DinoV2 feature extractor + - Head: DPT or DualDPT for depth prediction + - Optional camera decoders for pose estimation + - Optional GSDPT for 3DGS prediction + + Args: + preset: Configuration preset containing network dimensions and settings + + Returns: + Dictionary containing: + - depth: Predicted depth map (B, H, W) + - depth_conf: Depth confidence map (B, H, W) + - extrinsics: Camera extrinsics (B, N, 4, 4) + - intrinsics: Camera intrinsics (B, N, 3, 3) + - gaussians: 3D Gaussian Splats (world space), type: model.gs_adapter.Gaussians + - aux: Auxiliary features for specified layers + """ + + # Patch size for feature extraction + PATCH_SIZE = 14 + + def __init__(self, net, head, cam_dec=None, cam_enc=None, gs_head=None, gs_adapter=None): + """ + Initialize DepthAnything3Net with given yaml-initialized configuration. + """ + super().__init__() + self.backbone = net if isinstance(net, nn.Module) else create_object(_wrap_cfg(net)) + self.head = head if isinstance(head, nn.Module) else create_object(_wrap_cfg(head)) + self.cam_dec, self.cam_enc = None, None + if cam_dec is not None: + self.cam_dec = ( + cam_dec if isinstance(cam_dec, nn.Module) else create_object(_wrap_cfg(cam_dec)) + ) + self.cam_enc = ( + cam_enc if isinstance(cam_enc, nn.Module) else create_object(_wrap_cfg(cam_enc)) + ) + self.gs_adapter, self.gs_head = None, None + if gs_head is not None and gs_adapter is not None: + self.gs_adapter = ( + gs_adapter + if isinstance(gs_adapter, nn.Module) + else create_object(_wrap_cfg(gs_adapter)) + ) + gs_out_dim = self.gs_adapter.d_in + 1 + if isinstance(gs_head, nn.Module): + assert ( + gs_head.out_dim == gs_out_dim + ), f"gs_head.out_dim should be {gs_out_dim}, got {gs_head.out_dim}" + self.gs_head = gs_head + else: + assert ( + gs_head["output_dim"] == gs_out_dim + ), f"gs_head output_dim should set to {gs_out_dim}, got {gs_head['output_dim']}" + self.gs_head = create_object(_wrap_cfg(gs_head)) + + def forward( + self, + x: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = [], + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> Dict[str, torch.Tensor]: + """ + Forward pass through the network. + + Args: + x: Input images (B, N, 3, H, W) + extrinsics: Camera extrinsics (B, N, 4, 4) + intrinsics: Camera intrinsics (B, N, 3, 3) + feat_layers: List of layer indices to extract features from + infer_gs: Enable Gaussian Splatting branch + use_ray_pose: Use ray-based pose estimation + ref_view_strategy: Strategy for selecting reference view + + Returns: + Dictionary containing predictions and auxiliary features + """ + # Extract features using backbone + if extrinsics is not None: + with torch.autocast(device_type=x.device.type, enabled=False): + cam_token = self.cam_enc(extrinsics, intrinsics, x.shape[-2:]) + else: + cam_token = None + + feats, aux_feats = self.backbone( + x, cam_token=cam_token, export_feat_layers=export_feat_layers, ref_view_strategy=ref_view_strategy + ) + # feats = [[item for item in feat] for feat in feats] + H, W = x.shape[-2], x.shape[-1] + + # Process features through depth head + with torch.autocast(device_type=x.device.type, enabled=False): + output = self._process_depth_head(feats, H, W) + if use_ray_pose: + output = self._process_ray_pose_estimation(output, H, W) + else: + output = self._process_camera_estimation(feats, H, W, output) + if infer_gs: + output = self._process_gs_head(feats, H, W, output, x, extrinsics, intrinsics) + + output = self._process_mono_sky_estimation(output) + + # Extract auxiliary features if requested + output.aux = self._extract_auxiliary_features(aux_feats, export_feat_layers, H, W) + + return output + + def _process_mono_sky_estimation( + self, output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Process mono sky estimation.""" + if "sky" not in output: + return output + non_sky_mask = compute_sky_mask(output.sky, threshold=0.3) + if non_sky_mask.sum() <= 10: + return output + if (~non_sky_mask).sum() <= 10: + return output + + non_sky_depth = output.depth[non_sky_mask] + if non_sky_depth.numel() > 100000: + idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device) + sampled_depth = non_sky_depth[idx] + else: + sampled_depth = non_sky_depth + non_sky_max = torch.quantile(sampled_depth, 0.99) + + # Set sky regions to maximum depth and high confidence + output.depth, _ = set_sky_regions_to_max_depth( + output.depth, None, non_sky_mask, max_depth=non_sky_max + ) + return output + + def _process_ray_pose_estimation( + self, output: Dict[str, torch.Tensor], height: int, width: int + ) -> Dict[str, torch.Tensor]: + """Process ray pose estimation if ray pose decoder is available.""" + if "ray" in output and "ray_conf" in output: + pred_extrinsic, pred_focal_lengths, pred_principal_points = get_extrinsic_from_camray( + output.ray, + output.ray_conf, + output.ray.shape[-3], + output.ray.shape[-2], + ) + pred_extrinsic = affine_inverse(pred_extrinsic) # w2c -> c2w + pred_extrinsic = pred_extrinsic[:, :, :3, :] + pred_intrinsic = torch.eye(3, 3)[None, None].repeat(pred_extrinsic.shape[0], pred_extrinsic.shape[1], 1, 1).clone().to(pred_extrinsic.device) + pred_intrinsic[:, :, 0, 0] = pred_focal_lengths[:, :, 0] / 2 * width + pred_intrinsic[:, :, 1, 1] = pred_focal_lengths[:, :, 1] / 2 * height + pred_intrinsic[:, :, 0, 2] = pred_principal_points[:, :, 0] * width * 0.5 + pred_intrinsic[:, :, 1, 2] = pred_principal_points[:, :, 1] * height * 0.5 + del output.ray + del output.ray_conf + output.extrinsics = pred_extrinsic + output.intrinsics = pred_intrinsic + return output + + def _process_depth_head( + self, feats: list[torch.Tensor], H: int, W: int + ) -> Dict[str, torch.Tensor]: + """Process features through the depth prediction head.""" + return self.head(feats, H, W, patch_start_idx=0) + + def _process_camera_estimation( + self, feats: list[torch.Tensor], H: int, W: int, output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Process camera pose estimation if camera decoder is available.""" + if self.cam_dec is not None: + pose_enc = self.cam_dec(feats[-1][1]) + # Remove ray information as it's not needed for pose estimation + if "ray" in output: + del output.ray + if "ray_conf" in output: + del output.ray_conf + + # Convert pose encoding to extrinsics and intrinsics + c2w, ixt = pose_encoding_to_extri_intri(pose_enc, (H, W)) + output.extrinsics = affine_inverse(c2w) + output.intrinsics = ixt + + return output + + def _process_gs_head( + self, + feats: list[torch.Tensor], + H: int, + W: int, + output: Dict[str, torch.Tensor], + in_images: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + ) -> Dict[str, torch.Tensor]: + """Process 3DGS parameters estimation if 3DGS head is available.""" + if self.gs_head is None or self.gs_adapter is None: + return output + assert output.get("depth", None) is not None, "must provide MV depth for the GS head." + + # The depth is defined in the DA3 model's camera space, + # so even with provided GT camera poses, + # we instead use the predicted camera poses for better alignment. + ctx_extr = output.get("extrinsics", None) + ctx_intr = output.get("intrinsics", None) + assert ( + ctx_extr is not None and ctx_intr is not None + ), "must process camera info first if GT is not available" + + gt_extr = extrinsics + # homo the extr if needed + ctx_extr = as_homogeneous(ctx_extr) + if gt_extr is not None: + gt_extr = as_homogeneous(gt_extr) + + # forward through the gs_dpt head to get 'camera space' parameters + gs_outs = self.gs_head( + feats=feats, + H=H, + W=W, + patch_start_idx=0, + images=in_images, + ) + raw_gaussians = gs_outs.raw_gs + densities = gs_outs.raw_gs_conf + + # convert to 'world space' 3DGS parameters; ready to export and render + # gt_extr could be None, and will be used to align the pose scale if available + gs_world = self.gs_adapter( + extrinsics=ctx_extr, + intrinsics=ctx_intr, + depths=output.depth, + opacities=map_pdf_to_opacity(densities), + raw_gaussians=raw_gaussians, + image_shape=(H, W), + gt_extrinsics=gt_extr, + ) + output.gaussians = gs_world + + return output + + def _extract_auxiliary_features( + self, feats: list[torch.Tensor], feat_layers: list[int], H: int, W: int + ) -> Dict[str, torch.Tensor]: + """Extract auxiliary features from specified layers.""" + aux_features = Dict() + assert len(feats) == len(feat_layers) + for feat, feat_layer in zip(feats, feat_layers): + # Reshape features to spatial dimensions + feat_reshaped = feat.reshape( + [ + feat.shape[0], + feat.shape[1], + H // self.PATCH_SIZE, + W // self.PATCH_SIZE, + feat.shape[-1], + ] + ) + aux_features[f"feat_layer_{feat_layer}"] = feat_reshaped + + return aux_features + + +class NestedDepthAnything3Net(nn.Module): + """ + Nested Depth Anything 3 network with metric scaling capabilities. + + This network combines two DepthAnything3Net branches: + - Main branch: Standard depth estimation + - Metric branch: Metric depth estimation for scaling alignment + + The network performs depth alignment using least squares scaling + and handles sky region masking for improved depth estimation. + + Args: + preset: Configuration for the main depth estimation branch + second_preset: Configuration for the metric depth branch + """ + + def __init__(self, anyview: DictConfig, metric: DictConfig): + """ + Initialize NestedDepthAnything3Net with two branches. + + Args: + preset: Configuration for main depth estimation branch + second_preset: Configuration for metric depth branch + """ + super().__init__() + self.da3 = create_object(anyview) + self.da3_metric = create_object(metric) + + def forward( + self, + x: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = [], + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> Dict[str, torch.Tensor]: + """ + Forward pass through both branches with metric scaling alignment. + + Args: + x: Input images (B, N, 3, H, W) + extrinsics: Camera extrinsics (B, N, 4, 4) - unused + intrinsics: Camera intrinsics (B, N, 3, 3) - unused + feat_layers: List of layer indices to extract features from + infer_gs: Enable Gaussian Splatting branch + use_ray_pose: Use ray-based pose estimation + ref_view_strategy: Strategy for selecting reference view + + Returns: + Dictionary containing aligned depth predictions and camera parameters + """ + # Get predictions from both branches + output = self.da3( + x, extrinsics, intrinsics, export_feat_layers=export_feat_layers, infer_gs=infer_gs, use_ray_pose=use_ray_pose, ref_view_strategy=ref_view_strategy + ) + metric_output = self.da3_metric(x) + + # Apply metric scaling and alignment + output = self._apply_metric_scaling(output, metric_output) + output = self._apply_depth_alignment(output, metric_output) + output = self._handle_sky_regions(output, metric_output) + + return output + + def _apply_metric_scaling( + self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Apply metric scaling to the metric depth output.""" + # Scale metric depth based on camera intrinsics + metric_output.depth = apply_metric_scaling( + metric_output.depth, + output.intrinsics, + ) + return output + + def _apply_depth_alignment( + self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Apply depth alignment using least squares scaling.""" + # Compute non-sky mask + non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3) + + # Ensure we have enough non-sky pixels + assert non_sky_mask.sum() > 10, "Insufficient non-sky pixels for alignment" + + # Sample depth confidence for quantile computation + depth_conf_ns = output.depth_conf[non_sky_mask] + depth_conf_sampled = sample_tensor_for_quantile(depth_conf_ns, max_samples=100000) + median_conf = torch.quantile(depth_conf_sampled, 0.5) + + # Compute alignment mask + align_mask = compute_alignment_mask( + output.depth_conf, non_sky_mask, output.depth, metric_output.depth, median_conf + ) + + # Compute scale factor using least squares + valid_depth = output.depth[align_mask] + valid_metric_depth = metric_output.depth[align_mask] + scale_factor = least_squares_scale_scalar(valid_metric_depth, valid_depth) + + # Apply scaling to depth and extrinsics + output.depth *= scale_factor + output.extrinsics[:, :, :3, 3] *= scale_factor + output.is_metric = 1 + output.scale_factor = scale_factor.item() + + return output + + def _handle_sky_regions( + self, + output: Dict[str, torch.Tensor], + metric_output: Dict[str, torch.Tensor], + sky_depth_def: float = 200.0, + ) -> Dict[str, torch.Tensor]: + """Handle sky regions by setting them to maximum depth.""" + non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3) + + # Compute maximum depth for non-sky regions + # Use sampling to safely compute quantile on large tensors + non_sky_depth = output.depth[non_sky_mask] + if non_sky_depth.numel() > 100000: + idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device) + sampled_depth = non_sky_depth[idx] + else: + sampled_depth = non_sky_depth + non_sky_max = min(torch.quantile(sampled_depth, 0.99), sky_depth_def) + + # Set sky regions to maximum depth and high confidence + output.depth, output.depth_conf = set_sky_regions_to_max_depth( + output.depth, output.depth_conf, non_sky_mask, max_depth=non_sky_max + ) + + return output diff --git a/depth_anything_3/model/dinov2/dinov2.py b/depth_anything_3/model/dinov2/dinov2.py new file mode 100644 index 0000000000000000000000000000000000000000..ee0d88bdd6d33edda6c0fb237ee0c77ffc3f6034 --- /dev/null +++ b/depth_anything_3/model/dinov2/dinov2.py @@ -0,0 +1,64 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + + +from typing import List +import torch.nn as nn + +from depth_anything_3.model.dinov2.vision_transformer import ( + vit_base, + vit_giant2, + vit_large, + vit_small, +) + + +class DinoV2(nn.Module): + def __init__( + self, + name: str, + out_layers: List[int], + alt_start: int = -1, + qknorm_start: int = -1, + rope_start: int = -1, + cat_token: bool = True, + **kwargs, + ): + super().__init__() + assert name in {"vits", "vitb", "vitl", "vitg"} + self.name = name + self.out_layers = out_layers + self.alt_start = alt_start + self.qknorm_start = qknorm_start + self.rope_start = rope_start + self.cat_token = cat_token + encoder_map = { + "vits": vit_small, + "vitb": vit_base, + "vitl": vit_large, + "vitg": vit_giant2, + } + encoder_fn = encoder_map[self.name] + ffn_layer = "swiglufused" if self.name == "vitg" else "mlp" + self.pretrained = encoder_fn( + img_size=518, + patch_size=14, + ffn_layer=ffn_layer, + alt_start=alt_start, + qknorm_start=qknorm_start, + rope_start=rope_start, + cat_token=cat_token, + ) + + def forward(self, x, **kwargs): + return self.pretrained.get_intermediate_layers( + x, + self.out_layers, + **kwargs, + ) diff --git a/depth_anything_3/model/dinov2/layers/__init__.py b/depth_anything_3/model/dinov2/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97dfba90c12b2ffc5f0f1f823b6384c9b4cd6fa2 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# from .attention import MemEffAttention +from .block import Block +from .layer_scale import LayerScale +from .mlp import Mlp +from .patch_embed import PatchEmbed +from .rope import PositionGetter, RotaryPositionEmbedding2D +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused + +__all__ = [ + Mlp, + PatchEmbed, + SwiGLUFFN, + SwiGLUFFNFused, + Block, + # MemEffAttention, + LayerScale, + PositionGetter, + RotaryPositionEmbedding2D, +] diff --git a/depth_anything_3/model/dinov2/layers/attention.py b/depth_anything_3/model/dinov2/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..096b9d41ddc95b9c4652597b18f53aee31a573b6 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/attention.py @@ -0,0 +1,100 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import logging +import torch.nn.functional as F +from torch import Tensor, nn + +logger = logging.getLogger("dinov2") + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + qk_norm: bool = False, + fused_attn: bool = True, # use F.scaled_dot_product_attention or not + rope=None, + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + self.fused_attn = fused_attn + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv[0], qkv[1], qkv[2] + q, k = self.q_norm(q), self.k_norm(k) + if self.rope is not None and pos is not None: + q = self.rope(q, pos) + k = self.rope(k, pos) + if self.fused_attn: + x = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.attn_drop.p if self.training else 0.0, + attn_mask=( + (attn_mask)[:, None].repeat(1, self.num_heads, 1, 1) + if attn_mask is not None + else None + ), + ) + else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + def _forward(self, x: Tensor) -> Tensor: + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + attn = q @ k.transpose(-2, -1) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x diff --git a/depth_anything_3/model/dinov2/layers/block.py b/depth_anything_3/model/dinov2/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..731519b68b16c8936b765dd1620fbb9c81087f96 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/block.py @@ -0,0 +1,143 @@ +# flake8: noqa: F821 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +import logging +from typing import Callable, Optional +import torch +from torch import Tensor, nn + +from .attention import Attention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + +logger = logging.getLogger("dinov2") +XFORMERS_AVAILABLE = True + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + qk_norm: bool = False, + rope=None, + ln_eps: float = 1e-6, + ) -> None: + super().__init__() + # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") + self.norm1 = norm_layer(dim, eps=ln_eps) + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + qk_norm=qk_norm, + rope=rope, + ) + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim, eps=ln_eps) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + # the overhead is compensated only for a drop path rate larger than 0.1 + x = drop_add_residual_stochastic_depth( + x, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + pos=pos, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x, pos=pos, attn_mask=attn_mask)) + x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 + else: + x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, + pos: Optional[Tensor] = None, +) -> Tensor: + # 1) extract subset using permutation + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + # 2) apply residual_func to get residual + if pos is not None: + # if necessary, apply rope to the subset + pos = pos[brange] + residual = residual_func(x_subset, pos=pos) + else: + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + # 3) add the residual + x_plus_residual = torch.index_add( + x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor + ) + return x_plus_residual.view_as(x) + + +def get_branges_scales(x, sample_drop_ratio=0.0): + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + residual_scale_factor = b / sample_subset_size + return brange, residual_scale_factor diff --git a/depth_anything_3/model/dinov2/layers/drop_path.py b/depth_anything_3/model/dinov2/layers/drop_path.py new file mode 100644 index 0000000000000000000000000000000000000000..1c2cc94e969711f1eb9f62093b79a0139b9bfb1e --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/drop_path.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/depth_anything_3/model/dinov2/layers/layer_scale.py b/depth_anything_3/model/dinov2/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..898ee12d8b4b65d30d8c041588a8277a8f13d4f2 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/layer_scale.py @@ -0,0 +1,31 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa: E501 + +from typing import Union +import torch +from torch import Tensor, nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.dim = dim + self.inplace = inplace + self.init_values = init_values + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + def extra_repr(self) -> str: + return f"{self.dim}, init_values={self.init_values}, inplace={self.inplace}" diff --git a/depth_anything_3/model/dinov2/layers/mlp.py b/depth_anything_3/model/dinov2/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..78ad0d8897ddb77e95d2e188a579f6e1d21e3fb5 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/mlp.py @@ -0,0 +1,40 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/depth_anything_3/model/dinov2/layers/patch_embed.py b/depth_anything_3/model/dinov2/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..64bf6be8994fe52b2fdf1753fdb9f7a691e14980 --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/patch_embed.py @@ -0,0 +1,94 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union +import torch.nn as nn +from torch import Tensor + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert ( + H % patch_H == 0 + ), f"Input image height {H} is not a multiple of patch height {patch_H}" + assert ( + W % patch_W == 0 + ), f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = ( + Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + ) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/depth_anything_3/model/dinov2/layers/rope.py b/depth_anything_3/model/dinov2/layers/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..f75ba37c160cd806a32576e2a1704b3352dec7af --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/rope.py @@ -0,0 +1,200 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + + +# Implementation of 2D Rotary Position Embeddings (RoPE). + +# This module provides a clean implementation of 2D Rotary Position Embeddings, +# which extends the original RoPE concept to handle 2D spatial positions. + +# Inspired by: +# https://github.com/meta-llama/codellama/blob/main/llama/model.py +# https://github.com/naver-ai/rope-vit + + +from typing import Dict, Tuple +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class PositionGetter: + """Generates and caches 2D spatial positions for patches in a grid. + + This class efficiently manages the generation of spatial coordinates for patches + in a 2D grid, caching results to avoid redundant computations. + + Attributes: + position_cache: Dictionary storing precomputed position tensors for different + grid dimensions. + """ + + def __init__(self): + """Initializes the position generator with an empty cache.""" + self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {} + + def __call__( + self, batch_size: int, height: int, width: int, device: torch.device + ) -> torch.Tensor: + """Generates spatial positions for a batch of patches. + + Args: + batch_size: Number of samples in the batch. + height: Height of the grid in patches. + width: Width of the grid in patches. + device: Target device for the position tensor. + + Returns: + Tensor of shape (batch_size, height*width, 2) containing y,x coordinates + for each position in the grid, repeated for each batch item. + """ + if (height, width) not in self.position_cache: + y_coords = torch.arange(height, device=device) + x_coords = torch.arange(width, device=device) + positions = torch.cartesian_prod(y_coords, x_coords) + self.position_cache[height, width] = positions + + cached_positions = self.position_cache[height, width] + return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone() + + +class RotaryPositionEmbedding2D(nn.Module): + """2D Rotary Position Embedding implementation. + + This module applies rotary position embeddings to input tokens based on their + 2D spatial positions. It handles the position-dependent rotation of features + separately for vertical and horizontal dimensions. + + Args: + frequency: Base frequency for the position embeddings. Default: 100.0 + scaling_factor: Scaling factor for frequency computation. Default: 1.0 + + Attributes: + base_frequency: Base frequency for computing position embeddings. + scaling_factor: Factor to scale the computed frequencies. + frequency_cache: Cache for storing precomputed frequency components. + """ + + def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0): + """Initializes the 2D RoPE module.""" + super().__init__() + self.base_frequency = frequency + self.scaling_factor = scaling_factor + self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {} + + def _compute_frequency_components( + self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Computes frequency components for rotary embeddings. + + Args: + dim: Feature dimension (must be even). + seq_len: Maximum sequence length. + device: Target device for computations. + dtype: Data type for the computed tensors. + + Returns: + Tuple of (cosine, sine) tensors for frequency components. + """ + cache_key = (dim, seq_len, device, dtype) + if cache_key not in self.frequency_cache: + # Compute frequency bands + exponents = torch.arange(0, dim, 2, device=device).float() / dim + inv_freq = 1.0 / (self.base_frequency**exponents) + + # Generate position-dependent frequencies + positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + angles = torch.einsum("i,j->ij", positions, inv_freq) + + # Compute and cache frequency components + angles = angles.to(dtype) + angles = torch.cat((angles, angles), dim=-1) + cos_components = angles.cos().to(dtype) + sin_components = angles.sin().to(dtype) + self.frequency_cache[cache_key] = (cos_components, sin_components) + + return self.frequency_cache[cache_key] + + @staticmethod + def _rotate_features(x: torch.Tensor) -> torch.Tensor: + """Performs feature rotation by splitting and recombining feature dimensions. + + Args: + x: Input tensor to rotate. + + Returns: + Rotated feature tensor. + """ + feature_dim = x.shape[-1] + x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_1d_rope( + self, + tokens: torch.Tensor, + positions: torch.Tensor, + cos_comp: torch.Tensor, + sin_comp: torch.Tensor, + ) -> torch.Tensor: + """Applies 1D rotary position embeddings along one dimension. + + Args: + tokens: Input token features. + positions: Position indices. + cos_comp: Cosine components for rotation. + sin_comp: Sine components for rotation. + + Returns: + Tokens with applied rotary position embeddings. + """ + # Embed positions with frequency components + cos = F.embedding(positions, cos_comp)[:, None, :, :] + sin = F.embedding(positions, sin_comp)[:, None, :, :] + # Apply rotation + return (tokens * cos) + (self._rotate_features(tokens) * sin) + + def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Applies 2D rotary position embeddings to input tokens. + + Args: + tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim). + The feature dimension (dim) must be divisible by 4. + positions: Position tensor of shape (batch_size, n_tokens, 2) containing + the y and x coordinates for each token. + + Returns: + Tensor of same shape as input with applied 2D rotary position embeddings. + + Raises: + AssertionError: If input dimensions are invalid or positions are malformed. + """ + # Validate inputs + assert tokens.size(-1) % 2 == 0, "Feature dimension must be even" + assert ( + positions.ndim == 3 and positions.shape[-1] == 2 + ), "Positions must have shape (batch_size, n_tokens, 2)" + + # Compute feature dimension for each spatial direction + feature_dim = tokens.size(-1) // 2 + + # Get frequency components + max_position = int(positions.max()) + 1 + cos_comp, sin_comp = self._compute_frequency_components( + feature_dim, max_position, tokens.device, tokens.dtype + ) + + # Split features for vertical and horizontal processing + vertical_features, horizontal_features = tokens.chunk(2, dim=-1) + + # Apply RoPE separately for each dimension + vertical_features = self._apply_1d_rope( + vertical_features, positions[..., 0], cos_comp, sin_comp + ) + horizontal_features = self._apply_1d_rope( + horizontal_features, positions[..., 1], cos_comp, sin_comp + ) + + # Combine processed features + return torch.cat((vertical_features, horizontal_features), dim=-1) diff --git a/depth_anything_3/model/dinov2/layers/swiglu_ffn.py b/depth_anything_3/model/dinov2/layers/swiglu_ffn.py new file mode 100644 index 0000000000000000000000000000000000000000..c8f58e5b265f74c82a4c40adc3ee33a965e503cf --- /dev/null +++ b/depth_anything_3/model/dinov2/layers/swiglu_ffn.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional +import torch.nn.functional as F +from torch import Tensor, nn + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +try: + from xformers.ops import SwiGLU + + XFORMERS_AVAILABLE = True +except ImportError: + SwiGLU = SwiGLUFFN + XFORMERS_AVAILABLE = False + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/depth_anything_3/model/dinov2/vision_transformer.py b/depth_anything_3/model/dinov2/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f0f8e23ad21d8c5a60626955d38d740f98bb5c48 --- /dev/null +++ b/depth_anything_3/model/dinov2/vision_transformer.py @@ -0,0 +1,456 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import math +from typing import Callable, List, Sequence, Tuple, Union +import numpy as np +import torch +import torch.nn as nn +import torch.utils.checkpoint +from einops import rearrange + +from depth_anything_3.utils.logger import logger + +from .layers import LayerScale # noqa: F401 +from .layers import Mlp # noqa: F401 +from .layers import ( # noqa: F401 + Block, + PatchEmbed, + PositionGetter, + RotaryPositionEmbedding2D, + SwiGLUFFNFused, +) +from depth_anything_3.model.reference_view_selector import ( + RefViewStrategy, + select_reference_view, + reorder_by_reference, + restore_original_order, +) +from depth_anything_3.utils.constants import THRESH_FOR_REF_SELECTION + +# logger = logging.getLogger("dinov2") + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def named_apply( + fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False +) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=1.0, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + alt_start=-1, + qknorm_start=-1, + rope_start=-1, + rope_freq=100, + plus_cam_token=False, + cat_token=True, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating + positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating + positional embeddings + """ + super().__init__() + self.patch_start_idx = 1 + norm_layer = nn.LayerNorm + self.num_features = self.embed_dim = ( + embed_dim # num_features for consistency with other models + ) + self.alt_start = alt_start + self.qknorm_start = qknorm_start + self.rope_start = rope_start + self.cat_token = cat_token + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + + self.patch_embed = embed_layer( + img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim + ) + num_patches = self.patch_embed.num_patches + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + if self.alt_start != -1: + self.camera_token = nn.Parameter(torch.randn(1, 2, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) + if num_register_tokens + else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + if self.rope_start != -1: + self.rope = RotaryPositionEmbedding2D(frequency=rope_freq) if rope_freq > 0 else None + self.position_getter = PositionGetter() if self.rope is not None else None + else: + self.rope = None + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + qk_norm=i >= qknorm_start if qknorm_start != -1 else False, + rope=self.rope if i >= rope_start and rope_start != -1 else None, + ) + for i in range(depth) + ] + self.blocks = nn.ModuleList(blocks_list) + self.norm = norm_layer(embed_dim) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if self.interpolate_offset: + # Historical kludge: add a small number to avoid floating point error in the + # interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using + # both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sx, sy) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (w0, h0) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + assert (w0, h0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_cls_token(self, B, S): + cls_token = self.cls_token.expand(B, S, -1) + cls_token = cls_token.reshape(B * S, -1, self.embed_dim) + return cls_token + + def prepare_tokens_with_masks(self, x, masks=None, cls_token=None, **kwargs): + B, S, nc, w, h = x.shape + x = rearrange(x, "b s c h w -> (b s) c h w") + x = self.patch_embed(x) + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + cls_token = self.prepare_cls_token(B, S) + x = torch.cat((cls_token, x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + x = rearrange(x, "(b s) n c -> b s n c", b=B, s=S) + return x + + def _prepare_rope(self, B, S, H, W, device): + pos = None + pos_nodiff = None + if self.rope is not None: + pos = self.position_getter( + B * S, H // self.patch_size, W // self.patch_size, device=device + ) + pos = rearrange(pos, "(b s) n c -> b s n c", b=B) + pos_nodiff = torch.zeros_like(pos).to(pos.dtype) + if self.patch_start_idx > 0: + pos = pos + 1 + pos_special = torch.zeros(B * S, self.patch_start_idx, 2).to(device).to(pos.dtype) + pos_special = rearrange(pos_special, "(b s) n c -> b s n c", b=B) + pos = torch.cat([pos_special, pos], dim=2) + pos_nodiff = pos_nodiff + 1 + pos_nodiff = torch.cat([pos_special, pos_nodiff], dim=2) + return pos, pos_nodiff + + def _get_intermediate_layers_not_chunked(self, x, n=1, export_feat_layers=[], **kwargs): + B, S, _, H, W = x.shape + x = self.prepare_tokens_with_masks(x) + output, total_block_len, aux_output = [], len(self.blocks), [] + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + pos, pos_nodiff = self._prepare_rope(B, S, H, W, x.device) + + for i, blk in enumerate(self.blocks): + if i < self.rope_start or self.rope is None: + g_pos, l_pos = None, None + else: + g_pos = pos_nodiff + l_pos = pos + + if self.alt_start != -1 and (i == self.alt_start - 1) and x.shape[1] >= THRESH_FOR_REF_SELECTION and kwargs.get("cam_token", None) is None: + # Select reference view using configured strategy + strategy = kwargs.get("ref_view_strategy", "saddle_balanced") + logger.info(f"Selecting reference view using strategy: {strategy}") + b_idx = select_reference_view(x, strategy=strategy) + # Reorder views to place reference view first + x = reorder_by_reference(x, b_idx) + local_x = reorder_by_reference(local_x, b_idx) + + if self.alt_start != -1 and i == self.alt_start: + if kwargs.get("cam_token", None) is not None: + logger.info("Using camera conditions provided by the user") + cam_token = kwargs.get("cam_token") + else: + ref_token = self.camera_token[:, :1].expand(B, -1, -1) + src_token = self.camera_token[:, 1:].expand(B, S - 1, -1) + cam_token = torch.cat([ref_token, src_token], dim=1) + x[:, :, 0] = cam_token + + if self.alt_start != -1 and i >= self.alt_start and i % 2 == 1: + x = self.process_attention( + x, blk, "global", pos=g_pos, attn_mask=kwargs.get("attn_mask", None) + ) + else: + x = self.process_attention(x, blk, "local", pos=l_pos) + local_x = x + + if i in blocks_to_take: + out_x = torch.cat([local_x, x], dim=-1) if self.cat_token else x + # Restore original view order if reordering was applied + if x.shape[1] >= THRESH_FOR_REF_SELECTION and self.alt_start != -1 and 'b_idx' in locals(): + out_x = restore_original_order(out_x, b_idx) + output.append((out_x[:, :, 0], out_x)) + if i in export_feat_layers: + aux_output.append(x) + return output, aux_output + + def process_attention(self, x, block, attn_type="global", pos=None, attn_mask=None): + b, s, n = x.shape[:3] + if attn_type == "local": + x = rearrange(x, "b s n c -> (b s) n c") + if pos is not None: + pos = rearrange(pos, "b s n c -> (b s) n c") + elif attn_type == "global": + x = rearrange(x, "b s n c -> b (s n) c") + if pos is not None: + pos = rearrange(pos, "b s n c -> b (s n) c") + else: + raise ValueError(f"Invalid attention type: {attn_type}") + + x = block(x, pos=pos, attn_mask=attn_mask) + + if attn_type == "local": + x = rearrange(x, "(b s) n c -> b s n c", b=b, s=s) + elif attn_type == "global": + x = rearrange(x, "b (s n) c -> b s n c", b=b, s=s) + return x + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + export_feat_layers: List[int] = [], + **kwargs, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + outputs, aux_outputs = self._get_intermediate_layers_not_chunked( + x, n, export_feat_layers=export_feat_layers, **kwargs + ) + camera_tokens = [out[0] for out in outputs] + if outputs[0][1].shape[-1] == self.embed_dim: + outputs = [self.norm(out[1]) for out in outputs] + elif outputs[0][1].shape[-1] == (self.embed_dim * 2): + outputs = [ + torch.cat( + [out[1][..., : self.embed_dim], self.norm(out[1][..., self.embed_dim :])], + dim=-1, + ) + for out in outputs + ] + else: + raise ValueError(f"Invalid output shape: {outputs[0][1].shape}") + aux_outputs = [self.norm(out) for out in aux_outputs] + outputs = [out[..., 1 + self.num_register_tokens :, :] for out in outputs] + aux_outputs = [out[..., 1 + self.num_register_tokens :, :] for out in aux_outputs] + return tuple(zip(outputs, camera_tokens)), aux_outputs + + +def vit_small(patch_size=16, num_register_tokens=0, depth=12, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=depth, + num_heads=6, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, depth=12, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=depth, + num_heads=12, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, depth=24, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=depth, + num_heads=16, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, depth=40, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=depth, + num_heads=24, + mlp_ratio=4, + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model diff --git a/depth_anything_3/model/dpt.py b/depth_anything_3/model/dpt.py new file mode 100644 index 0000000000000000000000000000000000000000..337a8a964a7f96fae997c8b12e3ae13e99dbaa58 --- /dev/null +++ b/depth_anything_3/model/dpt.py @@ -0,0 +1,458 @@ +# flake8: noqa E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict as TyDict +from typing import List, Sequence, Tuple +import torch +import torch.nn as nn +from addict import Dict +from einops import rearrange + +from depth_anything_3.model.utils.head_utils import ( + Permute, + create_uv_grid, + custom_interpolate, + position_grid_to_embed, +) + + +class DPT(nn.Module): + """ + DPT for dense prediction (main head + optional sky head, sky always 1 channel). + + Returns: + - Main head: + * If output_dim>1: { head_name, f"{head_name}_conf" } + * If output_dim==1: { head_name } + - Sky head (if use_sky_head=True): { sky_name } # [B, S, 1, H/down_ratio, W/down_ratio] + """ + + def __init__( + self, + dim_in: int, + *, + patch_size: int = 14, + output_dim: int = 1, + activation: str = "exp", + conf_activation: str = "expp1", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = False, + down_ratio: int = 1, + head_name: str = "depth", + # ---- sky head (fixed 1 channel) ---- + use_sky_head: bool = True, + sky_name: str = "sky", + sky_activation: str = "relu", # 'sigmoid' / 'relu' / 'linear' + use_ln_for_heads: bool = False, # If needed, apply LayerNorm on intermediate features of both heads + norm_type: str = "idt", # use to match legacy GS-DPT head, "idt" / "layer" + fusion_block_inplace: bool = False, + ) -> None: + super().__init__() + + # -------------------- configuration -------------------- + self.patch_size = patch_size + self.activation = activation + self.conf_activation = conf_activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + + # Names + self.head_main = head_name + self.sky_name = sky_name + + # Main head: output dimension and confidence switch + self.out_dim = output_dim + self.has_conf = output_dim > 1 + + # Sky head parameters (always 1 channel) + self.use_sky_head = use_sky_head + self.sky_activation = sky_activation + + # Fixed 4 intermediate outputs + self.intermediate_layer_idx: Tuple[int, int, int, int] = (0, 1, 2, 3) + + # -------------------- token pre-norm + per-stage projection -------------------- + if norm_type == "layer": + self.norm = nn.LayerNorm(dim_in) + elif norm_type == "idt": + self.norm = nn.Identity() + else: + raise Exception(f"Unknown norm_type {norm_type}, should be 'layer' or 'idt'.") + self.projects = nn.ModuleList( + [nn.Conv2d(dim_in, oc, kernel_size=1, stride=1, padding=0) for oc in out_channels] + ) + + # -------------------- Spatial re-size (align to common scale before fusion) -------------------- + # Design consistent with original: relative to patch grid (x4, x2, x1, /2) + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + out_channels[0], out_channels[0], kernel_size=4, stride=4, padding=0 + ), + nn.ConvTranspose2d( + out_channels[1], out_channels[1], kernel_size=2, stride=2, padding=0 + ), + nn.Identity(), + nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1), + ] + ) + + # -------------------- scratch: stage adapters + main fusion chain -------------------- + self.scratch = _make_scratch(list(out_channels), features, expand=False) + + # Main fusion chain + self.scratch.refinenet1 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet2 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet3 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet4 = _make_fusion_block( + features, has_residual=False, inplace=fusion_block_inplace + ) + + # Heads (shared neck1; then split into two heads) + head_features_1 = features + head_features_2 = 32 + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1 + ) + + ln_seq = ( + [Permute((0, 2, 3, 1)), nn.LayerNorm(head_features_2), Permute((0, 3, 1, 2))] + if use_ln_for_heads + else [] + ) + + # Main head + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0), + ) + + # Sky head (fixed 1 channel) + if self.use_sky_head: + self.scratch.sky_output_conv2 = nn.Sequential( + nn.Conv2d( + head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1 + ), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), + ) + + # ------------------------------------------------------------------------- + # Public forward (supports frame chunking to save memory) + # ------------------------------------------------------------------------- + def forward( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + chunk_size: int = 8, + **kwargs, + ) -> Dict: + """ + Args: + feats: List of 4 entries, each entry is a tensor like [B, S, T, C] (or the 0th element of tuple/list is that tensor). + H, W: Original image dimensions + patch_start_idx: Starting index of patch tokens in sequence (for cropping non-patch tokens) + chunk_size: Chunk size along time dimension S + + Returns: + Dict[str, Tensor] + """ + B, S, N, C = feats[0][0].shape + feats = [feat[0].reshape(B * S, N, C) for feat in feats] + + # update image info, used by the GS-DPT head + extra_kwargs = {} + if "images" in kwargs: + extra_kwargs.update({"images": rearrange(kwargs["images"], "B S ... -> (B S) ...")}) + + if chunk_size is None or chunk_size >= S: + out_dict = self._forward_impl(feats, H, W, patch_start_idx, **extra_kwargs) + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + out_dicts: List[TyDict[str, torch.Tensor]] = [] + for s0 in range(0, S, chunk_size): + s1 = min(s0 + chunk_size, S) + kw = {} + if "images" in extra_kwargs: + kw.update({"images": extra_kwargs["images"][s0:s1]}) + out_dicts.append( + self._forward_impl([f[s0:s1] for f in feats], H, W, patch_start_idx, **kw) + ) + out_dict = {k: torch.cat([od[k] for od in out_dicts], dim=0) for k in out_dicts[0].keys()} + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + ) -> TyDict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] # [B*S, N_patch, C] + x = self.norm(x) + # permute -> contiguous before reshape to keep conv input contiguous + x = x.permute(0, 2, 1).contiguous().reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # Align scale + resized_feats.append(x) + + # 2) Fusion pyramid (main branch only) + fused = self._fuse(resized_feats) + + # 3) Upsample to target resolution, optionally add position encoding again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused = self.scratch.output_conv1(fused) + fused = custom_interpolate(fused, (h_out, w_out), mode="bilinear", align_corners=True) + if self.pos_embed: + fused = self._add_pos_embed(fused, W, H) + + # 4) Shared neck1 + feat = fused + + # 5) Main head: logits -> activation + main_logits = self.scratch.output_conv2(feat) + outs: TyDict[str, torch.Tensor] = {} + if self.has_conf: + fmap = main_logits.permute(0, 2, 3, 1) + pred = self._apply_activation_single(fmap[..., :-1], self.activation) + conf = self._apply_activation_single(fmap[..., -1], self.conf_activation) + outs[self.head_main] = pred.squeeze(1) + outs[f"{self.head_main}_conf"] = conf.squeeze(1) + else: + outs[self.head_main] = self._apply_activation_single( + main_logits, self.activation + ).squeeze(1) + + # 6) Sky head (fixed 1 channel) + if self.use_sky_head: + sky_logits = self.scratch.sky_output_conv2(feat) + outs[self.sky_name] = self._apply_sky_activation(sky_logits).squeeze(1) + + return outs + + # ------------------------------------------------------------------------- + # Subroutines + # ------------------------------------------------------------------------- + def _fuse(self, feats: List[torch.Tensor]) -> torch.Tensor: + """ + 4-layer top-down fusion, returns finest scale features (after fusion, before neck1). + """ + l1, l2, l3, l4 = feats + + l1_rn = self.scratch.layer1_rn(l1) + l2_rn = self.scratch.layer2_rn(l2) + l3_rn = self.scratch.layer3_rn(l3) + l4_rn = self.scratch.layer4_rn(l4) + + # 4 -> 3 -> 2 -> 1 + out = self.scratch.refinenet4(l4_rn, size=l3_rn.shape[2:]) + out = self.scratch.refinenet3(out, l3_rn, size=l2_rn.shape[2:]) + out = self.scratch.refinenet2(out, l2_rn, size=l1_rn.shape[2:]) + out = self.scratch.refinenet1(out, l1_rn) + return out + + def _apply_activation_single( + self, x: torch.Tensor, activation: str = "linear" + ) -> torch.Tensor: + """ + Apply activation to single channel output, maintaining semantic consistency with value branch in multi-channel case. + Supports: exp / relu / sigmoid / softplus / tanh / linear / expp1 + """ + act = activation.lower() if isinstance(activation, str) else activation + if act == "exp": + return torch.exp(x) + if act == "expp1": + return torch.exp(x) + 1 + if act == "expm1": + return torch.expm1(x) + if act == "relu": + return torch.relu(x) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "softplus": + return torch.nn.functional.softplus(x) + if act == "tanh": + return torch.tanh(x) + # Default linear + return x + + def _apply_sky_activation(self, x: torch.Tensor) -> torch.Tensor: + """ + Sky head activation (fixed 1 channel): + * 'sigmoid' -> Sigmoid probability map + * 'relu' -> ReLU positive domain output + * 'linear' -> Original value (logits) + """ + act = ( + self.sky_activation.lower() + if isinstance(self.sky_activation, str) + else self.sky_activation + ) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "relu": + return torch.relu(x) + # 'linear' + return x + + def _add_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor: + """Simple UV position encoding directly added to feature map.""" + pw, ph = x.shape[-1], x.shape[-2] + pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype, device=x.device) + pe = position_grid_to_embed(pe, x.shape[1]) * ratio + pe = pe.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1) + return x + pe + + +# ----------------------------------------------------------------------------- +# Building blocks (preserved, consistent with original) +# ----------------------------------------------------------------------------- +def _make_fusion_block( + features: int, + size: Tuple[int, int] = None, + has_residual: bool = True, + groups: int = 1, + inplace: bool = False, +) -> nn.Module: + return FeatureFusionBlock( + features=features, + activation=nn.ReLU(inplace=inplace), + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=size, + has_residual=has_residual, + groups=groups, + ) + + +def _make_scratch( + in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False +) -> nn.Module: + scratch = nn.Module() + # Optional expansion by stage + c1 = out_shape + c2 = out_shape * (2 if expand else 1) + c3 = out_shape * (4 if expand else 1) + c4 = out_shape * (8 if expand else 1) + + scratch.layer1_rn = nn.Conv2d(in_shape[0], c1, 3, 1, 1, bias=False, groups=groups) + scratch.layer2_rn = nn.Conv2d(in_shape[1], c2, 3, 1, 1, bias=False, groups=groups) + scratch.layer3_rn = nn.Conv2d(in_shape[2], c3, 3, 1, 1, bias=False, groups=groups) + scratch.layer4_rn = nn.Conv2d(in_shape[3], c4, 3, 1, 1, bias=False, groups=groups) + return scratch + + +class ResidualConvUnit(nn.Module): + """Lightweight residual convolution block for fusion""" + + def __init__(self, features: int, activation: nn.Module, bn: bool, groups: int = 1) -> None: + super().__init__() + self.bn = bn + self.groups = groups + self.conv1 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) + self.conv2 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) + self.norm1 = None + self.norm2 = None + self.activation = activation + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + out = self.activation(x) + out = self.conv1(out) + if self.norm1 is not None: + out = self.norm1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.norm2 is not None: + out = self.norm2(out) + + return self.skip_add.add(out, x) + + +class FeatureFusionBlock(nn.Module): + """Top-down fusion block: (optional) residual merge + upsampling + 1x1 contraction""" + + def __init__( + self, + features: int, + activation: nn.Module, + deconv: bool = False, + bn: bool = False, + expand: bool = False, + align_corners: bool = True, + size: Tuple[int, int] = None, + has_residual: bool = True, + groups: int = 1, + ) -> None: + super().__init__() + self.align_corners = align_corners + self.size = size + self.has_residual = has_residual + + self.resConfUnit1 = ( + ResidualConvUnit(features, activation, bn, groups=groups) if has_residual else None + ) + self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=groups) + + out_features = (features // 2) if expand else features + self.out_conv = nn.Conv2d(features, out_features, 1, 1, 0, bias=True, groups=groups) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, *xs: torch.Tensor, size: Tuple[int, int] = None) -> torch.Tensor: # type: ignore[override] + """ + xs: + - xs[0]: Top branch input + - xs[1]: Lateral input (can do residual addition with top branch) + """ + y = xs[0] + if self.has_residual and len(xs) > 1 and self.resConfUnit1 is not None: + y = self.skip_add.add(y, self.resConfUnit1(xs[1])) + + y = self.resConfUnit2(y) + + # Upsampling + if (size is None) and (self.size is None): + up_kwargs = {"scale_factor": 2} + elif size is None: + up_kwargs = {"size": self.size} + else: + up_kwargs = {"size": size} + + y = custom_interpolate(y, **up_kwargs, mode="bilinear", align_corners=self.align_corners) + y = self.out_conv(y) + return y diff --git a/depth_anything_3/model/dualdpt.py b/depth_anything_3/model/dualdpt.py new file mode 100644 index 0000000000000000000000000000000000000000..229e1a990472de96f003a4a2d75fdf4ca94b0a3a --- /dev/null +++ b/depth_anything_3/model/dualdpt.py @@ -0,0 +1,488 @@ +# flake8: noqa E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Sequence, Tuple +import torch +import torch.nn as nn +from addict import Dict + +from depth_anything_3.model.dpt import _make_fusion_block, _make_scratch +from depth_anything_3.model.utils.head_utils import ( + Permute, + create_uv_grid, + custom_interpolate, + position_grid_to_embed, +) + + +class DualDPT(nn.Module): + """ + Dual-head DPT for dense prediction with an always-on auxiliary head. + + Architectural notes: + - Sky/object branches are removed. + - `intermediate_layer_idx` is fixed to (0, 1, 2, 3). + - Auxiliary head has its **own** fusion blocks (no fusion_inplace / no sharing). + - Auxiliary head is internally multi-level; **only the final level** is returned. + - Returns a **dict** with keys from `head_names`, e.g.: + { main_name, f"{main_name}_conf", aux_name, f"{aux_name}_conf" } + - `feature_only` is fixed to False. + """ + + def __init__( + self, + dim_in: int, + *, + patch_size: int = 14, + output_dim: int = 2, + activation: str = "exp", + conf_activation: str = "expp1", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = True, + down_ratio: int = 1, + aux_pyramid_levels: int = 4, + aux_out1_conv_num: int = 5, + head_names: Tuple[str, str] = ("depth", "ray"), + ) -> None: + super().__init__() + + # -------------------- configuration -------------------- + self.patch_size = patch_size + self.activation = activation + self.conf_activation = conf_activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + + self.aux_levels = aux_pyramid_levels + self.aux_out1_conv_num = aux_out1_conv_num + + # names ONLY come from config (no hard-coded strings elsewhere) + self.head_main, self.head_aux = head_names + + # Always expect 4 scales; enforce intermediate idx = (0, 1, 2, 3) + self.intermediate_layer_idx: Tuple[int, int, int, int] = (0, 1, 2, 3) + + # -------------------- token pre-norm + per-stage projection -------------------- + self.norm = nn.LayerNorm(dim_in) + self.projects = nn.ModuleList( + [nn.Conv2d(dim_in, oc, kernel_size=1, stride=1, padding=0) for oc in out_channels] + ) + + # -------------------- spatial re-sizers (align to common scale before fusion) -------------------- + # design: stage strides (x4, x2, x1, /2) relative to patch grid to align to a common pivot scale + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + out_channels[0], out_channels[0], kernel_size=4, stride=4, padding=0 + ), + nn.ConvTranspose2d( + out_channels[1], out_channels[1], kernel_size=2, stride=2, padding=0 + ), + nn.Identity(), + nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1), + ] + ) + + # -------------------- scratch: stage adapters + fusion (main & aux are separate) -------------------- + self.scratch = _make_scratch(list(out_channels), features, expand=False) + + # Main fusion chain (independent) + self.scratch.refinenet1 = _make_fusion_block(features) + self.scratch.refinenet2 = _make_fusion_block(features) + self.scratch.refinenet3 = _make_fusion_block(features) + self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False) + + # Primary head neck + head (independent) + head_features_1 = features + head_features_2 = 32 + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1 + ) + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0), + ) + + # Auxiliary fusion chain (completely separate; no sharing, i.e., "fusion_inplace=False") + self.scratch.refinenet1_aux = _make_fusion_block(features) + self.scratch.refinenet2_aux = _make_fusion_block(features) + self.scratch.refinenet3_aux = _make_fusion_block(features) + self.scratch.refinenet4_aux = _make_fusion_block(features, has_residual=False) + + # Aux pre-head per level (we will only *return final level*) + self.scratch.output_conv1_aux = nn.ModuleList( + [self._make_aux_out1_block(head_features_1) for _ in range(self.aux_levels)] + ) + + # Aux final projection per level + use_ln = True + ln_seq = ( + [Permute((0, 2, 3, 1)), nn.LayerNorm(head_features_2), Permute((0, 3, 1, 2))] + if use_ln + else [] + ) + self.scratch.output_conv2_aux = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d( + head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1 + ), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, 7, kernel_size=1, stride=1, padding=0), + ) + for _ in range(self.aux_levels) + ] + ) + + # ------------------------------------------------------------------------- + # Public forward (supports frame chunking for memory) + # ------------------------------------------------------------------------- + + def forward( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + chunk_size: int = 8, + ) -> Dict[str, torch.Tensor]: + """ + Args: + aggregated_tokens_list: List of 4 tensors [B, S, T, C] from transformer. + images: [B, S, 3, H, W], in [0, 1]. + patch_start_idx: Patch-token start in the token sequence (to drop non-patch tokens). + frames_chunk_size: Optional chunking along S for memory. + + Returns: + Dict[str, Tensor] with keys based on `head_names`, e.g.: + self.head_main, f"{self.head_main}_conf", + self.head_aux, f"{self.head_aux}_conf" + Shapes: + main: [B, S, out_dim, H/down_ratio, W/down_ratio] + main_cf: [B, S, 1, H/down_ratio, W/down_ratio] + aux: [B, S, 7, H/down_ratio, W/down_ratio] + aux_cf: [B, S, 1, H/down_ratio, W/down_ratio] + """ + B, S, N, C = feats[0][0].shape + feats = [feat[0].reshape(B * S, N, C) for feat in feats] + if chunk_size is None or chunk_size >= S: + out_dict = self._forward_impl(feats, H, W, patch_start_idx) + out_dict = {k: v.reshape(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + out_dicts = [] + for s0 in range(0, B * S, chunk_size): + s1 = min(s0 + chunk_size, B * S) + out_dict = self._forward_impl( + [feat[s0:s1] for feat in feats], + H, + W, + patch_start_idx, + ) + out_dicts.append(out_dict) + out_dict = { + k: torch.cat([out_dict[k] for out_dict in out_dicts], dim=0) + for k in out_dicts[0].keys() + } + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + ) -> Dict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] + x = self.norm(x) + x = x.permute(0, 2, 1).reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # align scales + resized_feats.append(x) + + # 2) Fuse pyramid (main & aux are completely independent) + fused_main, fused_aux_pyr = self._fuse(resized_feats) + + # 3) Upsample to target resolution and (optional) add pos-embed again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused_main = custom_interpolate( + fused_main, (h_out, w_out), mode="bilinear", align_corners=True + ) + if self.pos_embed: + fused_main = self._add_pos_embed(fused_main, W, H) + + # Primary head: conv1 -> conv2 -> activate + # fused_main = self.scratch.output_conv1(fused_main) + main_logits = self.scratch.output_conv2(fused_main) + fmap = main_logits.permute(0, 2, 3, 1) + main_pred = self._apply_activation_single(fmap[..., :-1], self.activation) + main_conf = self._apply_activation_single(fmap[..., -1], self.conf_activation) + + # Auxiliary head (multi-level inside) -> only last level returned (after activation) + last_aux = fused_aux_pyr[-1] + if self.pos_embed: + last_aux = self._add_pos_embed(last_aux, W, H) + # neck (per-level pre-conv) then final projection (only for last level) + # last_aux = self.scratch.output_conv1_aux[-1](last_aux) + last_aux_logits = self.scratch.output_conv2_aux[-1](last_aux) + fmap_last = last_aux_logits.permute(0, 2, 3, 1) + aux_pred = self._apply_activation_single(fmap_last[..., :-1], "linear") + aux_conf = self._apply_activation_single(fmap_last[..., -1], self.conf_activation) + return { + self.head_main: main_pred.squeeze(-1), + f"{self.head_main}_conf": main_conf, + self.head_aux: aux_pred, + f"{self.head_aux}_conf": aux_conf, + } + + # ------------------------------------------------------------------------- + # Subroutines + # ------------------------------------------------------------------------- + + def _fuse(self, feats: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """ + Feature pyramid fusion. + Returns: + fused_main: Tensor at finest scale (after refinenet1) + aux_pyr: List of aux tensors at each level (pre out_conv1_aux) + """ + l1, l2, l3, l4 = feats + + l1_rn = self.scratch.layer1_rn(l1) + l2_rn = self.scratch.layer2_rn(l2) + l3_rn = self.scratch.layer3_rn(l3) + l4_rn = self.scratch.layer4_rn(l4) + + # level 4 -> 3 + out = self.scratch.refinenet4(l4_rn, size=l3_rn.shape[2:]) + aux_out = self.scratch.refinenet4_aux(l4_rn, size=l3_rn.shape[2:]) + aux_list: List[torch.Tensor] = [] + if self.aux_levels >= 4: + aux_list.append(aux_out) + + # level 3 -> 2 + out = self.scratch.refinenet3(out, l3_rn, size=l2_rn.shape[2:]) + aux_out = self.scratch.refinenet3_aux(aux_out, l3_rn, size=l2_rn.shape[2:]) + if self.aux_levels >= 3: + aux_list.append(aux_out) + + # level 2 -> 1 + out = self.scratch.refinenet2(out, l2_rn, size=l1_rn.shape[2:]) + aux_out = self.scratch.refinenet2_aux(aux_out, l2_rn, size=l1_rn.shape[2:]) + if self.aux_levels >= 2: + aux_list.append(aux_out) + + # level 1 (final) + out = self.scratch.refinenet1(out, l1_rn) + aux_out = self.scratch.refinenet1_aux(aux_out, l1_rn) + aux_list.append(aux_out) + + out = self.scratch.output_conv1(out) + aux_list = [self.scratch.output_conv1_aux[i](aux) for i, aux in enumerate(aux_list)] + + return out, aux_list + + def _add_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor: + """Simple UV positional embedding added to feature maps.""" + pw, ph = x.shape[-1], x.shape[-2] + pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype, device=x.device) + pe = position_grid_to_embed(pe, x.shape[1]) * ratio + pe = pe.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1) + return x + pe + + def _make_aux_out1_block(self, in_ch: int) -> nn.Sequential: + """Factory for the aux pre-head stack before the final 1x1 projection.""" + if self.aux_out1_conv_num == 5: + return nn.Sequential( + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + ) + if self.aux_out1_conv_num == 3: + return nn.Sequential( + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + ) + if self.aux_out1_conv_num == 1: + return nn.Sequential(nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1)) + raise ValueError(f"aux_out1_conv_num {self.aux_out1_conv_num} not supported") + + def _apply_activation_single( + self, x: torch.Tensor, activation: str = "linear" + ) -> torch.Tensor: + """ + Apply activation to single channel output, maintaining semantic consistency with value branch in multi-channel case. + Supports: exp / relu / sigmoid / softplus / tanh / linear / expp1 + """ + act = activation.lower() if isinstance(activation, str) else activation + if act == "exp": + return torch.exp(x) + if act == "expm1": + return torch.expm1(x) + if act == "expp1": + return torch.exp(x) + 1 + if act == "relu": + return torch.relu(x) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "softplus": + return torch.nn.functional.softplus(x) + if act == "tanh": + return torch.tanh(x) + # Default linear + return x + + +# # ----------------------------------------------------------------------------- +# # Building blocks (tidy) +# # ----------------------------------------------------------------------------- + + +# def _make_fusion_block( +# features: int, +# size: Tuple[int, int] = None, +# has_residual: bool = True, +# groups: int = 1, +# inplace: bool = False, # <- activation uses inplace=True by default; not related to "fusion_inplace" +# ) -> nn.Module: +# return FeatureFusionBlock( +# features=features, +# activation=nn.ReLU(inplace=inplace), +# deconv=False, +# bn=False, +# expand=False, +# align_corners=True, +# size=size, +# has_residual=has_residual, +# groups=groups, +# ) + + +# def _make_scratch( +# in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False +# ) -> nn.Module: +# scratch = nn.Module() +# # optionally expand widths by stage +# c1 = out_shape +# c2 = out_shape * (2 if expand else 1) +# c3 = out_shape * (4 if expand else 1) +# c4 = out_shape * (8 if expand else 1) + +# scratch.layer1_rn = nn.Conv2d(in_shape[0], c1, 3, 1, 1, bias=False, groups=groups) +# scratch.layer2_rn = nn.Conv2d(in_shape[1], c2, 3, 1, 1, bias=False, groups=groups) +# scratch.layer3_rn = nn.Conv2d(in_shape[2], c3, 3, 1, 1, bias=False, groups=groups) +# scratch.layer4_rn = nn.Conv2d(in_shape[3], c4, 3, 1, 1, bias=False, groups=groups) +# return scratch + + +# class ResidualConvUnit(nn.Module): +# """Lightweight residual conv block used within fusion.""" + +# def __init__(self, features: int, activation: nn.Module, bn: bool, groups: int = 1) -> None: +# super().__init__() +# self.bn = bn +# self.groups = groups +# self.conv1 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) +# self.conv2 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) +# self.norm1 = None +# self.norm2 = None +# self.activation = activation +# self.skip_add = nn.quantized.FloatFunctional() + +# def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] +# out = self.activation(x) +# out = self.conv1(out) +# if self.norm1 is not None: +# out = self.norm1(out) + +# out = self.activation(out) +# out = self.conv2(out) +# if self.norm2 is not None: +# out = self.norm2(out) + +# return self.skip_add.add(out, x) + + +# class FeatureFusionBlock(nn.Module): +# """Top-down fusion block: (optional) residual merge + upsample + 1x1 shrink.""" + +# def __init__( +# self, +# features: int, +# activation: nn.Module, +# deconv: bool = False, +# bn: bool = False, +# expand: bool = False, +# align_corners: bool = True, +# size: Tuple[int, int] = None, +# has_residual: bool = True, +# groups: int = 1, +# ) -> None: +# super().__init__() +# self.align_corners = align_corners +# self.size = size +# self.has_residual = has_residual + +# self.resConfUnit1 = ( +# ResidualConvUnit(features, activation, bn, groups=groups) if has_residual else None +# ) +# self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=groups) + +# out_features = (features // 2) if expand else features +# self.out_conv = nn.Conv2d(features, out_features, 1, 1, 0, bias=True, groups=groups) +# self.skip_add = nn.quantized.FloatFunctional() + +# def forward(self, *xs: torch.Tensor, size: Tuple[int, int] = None) -> torch.Tensor: # type: ignore[override] +# """ +# xs: +# - xs[0]: top input +# - xs[1]: (optional) lateral (to be added with residual) +# """ +# y = xs[0] +# if self.has_residual and len(xs) > 1 and self.resConfUnit1 is not None: +# y = self.skip_add.add(y, self.resConfUnit1(xs[1])) + +# y = self.resConfUnit2(y) + +# # upsample +# if (size is None) and (self.size is None): +# up_kwargs = {"scale_factor": 2} +# elif size is None: +# up_kwargs = {"size": self.size} +# else: +# up_kwargs = {"size": size} + +# y = custom_interpolate(y, **up_kwargs, mode="bilinear", align_corners=self.align_corners) +# y = self.out_conv(y) +# return y diff --git a/depth_anything_3/model/gs_adapter.py b/depth_anything_3/model/gs_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..1335ce387476c339fdd7a8b3cb5349dfce56a7e0 --- /dev/null +++ b/depth_anything_3/model/gs_adapter.py @@ -0,0 +1,200 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional +import torch +from einops import einsum, rearrange, repeat +from torch import nn + +from depth_anything_3.model.utils.transform import cam_quat_xyzw_to_world_quat_wxyz +from depth_anything_3.specs import Gaussians +from depth_anything_3.utils.geometry import affine_inverse, get_world_rays, sample_image_grid +from depth_anything_3.utils.pose_align import batch_align_poses_umeyama +from depth_anything_3.utils.sh_helpers import rotate_sh + + +class GaussianAdapter(nn.Module): + + def __init__( + self, + sh_degree: int = 0, + pred_color: bool = False, + pred_offset_depth: bool = False, + pred_offset_xy: bool = True, + gaussian_scale_min: float = 1e-5, + gaussian_scale_max: float = 30.0, + ): + super().__init__() + self.sh_degree = sh_degree + self.pred_color = pred_color + self.pred_offset_depth = pred_offset_depth + self.pred_offset_xy = pred_offset_xy + self.gaussian_scale_min = gaussian_scale_min + self.gaussian_scale_max = gaussian_scale_max + + # Create a mask for the spherical harmonics coefficients. This ensures that at + # initialization, the coefficients are biased towards having a large DC + # component and small view-dependent components. + if not pred_color: + self.register_buffer( + "sh_mask", + torch.ones((self.d_sh,), dtype=torch.float32), + persistent=False, + ) + for degree in range(1, sh_degree + 1): + self.sh_mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree + + def forward( + self, + extrinsics: torch.Tensor, # "*#batch 4 4" + intrinsics: torch.Tensor, # "*#batch 3 3" + depths: torch.Tensor, # "*#batch" + opacities: torch.Tensor, # "*#batch" | "*#batch _" + raw_gaussians: torch.Tensor, # "*#batch _" + image_shape: tuple[int, int], + eps: float = 1e-8, + gt_extrinsics: Optional[torch.Tensor] = None, # "*#batch 4 4" + **kwargs, + ) -> Gaussians: + device = extrinsics.device + dtype = raw_gaussians.dtype + H, W = image_shape + b, v = raw_gaussians.shape[:2] + + # get cam2worlds and intr_normed to adapt to 3DGS codebase + cam2worlds = affine_inverse(extrinsics) + intr_normed = intrinsics.clone().detach() + intr_normed[..., 0, :] /= W + intr_normed[..., 1, :] /= H + + # 1. compute 3DGS means + # 1.1) offset the predicted depth if needed + if self.pred_offset_depth: + gs_depths = depths + raw_gaussians[..., -1] + raw_gaussians = raw_gaussians[..., :-1] + else: + gs_depths = depths + # 1.2) align predicted poses with GT if needed + if gt_extrinsics is not None and not torch.equal(extrinsics, gt_extrinsics): + try: + _, _, pose_scales = batch_align_poses_umeyama( + gt_extrinsics.detach().float(), + extrinsics.detach().float(), + ) + except Exception: + pose_scales = torch.ones_like(extrinsics[:, 0, 0, 0]) + pose_scales = torch.clamp(pose_scales, min=1 / 3.0, max=3.0) + cam2worlds[:, :, :3, 3] = cam2worlds[:, :, :3, 3] * rearrange( + pose_scales, "b -> b () ()" + ) # [b, i, j] + gs_depths = gs_depths * rearrange(pose_scales, "b -> b () () ()") # [b, v, h, w] + # 1.3) casting xy in image space + xy_ray, _ = sample_image_grid((H, W), device) + xy_ray = xy_ray[None, None, ...].expand(b, v, -1, -1, -1) # b v h w xy + # offset xy if needed + if self.pred_offset_xy: + pixel_size = 1 / torch.tensor((W, H), dtype=xy_ray.dtype, device=device) + offset_xy = raw_gaussians[..., :2] + xy_ray = xy_ray + offset_xy * pixel_size + raw_gaussians = raw_gaussians[..., 2:] # skip the offset_xy + # 1.4) unproject depth + xy to world ray + origins, directions = get_world_rays( + xy_ray, + repeat(cam2worlds, "b v i j -> b v h w i j", h=H, w=W), + repeat(intr_normed, "b v i j -> b v h w i j", h=H, w=W), + ) + gs_means_world = origins + directions * gs_depths[..., None] + gs_means_world = rearrange(gs_means_world, "b v h w d -> b (v h w) d") + + # 2. compute other GS attributes + scales, rotations, sh = raw_gaussians.split((3, 4, 3 * self.d_sh), dim=-1) + + # 2.1) 3DGS scales + # make the scale invarient to resolution + scale_min = self.gaussian_scale_min + scale_max = self.gaussian_scale_max + scales = scale_min + (scale_max - scale_min) * scales.sigmoid() + pixel_size = 1 / torch.tensor((W, H), dtype=dtype, device=device) + multiplier = self.get_scale_multiplier(intr_normed, pixel_size) + gs_scales = scales * gs_depths[..., None] * multiplier[..., None, None, None] + gs_scales = rearrange(gs_scales, "b v h w d -> b (v h w) d") + + # 2.2) 3DGS quaternion (world space) + # due to historical issue, assume quaternion in order xyzw, not wxyz + # Normalize the quaternion features to yield a valid quaternion. + rotations = rotations / (rotations.norm(dim=-1, keepdim=True) + eps) + # rotate them to world space + cam_quat_xyzw = rearrange(rotations, "b v h w c -> b (v h w) c") + c2w_mat = repeat( + cam2worlds, + "b v i j -> b (v h w) i j", + h=H, + w=W, + ) + world_quat_wxyz = cam_quat_xyzw_to_world_quat_wxyz(cam_quat_xyzw, c2w_mat) + gs_rotations_world = world_quat_wxyz # b (v h w) c + + # 2.3) 3DGS color / SH coefficient (world space) + sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3) + if not self.pred_color: + sh = sh * self.sh_mask + + if self.pred_color or self.sh_degree == 0: + # predict pre-computed color or predict only DC band, no need to transform + gs_sh_world = sh + else: + gs_sh_world = rotate_sh(sh, cam2worlds[:, :, None, None, None, :3, :3]) + gs_sh_world = rearrange(gs_sh_world, "b v h w xyz d_sh -> b (v h w) xyz d_sh") + + # 2.4) 3DGS opacity + gs_opacities = rearrange(opacities, "b v h w ... -> b (v h w) ...") + + return Gaussians( + means=gs_means_world, + harmonics=gs_sh_world, + opacities=gs_opacities, + scales=gs_scales, + rotations=gs_rotations_world, + ) + + def get_scale_multiplier( + self, + intrinsics: torch.Tensor, # "*#batch 3 3" + pixel_size: torch.Tensor, # "*#batch 2" + multiplier: float = 0.1, + ) -> torch.Tensor: # " *batch" + xy_multipliers = multiplier * einsum( + intrinsics[..., :2, :2].float().inverse().to(intrinsics), + pixel_size, + "... i j, j -> ... i", + ) + return xy_multipliers.sum(dim=-1) + + @property + def d_sh(self) -> int: + return 1 if self.pred_color else (self.sh_degree + 1) ** 2 + + @property + def d_in(self) -> int: + # provided as reference to the gs_dpt output dim + raw_gs_dim = 0 + if self.pred_offset_xy: + raw_gs_dim += 2 + raw_gs_dim += 3 # scales + raw_gs_dim += 4 # quaternion + raw_gs_dim += 3 * self.d_sh # color + if self.pred_offset_depth: + raw_gs_dim += 1 + + return raw_gs_dim diff --git a/depth_anything_3/model/gsdpt.py b/depth_anything_3/model/gsdpt.py new file mode 100644 index 0000000000000000000000000000000000000000..70448b7235982a3b47c0c07f15edd4c38677e72e --- /dev/null +++ b/depth_anything_3/model/gsdpt.py @@ -0,0 +1,133 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict as TyDict +from typing import List, Sequence +import torch +import torch.nn as nn + +from depth_anything_3.model.dpt import DPT +from depth_anything_3.model.utils.head_utils import activate_head_gs, custom_interpolate + + +class GSDPT(DPT): + + def __init__( + self, + dim_in: int, + patch_size: int = 14, + output_dim: int = 4, + activation: str = "linear", + conf_activation: str = "sigmoid", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = True, + feature_only: bool = False, + down_ratio: int = 1, + conf_dim: int = 1, + norm_type: str = "idt", # use to match legacy GS-DPT head, "idt" / "layer" + fusion_block_inplace: bool = False, + ) -> None: + super().__init__( + dim_in=dim_in, + patch_size=patch_size, + output_dim=output_dim, + activation=activation, + conf_activation=conf_activation, + features=features, + out_channels=out_channels, + pos_embed=pos_embed, + down_ratio=down_ratio, + head_name="raw_gs", + use_sky_head=False, + norm_type=norm_type, + fusion_block_inplace=fusion_block_inplace, + ) + self.conf_dim = conf_dim + if conf_dim and conf_dim > 1: + assert ( + conf_activation == "linear" + ), "use linear prediction when using view-dependent opacity" + + merger_out_dim = features if feature_only else features // 2 + self.images_merger = nn.Sequential( + nn.Conv2d(3, merger_out_dim // 4, 3, 1, 1), # fewer channels first + nn.GELU(), + nn.Conv2d(merger_out_dim // 4, merger_out_dim // 2, 3, 1, 1), + nn.GELU(), + nn.Conv2d(merger_out_dim // 2, merger_out_dim, 3, 1, 1), + nn.GELU(), + ) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + images: torch.Tensor, + ) -> TyDict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] # [B*S, N_patch, C] + x = self.norm(x) + x = x.permute(0, 2, 1).reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # Align scale + resized_feats.append(x) + + # 2) Fusion pyramid (main branch only) + fused = self._fuse(resized_feats) + fused = self.scratch.output_conv1(fused) + + # 3) Upsample to target resolution, optionally add position encoding again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused = custom_interpolate(fused, (h_out, w_out), mode="bilinear", align_corners=True) + + # inject the image information here + fused = fused + self.images_merger(images) + + if self.pos_embed: + fused = self._add_pos_embed(fused, W, H) + + # 4) Shared neck1 + # feat = self.scratch.output_conv1(fused) + feat = fused + + # 5) Main head: logits -> activate_head or single channel activation + main_logits = self.scratch.output_conv2(feat) + outs: TyDict[str, torch.Tensor] = {} + if self.has_conf: + pred, conf = activate_head_gs( + main_logits, + activation=self.activation, + conf_activation=self.conf_activation, + conf_dim=self.conf_dim, + ) + outs[self.head_main] = pred.squeeze(1) + outs[f"{self.head_main}_conf"] = conf.squeeze(1) + else: + outs[self.head_main] = self._apply_activation_single(main_logits).squeeze(1) + + return outs diff --git a/depth_anything_3/model/reference_view_selector.py b/depth_anything_3/model/reference_view_selector.py new file mode 100644 index 0000000000000000000000000000000000000000..f406f5ff2e0dad5ce2d12f3a1d3919773fc21e66 --- /dev/null +++ b/depth_anything_3/model/reference_view_selector.py @@ -0,0 +1,223 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Reference View Selection Strategies + +This module provides different strategies for selecting a reference view +from multiple input views in multi-view depth estimation. +""" + +import torch +from typing import Literal + + +RefViewStrategy = Literal["first", "middle", "saddle_balanced", "saddle_sim_range"] + + +def select_reference_view( + x: torch.Tensor, + strategy: RefViewStrategy = "saddle_balanced", +) -> torch.Tensor: + """ + Select a reference view from multiple views using the specified strategy. + + Args: + x: Input tensor of shape (B, S, N, C) where + B = batch size + S = number of views + N = number of tokens + C = channel dimension + strategy: Selection strategy, one of: + - "first": Always select the first view + - "middle": Select the middle view + - "saddle_balanced": Select view with balanced features across multiple metrics + - "saddle_sim_range": Select view with largest similarity range + + Returns: + b_idx: Tensor of shape (B,) containing the selected view index for each batch + """ + B, S, N, C = x.shape + + # For single view, no reordering needed + if S <= 1: + return torch.zeros(B, dtype=torch.long, device=x.device) + + # Simple position-based strategies + if strategy == "first": + return torch.zeros(B, dtype=torch.long, device=x.device) + + elif strategy == "middle": + return torch.full((B,), S // 2, dtype=torch.long, device=x.device) + + # Feature-based strategies require normalized class tokens + # Extract and normalize class tokens (first token of each view) + img_class_feat = x[:, :, 0] / x[:, :, 0].norm(dim=-1, keepdim=True) # B S C + + if strategy == "saddle_balanced": + # Select view with balanced features across multiple metrics + # Compute similarity matrix + sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S + sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) + sim_score = sim_no_diag.sum(dim=-1) / (S - 1) # B S + + feat_norm = x[:, :, 0].norm(dim=-1) # B S + feat_var = img_class_feat.var(dim=-1) # B S + + # Normalize all metrics to [0, 1] + def normalize_metric(metric): + min_val = metric.min(dim=1, keepdim=True).values + max_val = metric.max(dim=1, keepdim=True).values + return (metric - min_val) / (max_val - min_val + 1e-8) + + sim_score_norm = normalize_metric(sim_score) + norm_norm = normalize_metric(feat_norm) + var_norm = normalize_metric(feat_var) + + # Select view closest to the median (0.5) across all metrics + balance_score = ( + (sim_score_norm - 0.5).abs() + + (norm_norm - 0.5).abs() + + (var_norm - 0.5).abs() + ) + b_idx = balance_score.argmin(dim=1) + + elif strategy == "saddle_sim_range": + # Select view with largest similarity range (max - min) + sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S + sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) + + sim_max = sim_no_diag.max(dim=-1).values # B S + sim_min = sim_no_diag.min(dim=-1).values # B S + sim_range = sim_max - sim_min + b_idx = sim_range.argmax(dim=1) + + else: + raise ValueError( + f"Unknown reference view selection strategy: {strategy}. " + f"Must be one of: 'first', 'middle', 'saddle_balanced', 'saddle_sim_range'" + ) + + return b_idx + + +def reorder_by_reference( + x: torch.Tensor, + b_idx: torch.Tensor, +) -> torch.Tensor: + """ + Reorder views to place the selected reference view first. + + Args: + x: Input tensor of shape (B, S, N, C) + b_idx: Reference view indices of shape (B,) + + Returns: + Reordered tensor with reference view at position 0 + + Example: + If b_idx = [2] and S = 5 (views [0,1,2,3,4]), + result order is [2,0,1,3,4] (ref_idx first, then others in order) + """ + B, S = x.shape[0], x.shape[1] + + # For single view, no reordering needed + if S <= 1: + return x + + # Create position indices: (B, S) where each row is [0, 1, 2, ..., S-1] + positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S + + # For each position, determine which original index it should take + # Position 0 gets ref_idx + # Position 1 to ref_idx gets indices 0 to ref_idx-1 + # Position ref_idx+1 to S-1 gets indices ref_idx+1 to S-1 + + b_idx_expanded = b_idx.unsqueeze(1) # B 1 + + # Create the reordering indices + # For positions 1 to ref_idx: map to indices 0 to ref_idx-1 (shift by -1) + # For positions > ref_idx: keep the same + reorder_indices = positions.clone() + reorder_indices = torch.where( + (positions > 0) & (positions <= b_idx_expanded), + positions - 1, + positions + ) + # Set position 0 to ref_idx + reorder_indices[:, 0] = b_idx + + # Gather using advanced indexing + batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 + x_reordered = x[batch_indices, reorder_indices] + + return x_reordered + + +def restore_original_order( + x: torch.Tensor, + b_idx: torch.Tensor, +) -> torch.Tensor: + """ + Restore original view order after processing. + + Args: + x: Reordered tensor of shape (B, S, ...) + b_idx: Original reference view indices of shape (B,) + + Returns: + Tensor with original view order restored + + Example: + If original order was [0, 1, 2, 3, 4] and b_idx=2, + reordered becomes [2, 0, 1, 3, 4] (reference at position 0), + restore should return [0, 1, 2, 3, 4] (original order). + """ + B, S = x.shape[0], x.shape[1] + + # For single view, no restoration needed + if S <= 1: + return x + + # Create target position indices: (B, S) where each row is [0, 1, 2, ..., S-1] + target_positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S + + # For each target position, determine which current position it comes from + # Target position 0 to ref_idx-1 <- Current position 1 to ref_idx (shift by +1) + # Target position ref_idx <- Current position 0 + # Target position ref_idx+1 to S-1 <- Current position ref_idx+1 to S-1 (no change) + + b_idx_expanded = b_idx.unsqueeze(1) # B 1 + + # Create the restore indices + restore_indices = torch.where( + target_positions < b_idx_expanded, + target_positions + 1, # Positions before ref_idx come from current position + 1 + target_positions # Positions after ref_idx stay the same + ) + # Target position = ref_idx comes from current position 0 + # Use scatter to set specific positions + restore_indices = torch.scatter( + restore_indices, + dim=1, + index=b_idx_expanded, + src=torch.zeros_like(b_idx_expanded) + ) + + # Gather using advanced indexing + batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 + x_restored = x[batch_indices, restore_indices] + + return x_restored + diff --git a/depth_anything_3/model/utils/attention.py b/depth_anything_3/model/utils/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..49c07a8c5e5c65fc8e0700aab9ae660552dba610 --- /dev/null +++ b/depth_anything_3/model/utils/attention.py @@ -0,0 +1,109 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa + +from typing import Callable, Optional, Union +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + qk_norm: bool = False, + rope=None, + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + # Debug breakpoint removed for production + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + q = self.rope(q, pos) if self.rope is not None else q + k = self.rope(k, pos) if self.rope is not None else k + x = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.attn_drop.p if self.training else 0.0, + attn_mask=attn_mask, + ) + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/depth_anything_3/model/utils/block.py b/depth_anything_3/model/utils/block.py new file mode 100644 index 0000000000000000000000000000000000000000..993fb4c0bdc02bd0976a6471adf1a15452d0f5c0 --- /dev/null +++ b/depth_anything_3/model/utils/block.py @@ -0,0 +1,81 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import Callable +from torch import Tensor, nn + +from .attention import Attention, LayerScale, Mlp + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + qk_norm: bool = False, + rope=None, + ) -> None: + super().__init__() + + self.norm1 = norm_layer(dim) + + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + qk_norm=qk_norm, + rope=rope, + ) + + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + + self.sample_drop_ratio = 0.0 # Equivalent to always having drop_path=0 + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + # drop_path is always 0, so always take the else branch + x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask) + x = x + ffn_residual_func(x) + return x diff --git a/depth_anything_3/model/utils/gs_renderer.py b/depth_anything_3/model/utils/gs_renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..91414d38e84cc030af5207682d06135eec72b125 --- /dev/null +++ b/depth_anything_3/model/utils/gs_renderer.py @@ -0,0 +1,340 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from math import isqrt +from typing import Literal, Optional +import torch +from einops import rearrange, repeat +from tqdm import tqdm + +from depth_anything_3.specs import Gaussians +from depth_anything_3.utils.camera_trj_helpers import ( + interpolate_extrinsics, + interpolate_intrinsics, + render_dolly_zoom_path, + render_stabilization_path, + render_wander_path, + render_wobble_inter_path, +) +from depth_anything_3.utils.geometry import affine_inverse, as_homogeneous, get_fov +from depth_anything_3.utils.logger import logger + +try: + from gsplat import rasterization +except ImportError: + logger.warn( + "Dependency `gsplat` is required for rendering 3DGS. " + "Install via: pip install git+https://github.com/nerfstudio-project/" + "gsplat.git@0b4dddf04cb687367602c01196913cde6a743d70" + ) + + +def render_3dgs( + extrinsics: torch.Tensor, # "batch_views 4 4", w2c + intrinsics: torch.Tensor, # "batch_views 3 3", normalized + image_shape: tuple[int, int], + gaussian: Gaussians, + background_color: Optional[torch.Tensor] = None, # "batch_views 3" + use_sh: bool = True, + num_view: int = 1, + color_mode: Literal["RGB+D", "RGB+ED"] = "RGB+D", + **kwargs, +) -> tuple[ + torch.Tensor, # "batch_views 3 height width" + torch.Tensor, # "batch_views height width" +]: + # extract gaussian params + gaussian_means = gaussian.means + gaussian_scales = gaussian.scales + gaussian_quats = gaussian.rotations + gaussian_opacities = gaussian.opacities + gaussian_sh_coefficients = gaussian.harmonics + b, _, _ = extrinsics.shape + + if background_color is None: + background_color = repeat(torch.tensor([0.0, 0.0, 0.0]), "c -> b c", b=b).to( + gaussian_sh_coefficients + ) + + if use_sh: + _, _, _, n = gaussian_sh_coefficients.shape + degree = isqrt(n) - 1 + shs = rearrange(gaussian_sh_coefficients, "b g xyz n -> b g n xyz").contiguous() + else: # use color + shs = ( + gaussian_sh_coefficients.squeeze(-1).sigmoid().contiguous() + ) # (b, g, c), normed to (0, 1) + + h, w = image_shape + + fov_x, fov_y = get_fov(intrinsics).unbind(dim=-1) + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + focal_length_x = w / (2 * tan_fov_x) + focal_length_y = h / (2 * tan_fov_y) + + view_matrix = extrinsics.float() + + all_images = [] + all_radii = [] + all_depths = [] + # render view in a batch based, each batch contains one scene + # assume the Gaussian parameters are originally repeated along the view dim + batch_scene = b // num_view + + def index_i_gs_attr(full_attr, idx): + # return rearrange(full_attr, "(b v) ... -> b v ...", v=num_view)[idx, 0] + return full_attr[idx] + + for i in range(batch_scene): + K = repeat( + torch.tensor( + [ + [0, 0, w / 2.0], + [0, 0, h / 2.0], + [0, 0, 1], + ] + ), + "i j -> v i j", + v=num_view, + ).to(gaussian_means) + K[:, 0, 0] = focal_length_x.reshape(batch_scene, num_view)[i] + K[:, 1, 1] = focal_length_y.reshape(batch_scene, num_view)[i] + + i_means = index_i_gs_attr(gaussian_means, i) # [N, 3] + i_scales = index_i_gs_attr(gaussian_scales, i) + i_quats = index_i_gs_attr(gaussian_quats, i) + i_opacities = index_i_gs_attr(gaussian_opacities, i) # [N,] + i_colors = index_i_gs_attr(shs, i) # [N, K, 3] + i_viewmats = rearrange(view_matrix, "(b v) ... -> b v ...", v=num_view)[i] # [v, 4, 4] + i_backgrounds = rearrange(background_color, "(b v) ... -> b v ...", v=num_view)[ + i + ] # [v, 3] + + render_colors, render_alphas, info = rasterization( + means=i_means, + quats=i_quats, # [N, 4] + scales=i_scales, # [N, 3] + opacities=i_opacities, + colors=i_colors, + viewmats=i_viewmats, # [v, 4, 4] + Ks=K, # [v, 3, 3] + backgrounds=i_backgrounds, + render_mode=color_mode, + width=w, + height=h, + packed=False, + sh_degree=degree if use_sh else None, + ) + depth = render_colors[..., -1].unbind(dim=0) + + image = rearrange(render_colors[..., :3], "v h w c -> v c h w").unbind(dim=0) + radii = info["radii"].unbind(dim=0) + try: + info["means2d"].retain_grad() # [1, N, 2] + except Exception: + pass + all_images.extend(image) + all_depths.extend(depth) + all_radii.extend(radii) + + return torch.stack(all_images), torch.stack(all_depths) + + +def run_renderer_in_chunk_w_trj_mode( + gaussians: Gaussians, + extrinsics: torch.Tensor, # world2cam, "batch view 4 4" | "batch view 3 4" + intrinsics: torch.Tensor, # unnormed intrinsics, "batch view 3 3" + image_shape: tuple[int, int], + chunk_size: Optional[int] = 8, + trj_mode: Literal[ + "original", + "smooth", + "interpolate", + "interpolate_smooth", + "wander", + "dolly_zoom", + "extend", + "wobble_inter", + ] = "smooth", + input_shape: Optional[tuple[int, int]] = None, + enable_tqdm: Optional[bool] = False, + **kwargs, +) -> tuple[ + torch.Tensor, # color, "batch view 3 height width" + torch.Tensor, # depth, "batch view height width" +]: + cam2world = affine_inverse(as_homogeneous(extrinsics)) + if input_shape is not None: + in_h, in_w = input_shape + else: + in_h, in_w = image_shape + intr_normed = intrinsics.clone().detach() + intr_normed[..., 0, :] /= in_w + intr_normed[..., 1, :] /= in_h + if extrinsics.shape[1] <= 1: + assert trj_mode in [ + "wander", + "dolly_zoom", + ], "Please set trj_mode to 'wander' or 'dolly_zoom' when n_views=1" + + def _smooth_trj_fn_batch(raw_c2ws, k_size=50): + try: + smooth_c2ws = torch.stack( + [render_stabilization_path(c2w_i, k_size) for c2w_i in raw_c2ws], + dim=0, + ) + except Exception as e: + print(f"[DEBUG] Path smoothing failed with error: {e}.") + smooth_c2ws = raw_c2ws + return smooth_c2ws + + # get rendered trj + if trj_mode == "original": + tgt_c2w = cam2world + tgt_intr = intr_normed + elif trj_mode == "smooth": + tgt_c2w = _smooth_trj_fn_batch(cam2world) + tgt_intr = intr_normed + elif trj_mode in ["interpolate", "interpolate_smooth", "extend"]: + inter_len = 8 + total_len = (cam2world.shape[1] - 1) * inter_len + if total_len > 24 * 18: # no more than 18s + inter_len = max(1, 24 * 10 // (cam2world.shape[1] - 1)) + if total_len < 24 * 2: # no less than 2s + inter_len = max(1, 24 * 2 // (cam2world.shape[1] - 1)) + + if inter_len > 2: + t = torch.linspace(0, 1, inter_len, dtype=torch.float32, device=cam2world.device) + t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + tgt_c2w_b = [] + tgt_intr_b = [] + for b_idx in range(cam2world.shape[0]): + tgt_c2w = [] + tgt_intr = [] + for cur_idx in range(cam2world.shape[1] - 1): + tgt_c2w.append( + interpolate_extrinsics( + cam2world[b_idx, cur_idx], cam2world[b_idx, cur_idx + 1], t + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_intr.append( + interpolate_intrinsics( + intr_normed[b_idx, cur_idx], intr_normed[b_idx, cur_idx + 1], t + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_c2w_b.append(torch.cat(tgt_c2w)) + tgt_intr_b.append(torch.cat(tgt_intr)) + tgt_c2w = torch.stack(tgt_c2w_b) # b v 4 4 + tgt_intr = torch.stack(tgt_intr_b) # b v 3 3 + else: + tgt_c2w = cam2world + tgt_intr = intr_normed + if trj_mode in ["interpolate_smooth", "extend"]: + tgt_c2w = _smooth_trj_fn_batch(tgt_c2w) + if trj_mode == "extend": + # apply dolly_zoom and wander in the middle frame + assert cam2world.shape[0] == 1, "extend only supports for batch_size=1 currently." + mid_idx = tgt_c2w.shape[1] // 2 + c2w_wd, intr_wd = render_wander_path( + tgt_c2w[0, mid_idx], + tgt_intr[0, mid_idx], + h=in_h, + w=in_w, + num_frames=max(36, min(60, mid_idx // 2)), + max_disp=24.0, + ) + c2w_dz, intr_dz = render_dolly_zoom_path( + tgt_c2w[0, mid_idx], + tgt_intr[0, mid_idx], + h=in_h, + w=in_w, + num_frames=max(36, min(60, mid_idx // 2)), + ) + tgt_c2w = torch.cat( + [ + tgt_c2w[:, :mid_idx], + c2w_wd.unsqueeze(0), + c2w_dz.unsqueeze(0), + tgt_c2w[:, mid_idx:], + ], + dim=1, + ) + tgt_intr = torch.cat( + [ + tgt_intr[:, :mid_idx], + intr_wd.unsqueeze(0), + intr_dz.unsqueeze(0), + tgt_intr[:, mid_idx:], + ], + dim=1, + ) + elif trj_mode in ["wander", "dolly_zoom"]: + if trj_mode == "wander": + render_fn = render_wander_path + extra_kwargs = {"max_disp": 24.0} + else: + render_fn = render_dolly_zoom_path + extra_kwargs = {"D_focus": 30.0, "max_disp": 2.0} + tgt_c2w = [] + tgt_intr = [] + for b_idx in range(cam2world.shape[0]): + c2w_i, intr_i = render_fn( + cam2world[b_idx, 0], intr_normed[b_idx, 0], h=in_h, w=in_w, **extra_kwargs + ) + tgt_c2w.append(c2w_i) + tgt_intr.append(intr_i) + tgt_c2w = torch.stack(tgt_c2w) + tgt_intr = torch.stack(tgt_intr) + elif trj_mode == "wobble_inter": + tgt_c2w, tgt_intr = render_wobble_inter_path( + cam2world=cam2world, + intr_normed=intr_normed, + inter_len=10, + n_skip=3, + ) + else: + raise Exception(f"trj mode [{trj_mode}] is not implemented.") + + _, v = tgt_c2w.shape[:2] + tgt_extr = affine_inverse(tgt_c2w) + if chunk_size is None: + chunk_size = v + chunk_size = min(v, chunk_size) + all_colors = [] + all_depths = [] + for chunk_idx in tqdm( + range(math.ceil(v / chunk_size)), + desc="Rendering novel views", + disable=(not enable_tqdm), + leave=False, + ): + s = int(chunk_idx * chunk_size) + e = int((chunk_idx + 1) * chunk_size) + cur_n_view = tgt_extr[:, s:e].shape[1] + color, depth = render_3dgs( + extrinsics=rearrange(tgt_extr[:, s:e], "b v ... -> (b v) ..."), # w2c + intrinsics=rearrange(tgt_intr[:, s:e], "b v ... -> (b v) ..."), # normed + image_shape=image_shape, + gaussian=gaussians, + num_view=cur_n_view, + **kwargs, + ) + all_colors.append(rearrange(color, "(b v) ... -> b v ...", v=cur_n_view)) + all_depths.append(rearrange(depth, "(b v) ... -> b v ...", v=cur_n_view)) + all_colors = torch.cat(all_colors, dim=1) + all_depths = torch.cat(all_depths, dim=1) + + return all_colors, all_depths diff --git a/depth_anything_3/model/utils/head_utils.py b/depth_anything_3/model/utils/head_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c1209582bc6e9b1658a0358df03f5ea79fe61daf --- /dev/null +++ b/depth_anything_3/model/utils/head_utils.py @@ -0,0 +1,230 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple, Union +import torch +import torch.nn as nn +import torch.nn.functional as F + +# ----------------------------------------------------------------------------- +# Activation functions +# ----------------------------------------------------------------------------- + + +def activate_head_gs(out, activation="norm_exp", conf_activation="expp1", conf_dim=None): + """ + Process network output to extract GS params and density values. + Density could be view-dependent as SH coefficient + + + Args: + out: Network output tensor (B, C, H, W) + activation: Activation type for 3D points + conf_activation: Activation type for confidence values + + Returns: + Tuple of (3D points tensor, confidence tensor) + """ + # Move channels from last dim to the 4th dimension => (B, H, W, C) + fmap = out.permute(0, 2, 3, 1) # B,H,W,C expected + + # Split into xyz (first C-1 channels) and confidence (last channel) + conf_dim = 1 if conf_dim is None else conf_dim + xyz = fmap[:, :, :, :-conf_dim] + conf = fmap[:, :, :, -1] if conf_dim == 1 else fmap[:, :, :, -conf_dim:] + + if activation == "norm_exp": + d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8) + xyz_normed = xyz / d + pts3d = xyz_normed * torch.expm1(d) + elif activation == "norm": + pts3d = xyz / xyz.norm(dim=-1, keepdim=True) + elif activation == "exp": + pts3d = torch.exp(xyz) + elif activation == "relu": + pts3d = F.relu(xyz) + elif activation == "sigmoid": + pts3d = torch.sigmoid(xyz) + elif activation == "linear": + pts3d = xyz + else: + raise ValueError(f"Unknown activation: {activation}") + + if conf_activation == "expp1": + conf_out = 1 + conf.exp() + elif conf_activation == "expp0": + conf_out = conf.exp() + elif conf_activation == "sigmoid": + conf_out = torch.sigmoid(conf) + elif conf_activation == "linear": + conf_out = conf + else: + raise ValueError(f"Unknown conf_activation: {conf_activation}") + + return pts3d, conf_out + + +# ----------------------------------------------------------------------------- +# Other utilities +# ----------------------------------------------------------------------------- + + +class Permute(nn.Module): + """nn.Module wrapper around Tensor.permute for cleaner nn.Sequential usage.""" + + dims: Tuple[int, ...] + + def __init__(self, dims: Tuple[int, ...]) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + return x.permute(*self.dims) + + +def position_grid_to_embed( + pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100 +) -> torch.Tensor: + """ + Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC) + + Args: + pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates + embed_dim: Output channel dimension for embeddings + + Returns: + Tensor of shape (H, W, embed_dim) with positional embeddings + """ + H, W, grid_dim = pos_grid.shape + assert grid_dim == 2 + pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2) + + # Process x and y coordinates separately + emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0=omega_0) # [1, H*W, D/2] + emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0=omega_0) # [1, H*W, D/2] + + # Combine and reshape + emb = torch.cat([emb_x, emb_y], dim=-1) # [1, H*W, D] + + return emb.view(H, W, embed_dim) # [H, W, D] + + +def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 100) -> torch.Tensor: + """ + This function generates a 1D positional embedding from a given grid using sine and cosine functions. # noqa + + Args: + - embed_dim: The embedding dimension. + - pos: The position to generate the embedding from. + + Returns: + - emb: The generated 1D positional embedding. + """ + assert embed_dim % 2 == 0 + omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device) + omega /= embed_dim / 2.0 + omega = 1.0 / omega_0**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = torch.sin(out) # (M, D/2) + emb_cos = torch.cos(out) # (M, D/2) + + emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D) + return emb.float() + + +# Inspired by https://github.com/microsoft/moge + + +def create_uv_grid( + width: int, + height: int, + aspect_ratio: float = None, + dtype: torch.dtype = None, + device: torch.device = None, +) -> torch.Tensor: + """ + Create a normalized UV grid of shape (width, height, 2). + + The grid spans horizontally and vertically according to an aspect ratio, + ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right + corner is at (x_span, y_span), normalized by the diagonal of the plane. + + Args: + width (int): Number of points horizontally. + height (int): Number of points vertically. + aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height. + dtype (torch.dtype, optional): Data type of the resulting tensor. + device (torch.device, optional): Device on which the tensor is created. + + Returns: + torch.Tensor: A (width, height, 2) tensor of UV coordinates. + """ + # Derive aspect ratio if not explicitly provided + if aspect_ratio is None: + aspect_ratio = float(width) / float(height) + + # Compute normalized spans for X and Y + diag_factor = (aspect_ratio**2 + 1.0) ** 0.5 + span_x = aspect_ratio / diag_factor + span_y = 1.0 / diag_factor + + # Establish the linspace boundaries + left_x = -span_x * (width - 1) / width + right_x = span_x * (width - 1) / width + top_y = -span_y * (height - 1) / height + bottom_y = span_y * (height - 1) / height + + # Generate 1D coordinates + x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device) + y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device) + + # Create 2D meshgrid (width x height) and stack into UV + uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy") + uv_grid = torch.stack((uu, vv), dim=-1) + + return uv_grid + + +# ----------------------------------------------------------------------------- +# Interpolation (safe interpolation, avoid INT_MAX overflow) +# ----------------------------------------------------------------------------- +def custom_interpolate( + x: torch.Tensor, + size: Union[Tuple[int, int], None] = None, + scale_factor: Union[float, None] = None, + mode: str = "bilinear", + align_corners: bool = True, +) -> torch.Tensor: + """ + Safe interpolation implementation to avoid INT_MAX overflow in torch.nn.functional.interpolate. + """ + if size is None: + assert scale_factor is not None, "Either size or scale_factor must be provided." + size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor)) + + INT_MAX = 1610612736 + total = size[0] * size[1] * x.shape[0] * x.shape[1] + + if total > INT_MAX: + chunks = torch.chunk(x, chunks=(total // INT_MAX) + 1, dim=0) + outs = [ + nn.functional.interpolate(c, size=size, mode=mode, align_corners=align_corners) + for c in chunks + ] + return torch.cat(outs, dim=0).contiguous() + + return nn.functional.interpolate(x, size=size, mode=mode, align_corners=align_corners) diff --git a/depth_anything_3/model/utils/transform.py b/depth_anything_3/model/utils/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..8d732b093e5ad1578bc0ba5eb0e31b1b69b766ed --- /dev/null +++ b/depth_anything_3/model/utils/transform.py @@ -0,0 +1,208 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn.functional as F + + +def extri_intri_to_pose_encoding( + extrinsics, + intrinsics, + image_size_hw=None, +): + """Convert camera extrinsics and intrinsics to a compact pose encoding.""" + + # extrinsics: BxSx3x4 + # intrinsics: BxSx3x3 + R = extrinsics[:, :, :3, :3] # BxSx3x3 + T = extrinsics[:, :, :3, 3] # BxSx3 + + quat = mat_to_quat(R) + # Note the order of h and w here + H, W = image_size_hw + fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) + fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0]) + pose_encoding = torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float() + + return pose_encoding + + +def pose_encoding_to_extri_intri( + pose_encoding, + image_size_hw=None, +): + """Convert a pose encoding back to camera extrinsics and intrinsics.""" + + T = pose_encoding[..., :3] + quat = pose_encoding[..., 3:7] + fov_h = pose_encoding[..., 7] + fov_w = pose_encoding[..., 8] + + R = quat_to_mat(quat) + extrinsics = torch.cat([R, T[..., None]], dim=-1) + + H, W = image_size_hw + fy = (H / 2.0) / torch.clamp(torch.tan(fov_h / 2.0), 1e-6) + fx = (W / 2.0) / torch.clamp(torch.tan(fov_w / 2.0), 1e-6) + intrinsics = torch.zeros(pose_encoding.shape[:2] + (3, 3), device=pose_encoding.device) + intrinsics[..., 0, 0] = fx + intrinsics[..., 1, 1] = fy + intrinsics[..., 0, 2] = W / 2 + intrinsics[..., 1, 2] = H / 2 + intrinsics[..., 2, 2] = 1.0 # Set the homogeneous coordinate to 1 + + return extrinsics, intrinsics + + +def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor: + """ + Quaternion Order: XYZW or say ijkr, scalar-last + + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + i, j, k, r = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + quat_by_rijk = torch.stack( + [ + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape( + batch_dim + (4,) + ) + + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def cam_quat_xyzw_to_world_quat_wxyz(cam_quat_xyzw, c2w): + # cam_quat_xyzw: (b, n, 4) in xyzw + # c2w: (b, n, 4, 4) + b, n = cam_quat_xyzw.shape[:2] + # 1. xyzw -> wxyz + cam_quat_wxyz = torch.cat( + [ + cam_quat_xyzw[..., 3:4], # w + cam_quat_xyzw[..., 0:1], # x + cam_quat_xyzw[..., 1:2], # y + cam_quat_xyzw[..., 2:3], # z + ], + dim=-1, + ) + # 2. Quaternion to matrix + cam_quat_wxyz_flat = cam_quat_wxyz.reshape(-1, 4) + rotmat_cam = quat_to_mat(cam_quat_wxyz_flat).reshape(b, n, 3, 3) + # 3. Transform to world space + rotmat_c2w = c2w[..., :3, :3] + rotmat_world = torch.matmul(rotmat_c2w, rotmat_cam) + # 4. Matrix to quaternion (wxyz) + rotmat_world_flat = rotmat_world.reshape(-1, 3, 3) + world_quat_wxyz_flat = mat_to_quat(rotmat_world_flat) + world_quat_wxyz = world_quat_wxyz_flat.reshape(b, n, 4) + return world_quat_wxyz diff --git a/depth_anything_3/registry.py b/depth_anything_3/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..96450717d696b395503fedfe93af812e975d2671 --- /dev/null +++ b/depth_anything_3/registry.py @@ -0,0 +1,50 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import OrderedDict +from pathlib import Path + + +def get_all_models() -> OrderedDict: + """ + Scans all YAML files in the configs directory and returns a sorted dictionary where: + - Keys are model names (YAML filenames without the .yaml extension) + - Values are absolute paths to the corresponding YAML files + """ + # Get path to the configs directory within the da3 package + # Works both in development and after pip installation + # configs_dir = files("depth_anything_3").joinpath("configs") + configs_dir = Path(__file__).resolve().parent / "configs" + + # Ensure path is a Path object for consistent cross-platform handling + configs_dir = Path(configs_dir) + + model_entries = [] + # Iterate through all items in the configs directory + for item in configs_dir.iterdir(): + # Filter for YAML files (excluding directories) + if item.is_file() and item.suffix == ".yaml": + # Extract model name (filename without .yaml extension) + model_name = item.stem + # Get absolute path (resolve() handles symlinks) + file_abs_path = str(item.resolve()) + model_entries.append((model_name, file_abs_path)) + + # Sort entries by model name and convert to OrderedDict + sorted_entries = sorted(model_entries, key=lambda x: x[0]) + return OrderedDict(sorted_entries) + + +# Global registry for external imports +MODEL_REGISTRY = get_all_models() diff --git a/depth_anything_3/services/__init__.py b/depth_anything_3/services/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..07ed8b6b96d68a3f3933dfa6651875f3f88b0aef --- /dev/null +++ b/depth_anything_3/services/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Services module for Depth Anything 3. +""" + +from depth_anything_3.services.backend import create_app, start_server + +__all__ = [ + start_server, + create_app, +] diff --git a/depth_anything_3/services/backend.py b/depth_anything_3/services/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..3192488b3de3d97c458a7c20d8682dec46d6affe --- /dev/null +++ b/depth_anything_3/services/backend.py @@ -0,0 +1,1559 @@ +# flake8: noqa: E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model backend service for Depth Anything 3. +Provides HTTP API for model inference with persistent model loading. +""" + +import hmac +import os +import posixpath +import re +import secrets +import time +import uuid + +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional +from urllib.parse import quote +import numpy as np + +import uvicorn +from fastapi import Depends, FastAPI, Header, HTTPException +from fastapi.responses import FileResponse, HTMLResponse +from pydantic import BaseModel, field_validator + +from ..api import DepthAnything3 +from ..utils.export import SUPPORTED_EXPORT_FORMATS +from ..utils.memory import ( + get_gpu_memory_info, + cleanup_cuda_memory, + check_memory_availability, + estimate_memory_requirement, +) + + +class InferenceRequest(BaseModel): + """Request model for inference API.""" + + image_paths: List[str] + export_dir: Optional[str] = None + export_format: str = "mini_npz-glb" + extrinsics: Optional[List[List[List[float]]]] = None + intrinsics: Optional[List[List[List[float]]]] = None + process_res: int = 504 + process_res_method: str = "upper_bound_resize" + export_feat_layers: List[int] = [] + align_to_input_ext_scale: bool = True + # GLB export parameters + conf_thresh_percentile: float = 40.0 + num_max_points: int = 1_000_000 + show_cameras: bool = True + # Feat_vis export parameters + feat_vis_fps: int = 15 + + @field_validator("export_format") + @classmethod + def _validate_export_format(cls, value: str) -> str: + """Restrict export_format to the server's known/trusted export dispatcher formats, + instead of letting an arbitrary client-supplied string reach the export() dispatcher.""" + tokens = value.split("-") if value else [value] + unknown = sorted(set(tokens) - SUPPORTED_EXPORT_FORMATS) + if unknown: + raise ValueError( + f"Unsupported export_format value(s): {unknown}. " + f"Supported formats: {sorted(SUPPORTED_EXPORT_FORMATS)}" + ) + return value + + +class InferenceResponse(BaseModel): + """Response model for inference API.""" + + success: bool + message: str + task_id: Optional[str] = None + export_dir: Optional[str] = None + export_format: str = "mini_npz-glb" + processing_time: Optional[float] = None + + +class TaskStatus(BaseModel): + """Task status model.""" + + task_id: str + status: str # "pending", "running", "completed", "failed" + message: str + progress: Optional[float] = None # 0.0 to 1.0 + created_at: float + started_at: Optional[float] = None + completed_at: Optional[float] = None + export_dir: Optional[str] = None + request: Optional[InferenceRequest] = None # Store the original request + + # Essential task parameters + num_images: Optional[int] = None # Number of input images + export_format: Optional[str] = None # Export format + process_res_method: Optional[str] = None # Processing resolution method + video_path: Optional[str] = None # Source video path + + +class ModelBackend: + """Model backend service with persistent model loading.""" + + def __init__(self, model_dir: str, device: str = "cuda"): + self.model_dir = model_dir + self.device = device + self.model = None + self.model_loaded = False + self.load_time = None + self.load_start_time = None # Time when model loading started + self.load_completed_time = None # Time when model loading completed + self.last_used = None + + def load_model(self): + """Load model if not already loaded.""" + if self.model_loaded and self.model is not None: + self.last_used = time.time() + return self.model + + try: + print(f"Loading model from {self.model_dir}...") + self.load_start_time = time.time() + start_time = time.time() + + self.model = DepthAnything3.from_pretrained(self.model_dir).to(self.device) + self.model.eval() + + self.model_loaded = True + self.load_time = time.time() - start_time + self.load_completed_time = time.time() + self.last_used = time.time() + + print(f"Model loaded successfully in {self.load_time:.2f}s") + return self.model + + except Exception as e: + print(f"Failed to load model: {e}") + raise e + + def get_model(self): + """Get model, loading if necessary.""" + if not self.model_loaded: + return self.load_model() + self.last_used = time.time() + return self.model + + def get_status(self) -> Dict[str, Any]: + """Get backend status information.""" + # Calculate uptime from when model loading completed + uptime = 0 + if self.model_loaded and self.load_completed_time: + uptime = time.time() - self.load_completed_time + + return { + "model_loaded": self.model_loaded, + "model_dir": self.model_dir, + "device": self.device, + "load_time": self.load_time, + "last_used": self.last_used, + "uptime": uptime, + } + + +# Global backend instance +_backend: Optional[ModelBackend] = None +_app: Optional[FastAPI] = None +_gallery_dir: Optional[str] = None +_tasks: Dict[str, TaskStatus] = {} +_executor = ThreadPoolExecutor(max_workers=1) # Restrict to single-task execution +_running_task_id: Optional[str] = None # Currently running task ID +_task_queue: List[str] = [] # Pending task queue + +# Task cleanup configuration +MAX_TASK_HISTORY = 100 # Maximum number of tasks to keep in memory +CLEANUP_INTERVAL = 300 # Cleanup interval in seconds (5 minutes) + + +def _process_next_task(): + """Process the next task in the queue.""" + global _task_queue, _running_task_id + + if not _task_queue or _running_task_id is not None: + return + + # Get next task from queue + task_id = _task_queue.pop(0) + + # Get task request from tasks dict (we need to store the request) + if task_id not in _tasks: + return + + # Submit task to executor + _executor.submit(_run_inference_task, task_id) + + +# get_gpu_memory_info imported from depth_anything_3.utils.memory + + +# cleanup_cuda_memory imported from depth_anything_3.utils.memory + + +# check_memory_availability imported from depth_anything_3.utils.memory + + +# estimate_memory_requirement imported from depth_anything_3.utils.memory + + +def _run_inference_task(task_id: str): + """Run inference task in background thread with OOM protection.""" + global _tasks, _backend, _running_task_id, _task_queue + + model = None + inference_started = False + start_time = time.time() + + try: + # Get task request + if task_id not in _tasks or _tasks[task_id].request is None: + print(f"[{task_id}] Task not found or request missing") + return + + request = _tasks[task_id].request + num_images = len(request.image_paths) + + # Set current running task + _running_task_id = task_id + + # Update task status to running + _tasks[task_id].status = "running" + _tasks[task_id].started_at = start_time + _tasks[task_id].message = f"[{task_id}] Starting inference on {num_images} frames..." + print(f"[{task_id}] Starting inference on {num_images} frames") + + # Pre-inference cleanup to ensure maximum available memory + print(f"[{task_id}] Pre-inference cleanup...") + cleanup_cuda_memory() + + # Check memory availability + estimated_memory = estimate_memory_requirement(num_images, request.process_res) + mem_available, mem_msg = check_memory_availability(estimated_memory) + print(f"[{task_id}] {mem_msg}") + + if not mem_available: + # Try aggressive cleanup + print(f"[{task_id}] Insufficient memory, attempting aggressive cleanup...") + cleanup_cuda_memory() + time.sleep(0.5) # Give system time to reclaim memory + + # Check again + mem_available, mem_msg = check_memory_availability(estimated_memory) + if not mem_available: + raise RuntimeError( + f"Insufficient GPU memory after cleanup. {mem_msg}\n" + f"Suggestions:\n" + f" 1. Reduce process_res (current: {request.process_res})\n" + f" 2. Process fewer images at once (current: {num_images})\n" + f" 3. Clear other GPU processes" + ) + + # Get model (with error handling) + print(f"[{task_id}] Loading model...") + _tasks[task_id].message = f"[{task_id}] Loading model..." + _tasks[task_id].progress = 0.1 + + try: + model = _backend.get_model() + except RuntimeError as e: + if "out of memory" in str(e).lower(): + cleanup_cuda_memory() + raise RuntimeError( + f"OOM during model loading: {str(e)}\n" + f"Try reducing the batch size or resolution." + ) + raise + + print(f"[{task_id}] Model loaded successfully") + _tasks[task_id].progress = 0.2 + + # Prepare inference parameters + inference_kwargs = { + "image": request.image_paths, + "export_format": request.export_format, + "process_res": request.process_res, + "process_res_method": request.process_res_method, + "export_feat_layers": request.export_feat_layers, + "align_to_input_ext_scale": request.align_to_input_ext_scale, + "conf_thresh_percentile": request.conf_thresh_percentile, + "num_max_points": request.num_max_points, + "show_cameras": request.show_cameras, + "feat_vis_fps": request.feat_vis_fps, + } + + if request.export_dir: + # Re-check against the current configuration: this task may have sat in the + # queue for a while after the initial request-time check, and configuration + # can differ if the process was restarted with a different gallery_dir. + validate_export_dir(request.export_dir, _gallery_dir) + inference_kwargs["export_dir"] = request.export_dir + + if request.extrinsics: + inference_kwargs["extrinsics"] = np.array(request.extrinsics, dtype=np.float32) + + if request.intrinsics: + inference_kwargs["intrinsics"] = np.array(request.intrinsics, dtype=np.float32) + + # Run inference with timing + inference_start_time = time.time() + print(f"[{task_id}] Running model inference...") + _tasks[task_id].message = f"[{task_id}] Running model inference on {num_images} images..." + _tasks[task_id].progress = 0.3 + + inference_started = True + + try: + model.inference(**inference_kwargs) + inference_time = time.time() - inference_start_time + avg_time_per_image = inference_time / num_images if num_images > 0 else 0 + + print( + f"[{task_id}] Inference completed in {inference_time:.2f}s " + f"({avg_time_per_image:.2f}s per image)" + ) + + except RuntimeError as e: + if "out of memory" in str(e).lower(): + cleanup_cuda_memory() + raise RuntimeError( + f"OOM during inference: {str(e)}\n" + f"Settings: {num_images} images, resolution={request.process_res}\n" + f"Suggestions:\n" + f" 1. Reduce process_res to {int(request.process_res * 0.75)}\n" + f" 2. Process images in smaller batches\n" + f" 3. Use process_res_method='resize' instead of 'upper_bound_resize'" + ) + raise + + _tasks[task_id].progress = 0.9 + + # Post-inference cleanup + print(f"[{task_id}] Post-inference cleanup...") + cleanup_cuda_memory() + + # Calculate total processing time + total_time = time.time() - start_time + + # Update task status to completed + _tasks[task_id].status = "completed" + _tasks[task_id].completed_at = time.time() + _tasks[task_id].message = ( + f"[{task_id}] Completed in {total_time:.2f}s " f"({avg_time_per_image:.2f}s per image)" + ) + _tasks[task_id].progress = 1.0 + _tasks[task_id].export_dir = request.export_dir + + # Clear running state + _running_task_id = None + + # Process next task in queue + _process_next_task() + + print(f"[{task_id}] Task completed successfully") + print( + f"[{task_id}] Total time: {total_time:.2f}s, " + f"Inference time: {inference_time:.2f}s, " + f"Avg per image: {avg_time_per_image:.2f}s" + ) + + except Exception as e: + # Update task status to failed + error_msg = str(e) + total_time = time.time() - start_time + + print(f"[{task_id}] Task failed after {total_time:.2f}s: {error_msg}") + + # Always attempt cleanup on failure + cleanup_cuda_memory() + + _tasks[task_id].status = "failed" + _tasks[task_id].completed_at = time.time() + _tasks[task_id].message = f"[{task_id}] Failed after {total_time:.2f}s: {error_msg}" + + # Clear running state + _running_task_id = None + + # Process next task in queue + _process_next_task() + + finally: + # Final cleanup in finally block to ensure it always runs + # This is critical for releasing resources even if unexpected errors occur + try: + if inference_started: + print(f"[{task_id}] Final cleanup in finally block...") + cleanup_cuda_memory() + except Exception as e: + print(f"[{task_id}] Warning: Finally block cleanup failed: {e}") + + # Schedule cleanup after task completion + _schedule_task_cleanup() + + +def _cleanup_old_tasks(): + """Clean up old completed/failed tasks to prevent memory buildup.""" + global _tasks + + current_time = time.time() + tasks_to_remove = [] + + # Find tasks to remove - more aggressive cleanup + for task_id, task in _tasks.items(): + # Remove completed/failed tasks older than 10 minutes (instead of 1 hour) + if ( + task.status in ["completed", "failed"] + and task.completed_at + and current_time - task.completed_at > 600 + ): # 10 minutes + tasks_to_remove.append(task_id) + + # Remove old tasks + for task_id in tasks_to_remove: + del _tasks[task_id] + print(f"[CLEANUP] Removed old task: {task_id}") + + # If still too many tasks, remove oldest completed/failed tasks + if len(_tasks) > MAX_TASK_HISTORY: + completed_tasks = [ + (task_id, task) + for task_id, task in _tasks.items() + if task.status in ["completed", "failed"] + ] + completed_tasks.sort(key=lambda x: x[1].completed_at or 0) + + excess_count = len(_tasks) - MAX_TASK_HISTORY + for i in range(min(excess_count, len(completed_tasks))): + task_id = completed_tasks[i][0] + del _tasks[task_id] + print(f"[CLEANUP] Removed excess task: {task_id}") + + # Count active tasks (only pending and running) + active_count = sum(1 for task in _tasks.values() if task.status in ["pending", "running"]) + print( + "[CLEANUP] Task cleanup completed. " + f"Total tasks: {len(_tasks)}, Active tasks: {active_count}" + ) + + +def _schedule_task_cleanup(): + """Schedule task cleanup in background.""" + + def cleanup_worker(): + try: + time.sleep(2) # Small delay to ensure task status is updated + _cleanup_old_tasks() + except Exception as e: + print(f"[CLEANUP] Cleanup worker failed: {e}") + + # Run cleanup in background thread + _executor.submit(cleanup_worker) + + +# ============================================================================ +# Gallery utilities (extracted from gallery.py) +# ============================================================================ + +GALLERY_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + + +def _load_gallery_html() -> str: + """ + Load and modify gallery HTML to work under /gallery/ subdirectory. + Replaces API paths from root to /gallery/ prefix. + """ + from ..services.gallery import HTML_PAGE + + # Replace API paths to be under /gallery/ subdirectory + html = ( + HTML_PAGE.replace("fetch('/manifest.json'", "fetch('/gallery/manifest.json'") + .replace("fetch('/manifest/'+", "fetch('/gallery/manifest/'+") + .replace( + "if(location.pathname!=\"/\")history.replaceState(null,'','/'+location.search)", + "if(!location.pathname.startsWith(\"/gallery\"))history.replaceState(null,'','/gallery/'+location.search)", + ) + ) + + return html + + +def _gallery_url_join(*parts: str) -> str: + """Join URL parts safely.""" + norm = posixpath.join(*[p.replace("\\", "/") for p in parts]) + segs = [s for s in norm.split("/") if s not in ("", ".")] + return "/".join(quote(s) for s in segs) + + +def _is_plain_name(name: str) -> bool: + """Check if name is safe for use in paths.""" + return all(c not in name for c in ("/", "\\")) and name not in (".", "..") + + +def build_group_list(root_dir: str) -> dict: + """Build list of groups from gallery directory.""" + groups = [] + try: + for gname in sorted(os.listdir(root_dir)): + gpath = os.path.join(root_dir, gname) + if not os.path.isdir(gpath): + continue + has_scene = False + try: + for sname in os.listdir(gpath): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + if os.path.exists(os.path.join(spath, "scene.glb")) and os.path.exists( + os.path.join(spath, "scene.jpg") + ): + has_scene = True + break + except Exception: + pass + if has_scene: + groups.append({"id": gname, "title": gname}) + except Exception as e: + print(f"[warn] build_group_list failed: {e}") + return {"groups": groups} + + +def build_group_manifest(root_dir: str, group: str) -> dict: + """Build manifest for a specific group.""" + items = [] + gpath = os.path.join(root_dir, group) + try: + if not os.path.isdir(gpath): + return {"group": group, "items": []} + for sname in sorted(os.listdir(gpath)): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + glb_fs = os.path.join(spath, "scene.glb") + jpg_fs = os.path.join(spath, "scene.jpg") + if not (os.path.exists(glb_fs) and os.path.exists(jpg_fs)): + continue + depth_images = [] + dpath = os.path.join(spath, "depth_vis") + if os.path.isdir(dpath): + files = [ + f + for f in os.listdir(dpath) + if os.path.splitext(f)[1].lower() in GALLERY_IMAGE_EXTS + ] + for fn in sorted(files): + depth_images.append( + "/gallery/" + _gallery_url_join(group, sname, "depth_vis", fn) + ) + items.append( + { + "id": sname, + "title": sname, + "model": "/gallery/" + _gallery_url_join(group, sname, "scene.glb"), + "thumbnail": "/gallery/" + _gallery_url_join(group, sname, "scene.jpg"), + "depth_images": depth_images, + } + ) + except Exception as e: + print(f"[warn] build_group_manifest failed for {group}: {e}") + return {"group": group, "items": items} + + +# ============================================================================ +# Inference request validation +# ============================================================================ + +# export_dir is expected to be a simple relative path; anything else is rejected below. +_EXPORT_DIR_METACHAR_RE = re.compile(r"[;&|`$\n\r<>\x00]") + +# Hosts considered "local only" for the purposes of deciding whether an unauthenticated +# backend is safe to start. +_LOOPBACK_HOSTS = {"127.0.0.1", "localhost", "::1", "0:0:0:0:0:0:0:1"} + + +def validate_export_dir(export_dir: Optional[str], gallery_dir: Optional[str]) -> Optional[str]: + """Validate a client-supplied export_dir before it reaches any inference/export code. + + export_dir must be a simple relative path that resolves within the server's + configured gallery directory. + + Raises: + ValueError: with a human-readable reason, if export_dir is rejected. + """ + if not export_dir: + return export_dir + + if _EXPORT_DIR_METACHAR_RE.search(export_dir): + raise ValueError("export_dir contains disallowed characters") + + if os.path.isabs(export_dir) or export_dir.startswith("~"): + raise ValueError("export_dir must be a relative path") + + if any(part == ".." for part in export_dir.replace("\\", "/").split("/")): + raise ValueError("export_dir must not contain '..' path traversal") + + if not gallery_dir: + raise ValueError( + "export_dir is not permitted: no gallery directory is configured on this server" + ) + + gallery_root = os.path.realpath(gallery_dir) + resolved = os.path.realpath(export_dir) + if resolved != gallery_root and not resolved.startswith(gallery_root + os.sep): + raise ValueError( + "export_dir must resolve within the server's configured gallery directory" + ) + + return export_dir + + +def create_app( + model_dir: str, + device: str = "cuda", + gallery_dir: Optional[str] = None, + api_key: Optional[str] = None, +) -> FastAPI: + """Create FastAPI application with model backend. + + Note: the non-loopback-without-authentication guard lives in start_server(), not + here. Callers that take this app and serve it themselves (e.g. via a custom + uvicorn.run() or by mounting it into a larger ASGI app) are responsible for their + own host-binding and authentication decisions. + """ + global _backend, _app, _gallery_dir + + _backend = ModelBackend(model_dir, device) + _app = FastAPI( + title="Depth Anything 3 Backend", + description="Model inference service for Depth Anything 3", + version="1.0.0", + ) + + # Module-level (not just closure-local) so background tasks can re-check it too. + _gallery_dir = gallery_dir + + # Optional shared-secret gate for task-submitting endpoints. Falls back to the + # DA3_BACKEND_API_KEY env var so the key never needs to appear in a command line. + _api_key = api_key or os.environ.get("DA3_BACKEND_API_KEY") + + async def _require_api_key(x_api_key: Optional[str] = Header(default=None, alias="X-API-Key")): + """FastAPI dependency enforcing the API key when one is configured.""" + if _api_key and not hmac.compare_digest(x_api_key or "", _api_key): + raise HTTPException(status_code=401, detail="Missing or invalid X-API-Key header") + + @_app.get("/", response_class=HTMLResponse) + async def root(): + """Home page with navigation to dashboard and gallery.""" + html_content = ( + """ + + + + + + Depth Anything 3 Backend + + + +
+

Depth Anything 3

+

Model Backend Service

+
+ +

📊 Dashboard

+

Monitor backend status, model information, and inference tasks in real-time.

+ + """ + + ( + '' + "

🎨 Gallery

" + "

Browse 3D reconstructions and depth visualizations from processed scenes.

" + "" + if _gallery_dir and os.path.exists(_gallery_dir) + else "" + ) + + """ +
+
+

Depth Anything 3 Backend API

+
+
+ + + """ + ) + return HTMLResponse(html_content) + + @_app.get("/dashboard", response_class=HTMLResponse) + async def dashboard(): + """HTML dashboard for monitoring backend status and tasks.""" + if _backend is None: + return HTMLResponse("

Backend not initialized

", status_code=500) + + # Get backend status + status = _backend.get_status() + + # Safely format status values + if status["load_time"] is not None: + load_time_str = f"{status['load_time']:.2f}s" + else: + load_time_str = "Not loaded" + + if status["uptime"] is not None: + uptime_str = f"{status['uptime']:.2f}s" + else: + uptime_str = "Not running" + + # Get tasks information + active_tasks = [task for task in _tasks.values() if task.status in ["pending", "running"]] + completed_tasks = [ + task for task in _tasks.values() if task.status in ["completed", "failed"] + ] + + # Generate task HTML + active_tasks_html = "" + if active_tasks: + for task in active_tasks: + task_details = f""" +
+
+ {task.task_id} + {task.status} +
+
{task.message}
+
+ + Images: {task.num_images or 'N/A'} | + Format: {task.export_format or 'N/A'} | + Method: {task.process_res_method or 'N/A'} | + Export Dir: {task.export_dir or 'N/A'} + + {f'
Video: {task.video_path}' if task.video_path else ''} +
+
+ """ + active_tasks_html += task_details + else: + active_tasks_html = "

No active tasks

" + + completed_tasks_html = "" + if completed_tasks: + for task in completed_tasks[-10:]: + task_details = f""" +
+
+ {task.task_id} + {task.status} +
+
{task.message}
+
+ + Images: {task.num_images or 'N/A'} | + Format: {task.export_format or 'N/A'} | + Method: {task.process_res_method or 'N/A'} | + Export Dir: {task.export_dir or 'N/A'} + + {f'
Video: {task.video_path}' if task.video_path else ''} +
+
+ """ + completed_tasks_html += task_details + else: + completed_tasks_html = "

No completed tasks

" + + # Generate HTML + html_content = f""" + + + + + + Depth Anything 3 Backend Dashboard + + + +
+
+

Depth Anything 3 Backend Dashboard

+

Real-time monitoring of model status and inference tasks

+
+ +
+
+

Model Status

+
+ Status: + + {'Online' if status['model_loaded'] else 'Offline'} + +
+
+ Model Directory: + {status['model_dir']} +
+
+ Device: + {status['device']} +
+
+ Load Time: + {load_time_str} +
+
+ Uptime: + {uptime_str} +
+
+ +
+

Task Summary

+
+ Active Tasks: + {len(active_tasks)} +
+
+ Completed Tasks: + {len(completed_tasks)} +
+
+ Total Tasks: + {len(_tasks)} +
+
+
+ +
+

Active Tasks

+ + +
Last updated: {time.strftime('%Y-%m-%d %H:%M:%S')}
+ + {active_tasks_html} +
+ +
+

Recent Completed Tasks

+ {completed_tasks_html} +
+
+ + + + + """ + + return HTMLResponse(html_content) + + @_app.get("/status") + async def get_status(): + """Get backend status with GPU memory information.""" + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + status = _backend.get_status() + + # Add GPU memory information + gpu_memory = get_gpu_memory_info() + if gpu_memory: + status["gpu_memory"] = { + "total_gb": round(gpu_memory["total_gb"], 2), + "allocated_gb": round(gpu_memory["allocated_gb"], 2), + "reserved_gb": round(gpu_memory["reserved_gb"], 2), + "free_gb": round(gpu_memory["free_gb"], 2), + "utilization_percent": round(gpu_memory["utilization"], 1), + } + else: + status["gpu_memory"] = None + + return status + + @_app.post( + "/inference", response_model=InferenceResponse, dependencies=[Depends(_require_api_key)] + ) + async def run_inference(request: InferenceRequest): + """Submit inference task and return task ID.""" + global _running_task_id + + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + if request.export_dir: + try: + validate_export_dir(request.export_dir, _gallery_dir) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + # Generate unique task ID + task_id = str(uuid.uuid4()) + + # Create task status + if _running_task_id is not None: + status_msg = f"[{task_id}] Task queued (waiting for {_running_task_id} to complete)" + else: + status_msg = f"[{task_id}] Task submitted" + + _tasks[task_id] = TaskStatus( + task_id=task_id, + status="pending", + message=status_msg, + created_at=time.time(), + export_dir=request.export_dir, + request=request, + # Record essential parameters + num_images=len(request.image_paths), + export_format=request.export_format, + process_res_method=request.process_res_method, + video_path=( + request.image_paths[0] if request.image_paths else None + ), # Use first image path as video reference + ) + + # Add task to queue + _task_queue.append(task_id) + + # If no task is running, start processing the queue + if _running_task_id is None: + _process_next_task() + + return InferenceResponse( + success=True, + message="Task submitted successfully", + task_id=task_id, + export_dir=request.export_dir, + export_format=request.export_format, + ) + + @_app.get("/task/{task_id}", response_model=TaskStatus) + async def get_task_status(task_id: str): + """Get task status by task ID.""" + if task_id not in _tasks: + raise HTTPException(status_code=404, detail="Task not found") + + return _tasks[task_id] + + @_app.get("/gpu-memory") + async def get_gpu_memory(): + """Get detailed GPU memory information.""" + gpu_memory = get_gpu_memory_info() + if gpu_memory is None: + return { + "available": False, + "message": "CUDA not available or memory info cannot be retrieved", + } + + return { + "available": True, + "total_gb": round(gpu_memory["total_gb"], 2), + "allocated_gb": round(gpu_memory["allocated_gb"], 2), + "reserved_gb": round(gpu_memory["reserved_gb"], 2), + "free_gb": round(gpu_memory["free_gb"], 2), + "utilization_percent": round(gpu_memory["utilization"], 1), + "status": ( + "healthy" + if gpu_memory["utilization"] < 80 + else "warning" if gpu_memory["utilization"] < 95 else "critical" + ), + } + + @_app.get("/tasks") + async def list_tasks(): + """List all tasks.""" + # Separate active and completed tasks + active_tasks = [task for task in _tasks.values() if task.status in ["pending", "running"]] + completed_tasks = [ + task for task in _tasks.values() if task.status in ["completed", "failed"] + ] + + return { + "tasks": list(_tasks.values()), + "active_tasks": active_tasks, + "completed_tasks": completed_tasks, + "active_count": len(active_tasks), + "total_count": len(_tasks), + } + + @_app.post("/cleanup") + async def manual_cleanup(): + """Manually trigger task cleanup.""" + try: + _cleanup_old_tasks() + return {"message": "Cleanup completed", "active_tasks": len(_tasks)} + except Exception as e: + raise HTTPException(status_code=500, detail=f"Cleanup failed: {str(e)}") + + @_app.delete("/task/{task_id}") + async def delete_task(task_id: str): + """Delete a specific task.""" + if task_id not in _tasks: + raise HTTPException(status_code=404, detail="Task not found") + + # Only allow deletion of completed/failed tasks + if _tasks[task_id].status not in ["completed", "failed"]: + raise HTTPException(status_code=400, detail="Cannot delete running or pending tasks") + + del _tasks[task_id] + return {"message": f"Task {task_id} deleted successfully"} + + @_app.post("/reload") + async def reload_model(): + """Reload the model.""" + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + try: + _backend.model = None + _backend.model_loaded = False + _backend.load_model() + return {"message": "Model reloaded successfully"} + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to reload model: {str(e)}") + + # ============================================================================ + # Gallery routes + # ============================================================================ + + if _gallery_dir and os.path.exists(_gallery_dir): + # Load gallery HTML page (with modified paths for /gallery/ subdirectory) + _gallery_html = _load_gallery_html() + + @_app.get("/gallery/", response_class=HTMLResponse) + @_app.get("/gallery", response_class=HTMLResponse) + async def gallery_home(): + """Gallery home page.""" + return HTMLResponse(_gallery_html) + + @_app.get("/gallery/manifest.json") + async def gallery_manifest(): + """Get gallery group list.""" + try: + return build_group_list(_gallery_dir) + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to build group list: {str(e)}" + ) + + @_app.get("/gallery/manifest/{group}.json") + async def gallery_group_manifest(group: str): + """Get manifest for a specific group.""" + if not _is_plain_name(group): + raise HTTPException(status_code=400, detail="Invalid group name") + try: + return build_group_manifest(_gallery_dir, group) + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to build group manifest: {str(e)}" + ) + + @_app.get("/gallery/{path:path}") + async def gallery_files(path: str): + """Serve gallery static files (GLB, JPG, etc.).""" + # Security check: prevent directory traversal + path_parts = path.split("/") + if any(not _is_plain_name(part) for part in path_parts if part): + raise HTTPException(status_code=400, detail="Invalid path") + + file_path = os.path.join(_gallery_dir, *path_parts) + + # Ensure the file is within gallery directory + real_file_path = os.path.realpath(file_path) + real_gallery_dir = os.path.realpath(_gallery_dir) + if not real_file_path.startswith(real_gallery_dir): + raise HTTPException(status_code=403, detail="Access denied") + + if not os.path.exists(file_path) or not os.path.isfile(file_path): + raise HTTPException(status_code=404, detail="File not found") + + return FileResponse(file_path) + + return _app + + +def start_server( + model_dir: str, + device: str = "cuda", + host: str = "127.0.0.1", + port: int = 8000, + gallery_dir: Optional[str] = None, + api_key: Optional[str] = None, + allow_unauthenticated: bool = False, +): + """Start the backend server.""" + api_key = api_key or os.environ.get("DA3_BACKEND_API_KEY") + auto_generated_key = False + + if host not in _LOOPBACK_HOSTS and not api_key and not allow_unauthenticated: + # No key was configured for a host reachable beyond this machine: generate one + # for this run rather than requiring the operator to come up with one first. + api_key = secrets.token_urlsafe(32) + auto_generated_key = True + + app = create_app(model_dir, device, gallery_dir, api_key=api_key) + + print("Starting Depth Anything 3 Backend...") + print(f"Model directory: {model_dir}") + print(f"Device: {device}") + print(f"Server: http://{host}:{port}") + print(f"Dashboard: http://{host}:{port}/dashboard") + print(f"API Status: http://{host}:{port}/status") + + if gallery_dir and os.path.exists(gallery_dir): + print(f"Gallery: http://{host}:{port}/gallery/") + + if auto_generated_key: + print("=" * 60) + print("No --api-key was set, so one was generated for this run:") + print(f" X-API-Key: {api_key}") + print("Set DA3_BACKEND_API_KEY to this value wherever you submit jobs from") + print("(e.g. `da3 ... --use-backend`), or pass --api-key / set") + print("DA3_BACKEND_API_KEY before starting the backend for a key that stays") + print("the same across restarts. Use --allow-unauthenticated to skip this.") + elif api_key: + print("Authentication: X-API-Key header required for /inference") + + print("=" * 60) + print("Backend is running! You can now:") + print(f" • Open home page: http://{host}:{port}") + print(f" • Open dashboard: http://{host}:{port}/dashboard") + print(f" • Check API status: http://{host}:{port}/status") + + if gallery_dir and os.path.exists(gallery_dir): + print(f" • Browse gallery: http://{host}:{port}/gallery/") + + print(" • Submit inference tasks via API") + print("=" * 60) + + uvicorn.run(app, host=host, port=port, log_level="info") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Depth Anything 3 Backend Server") + parser.add_argument("--model-dir", required=True, help="Model directory path") + parser.add_argument("--device", default="cuda", help="Device to use") + parser.add_argument("--host", default="127.0.0.1", help="Host to bind to") + parser.add_argument("--port", type=int, default=8000, help="Port to bind to") + parser.add_argument("--gallery-dir", help="Gallery directory path (optional)") + parser.add_argument( + "--api-key", + default=None, + help="Require this API key (X-API-Key header) on inference requests. Falls back " + "to the DA3_BACKEND_API_KEY env var, or an auto-generated key when binding to a " + "non-loopback host.", + ) + parser.add_argument( + "--allow-unauthenticated", + action="store_true", + help="Skip authentication entirely on a non-loopback host, instead of using an " + "auto-generated API key.", + ) + + args = parser.parse_args() + start_server( + args.model_dir, + args.device, + args.host, + args.port, + args.gallery_dir, + api_key=args.api_key, + allow_unauthenticated=args.allow_unauthenticated, + ) diff --git a/depth_anything_3/services/gallery.py b/depth_anything_3/services/gallery.py new file mode 100644 index 0000000000000000000000000000000000000000..f72bb5e5f6defbc24cf2278a53b7162a8ad5519d --- /dev/null +++ b/depth_anything_3/services/gallery.py @@ -0,0 +1,806 @@ +#!/usr/bin/env python3 +# flake8: noqa: E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Depth Anything 3 Gallery Server (two-level, single-file) +Now supports paginated depth preview (4 per page). +""" + +import argparse +import json +import mimetypes +import os +import posixpath +import sys +from functools import partial +from http import HTTPStatus +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import quote, unquote + +# ------------------------------ Embedded HTML ------------------------------ # + +HTML_PAGE = r""" + + + + Depth Anything 3 Gallery + + + + + + + +
+
+ +

Depth Anything 3 Gallery

+ + +
+
Level 1 shows groups only; click a group to browse scenes and previews.
+
+ +
+ +
+

+ 🎯 Depth Anything 3 Gallery +

+

+ Explore 3D reconstructions and depth visualizations from Depth Anything 3. + Browse through groups of scenes, preview 3D models, and examine depth maps interactively. +

+
+ +
+
    + +
    + + +
    + + + +
    Depth Anything 3 Gallery. Copyright 2025 Depth Anything 3 authors.
    + + + + +""" + +# ------------------------------ Utilities ------------------------------ # + +IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + + +def _url_join(*parts: str) -> str: + norm = posixpath.join(*[p.replace("\\", "/") for p in parts]) + segs = [s for s in norm.split("/") if s not in ("", ".")] + return "/".join(quote(s) for s in segs) + + +def _is_plain_name(name: str) -> bool: + return all(c not in name for c in ("/", "\\")) and name not in (".", "..") + + +def build_group_list(root_dir: str) -> dict: + groups = [] + try: + for gname in sorted(os.listdir(root_dir)): + gpath = os.path.join(root_dir, gname) + if not os.path.isdir(gpath): + continue + has_scene = False + try: + for sname in os.listdir(gpath): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + if os.path.exists(os.path.join(spath, "scene.glb")) and os.path.exists( + os.path.join(spath, "scene.jpg") + ): + has_scene = True + break + except Exception: + pass + if has_scene: + groups.append({"id": gname, "title": gname}) + except Exception as e: + print(f"[warn] build_group_list failed: {e}", file=sys.stderr) + return {"groups": groups} + + +def build_group_manifest(root_dir: str, group: str) -> dict: + items = [] + gpath = os.path.join(root_dir, group) + try: + if not os.path.isdir(gpath): + return {"group": group, "items": []} + for sname in sorted(os.listdir(gpath)): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + glb_fs = os.path.join(spath, "scene.glb") + jpg_fs = os.path.join(spath, "scene.jpg") + if not (os.path.exists(glb_fs) and os.path.exists(jpg_fs)): + continue + depth_images = [] + dpath = os.path.join(spath, "depth_vis") + if os.path.isdir(dpath): + files = [ + f for f in os.listdir(dpath) if os.path.splitext(f)[1].lower() in IMAGE_EXTS + ] + for fn in sorted(files): + depth_images.append("/" + _url_join(group, sname, "depth_vis", fn)) + items.append( + { + "id": sname, + "title": sname, + "model": "/" + _url_join(group, sname, "scene.glb"), + "thumbnail": "/" + _url_join(group, sname, "scene.jpg"), + "depth_images": depth_images, + } + ) + except Exception as e: + print(f"[warn] build_group_manifest failed for {group}: {e}", file=sys.stderr) + return {"group": group, "items": items} + + +class GalleryHandler(SimpleHTTPRequestHandler): + def __init__(self, *args, directory=None, **kwargs): + super().__init__(*args, directory=directory, **kwargs) + + def do_GET(self): + if self.path in ("/", "/index.html") or self.path.startswith("/?"): + content = HTML_PAGE.encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(content))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(content) + return + if self.path == "/manifest.json": + data = json.dumps( + build_group_list(self.directory), ensure_ascii=False, indent=2 + ).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(data))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(data) + return + if self.path.startswith("/manifest/") and self.path.endswith(".json"): + group_enc = self.path[len("/manifest/") : -len(".json")] + try: + group = unquote(group_enc) + except Exception: + group = group_enc + if not _is_plain_name(group): + self.send_error(HTTPStatus.BAD_REQUEST, "Invalid group name") + return + data = json.dumps( + build_group_manifest(self.directory, group), ensure_ascii=False, indent=2 + ).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(data))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(data) + return + if self.path == "/favicon.ico": + self.send_response(HTTPStatus.NO_CONTENT) + self.end_headers() + return + return super().do_GET() + + def list_directory(self, path): + self.send_error(HTTPStatus.NOT_FOUND, "Directory listing disabled") + return None + + +def gallery(): + parser = argparse.ArgumentParser( + description="Depth Anything 3 Gallery Server (two-level, with pagination)" + ) + parser.add_argument( + "-d", "--dir", required=True, help="Gallery root directory (two-level: group/scene)" + ) + parser.add_argument("-p", "--port", type=int, default=8000, help="Port (default 8000)") + parser.add_argument("--host", default="127.0.0.1", help="Host address (default 127.0.0.1)") + parser.add_argument("--open", action="store_true", help="Open browser after launch") + args = parser.parse_args() + + root_dir = os.path.abspath(args.dir) + if not os.path.isdir(root_dir): + print(f"[error] Directory not found: {root_dir}", file=sys.stderr) + sys.exit(1) + + Handler = partial(GalleryHandler, directory=root_dir) + server = ThreadingHTTPServer((args.host, args.port), Handler) + + addr = f"http://{args.host}:{args.port}/" + print(f"[info] Serving gallery from: {root_dir}") + print(f"[info] Open: {addr}") + + if args.open: + try: + import webbrowser + + webbrowser.open(addr) + except Exception as e: + print(f"[warn] Failed to open browser: {e}", file=sys.stderr) + + try: + server.serve_forever() + except KeyboardInterrupt: + print("\n[info] Shutting down...") + finally: + server.server_close() + + +def main(): + """Main entry point for gallery server.""" + mimetypes.add_type("model/gltf-binary", ".glb") + gallery() + + +if __name__ == "__main__": + main() diff --git a/depth_anything_3/services/inference_service.py b/depth_anything_3/services/inference_service.py new file mode 100644 index 0000000000000000000000000000000000000000..3f3ad512360179215739e26532b9269bf9580d11 --- /dev/null +++ b/depth_anything_3/services/inference_service.py @@ -0,0 +1,247 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unified Inference Service +Provides unified interface for local and remote inference +""" + +import os + +from typing import Any, Dict, List, Optional, Union +import numpy as np +import requests +import typer + +from ..api import DepthAnything3 + + +class InferenceService: + """Unified inference service class""" + + def __init__(self, model_dir: str, device: str = "cuda"): + self.model_dir = model_dir + self.device = device + self.model = None + + def load_model(self): + """Load model""" + if self.model is None: + typer.echo(f"Loading model from {self.model_dir}...") + self.model = DepthAnything3.from_pretrained(self.model_dir).to(self.device) + return self.model + + def run_local_inference( + self, + image_paths: List[str], + export_dir: str, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, + ) -> Any: + """Run local inference""" + if export_feat_layers is None: + export_feat_layers = [] + + model = self.load_model() + + # Prepare inference parameters + inference_kwargs = { + "image": image_paths, + "export_dir": export_dir, + "export_format": export_format, + "process_res": process_res, + "process_res_method": process_res_method, + "export_feat_layers": export_feat_layers, + "align_to_input_ext_scale": align_to_input_ext_scale, + "use_ray_pose": use_ray_pose, + "ref_view_strategy": ref_view_strategy, + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + "feat_vis_fps": feat_vis_fps, + } + + # Add pose data (if exists) + if extrinsics is not None: + inference_kwargs["extrinsics"] = extrinsics + if intrinsics is not None: + inference_kwargs["intrinsics"] = intrinsics + + # Run inference + typer.echo(f"Running inference on {len(image_paths)} images...") + prediction = model.inference(**inference_kwargs) + + typer.echo(f"Results saved to {export_dir}") + typer.echo(f"Export format: {export_format}") + + return prediction + + def run_backend_inference( + self, + image_paths: List[str], + export_dir: str, + backend_url: str, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, + ) -> Dict[str, Any]: + """Run backend inference""" + if export_feat_layers is None: + export_feat_layers = [] + + # Check backend status + if not self._check_backend_status(backend_url): + raise typer.BadParameter(f"Backend service is not running at {backend_url}") + + # Prepare payload + payload = { + "image_paths": image_paths, + "export_dir": export_dir, + "export_format": export_format, + "process_res": process_res, + "process_res_method": process_res_method, + "export_feat_layers": export_feat_layers, + "align_to_input_ext_scale": align_to_input_ext_scale, + "use_ray_pose": use_ray_pose, + "ref_view_strategy": ref_view_strategy, + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + "feat_vis_fps": feat_vis_fps, + } + + # Add pose data (if exists) + if extrinsics is not None: + payload["extrinsics"] = [ext.astype(np.float64).tolist() for ext in extrinsics] + if intrinsics is not None: + payload["intrinsics"] = [intr.astype(np.float64).tolist() for intr in intrinsics] + + # Submit task + typer.echo("Submitting inference task to backend...") + # Sent automatically when the backend was started with a matching --api-key / + # DA3_BACKEND_API_KEY, so client and server share the same env-var convention. + api_key = os.environ.get("DA3_BACKEND_API_KEY") + headers = {"X-API-Key": api_key} if api_key else None + try: + response = requests.post( + f"{backend_url}/inference", json=payload, headers=headers, timeout=30 + ) + response.raise_for_status() + result = response.json() + + if result["success"]: + task_id = result["task_id"] + typer.echo("Task submitted successfully!") + typer.echo(f"Task ID: {task_id}") + typer.echo(f"Results will be saved to: {export_dir}") + typer.echo(f"Check backend logs for progress updates with task ID: {task_id}") + return result + else: + raise typer.BadParameter( + f"Backend inference submission failed: {result['message']}" + ) + except requests.exceptions.RequestException as e: + raise typer.BadParameter(f"Backend inference submission failed: {e}") + + def _check_backend_status(self, backend_url: str) -> bool: + """Check backend status""" + try: + response = requests.get(f"{backend_url}/status", timeout=5) + return response.status_code == 200 + except Exception: + return False + + +def run_inference( + image_paths: List[str], + export_dir: str, + model_dir: str, + device: str = "cuda", + backend_url: Optional[str] = None, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, +) -> Union[Any, Dict[str, Any]]: + """Unified inference interface""" + + service = InferenceService(model_dir, device) + + if backend_url: + return service.run_backend_inference( + image_paths=image_paths, + export_dir=export_dir, + backend_url=backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + else: + return service.run_local_inference( + image_paths=image_paths, + export_dir=export_dir, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) diff --git a/depth_anything_3/services/input_handlers.py b/depth_anything_3/services/input_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..dc0536b51e5c71bb6fb9d10f8e74fbb12fc42d31 --- /dev/null +++ b/depth_anything_3/services/input_handlers.py @@ -0,0 +1,266 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Input Processing Service +Handles different types of inputs (image, images, colmap, video) +""" + +import glob +import os +from typing import List, Tuple +import cv2 +import numpy as np +import typer + +from ..utils.read_write_model import read_model + + +class InputHandler: + """Base input handler class""" + + @staticmethod + def validate_path(path: str, path_type: str = "file") -> str: + """Validate path""" + if not os.path.exists(path): + raise typer.BadParameter(f"{path_type} not found: {path}") + return path + + @staticmethod + def handle_export_dir(export_dir: str, auto_cleanup: bool = False) -> str: + """Handle export directory""" + if os.path.exists(export_dir): + if auto_cleanup: + typer.echo(f"Auto-cleaning existing export directory: {export_dir}") + import shutil + + shutil.rmtree(export_dir) + os.makedirs(export_dir, exist_ok=True) + else: + typer.echo(f"Export directory '{export_dir}' already exists.") + if typer.confirm("Do you want to clean it and continue?"): + import shutil + + shutil.rmtree(export_dir) + os.makedirs(export_dir, exist_ok=True) + typer.echo(f"Cleaned export directory: {export_dir}") + else: + typer.echo("Operation cancelled.") + raise typer.Exit(0) + else: + os.makedirs(export_dir, exist_ok=True) + return export_dir + + +class ImageHandler(InputHandler): + """Single image handler""" + + @staticmethod + def process(image_path: str) -> List[str]: + """Process single image""" + InputHandler.validate_path(image_path, "Image file") + return [image_path] + + +class ImagesHandler(InputHandler): + """Image directory handler""" + + @staticmethod + def process(images_dir: str, image_extensions: str = "png,jpg,jpeg") -> List[str]: + """Process image directory""" + InputHandler.validate_path(images_dir, "Images directory") + + # Parse extensions + extensions = [ext.strip().lower() for ext in image_extensions.split(",")] + extensions = [ext if ext.startswith(".") else f".{ext}" for ext in extensions] + + # Find image files + image_files = [] + for ext in extensions: + pattern = f"*{ext}" + image_files.extend(glob.glob(os.path.join(images_dir, pattern))) + image_files.extend(glob.glob(os.path.join(images_dir, pattern.upper()))) + + image_files = sorted(list(set(image_files))) # Remove duplicates and sort + + if not image_files: + raise typer.BadParameter( + f"No image files found in {images_dir} with extensions: {extensions}" + ) + + typer.echo(f"Found {len(image_files)} images to process") + return image_files + + +class ColmapHandler(InputHandler): + """COLMAP data handler""" + + @staticmethod + def process( + colmap_dir: str, sparse_subdir: str = "" + ) -> Tuple[List[str], np.ndarray, np.ndarray]: + """Process COLMAP data""" + InputHandler.validate_path(colmap_dir, "COLMAP directory") + + # Build paths + images_dir = os.path.join(colmap_dir, "images") + if sparse_subdir: + sparse_dir = os.path.join(colmap_dir, "sparse", sparse_subdir) + else: + sparse_dir = os.path.join(colmap_dir, "sparse") + + InputHandler.validate_path(images_dir, "Images directory") + InputHandler.validate_path(sparse_dir, "Sparse reconstruction directory") + + # Load COLMAP data + typer.echo("Loading COLMAP reconstruction data...") + try: + cameras, images, points3D = read_model(sparse_dir) + + typer.echo( + f"Loaded COLMAP data: {len(cameras)} cameras, {len(images)} images, " + f"{len(points3D)} 3D points." + ) + + # Get image files and pose data + image_files = [] + extrinsics = [] + intrinsics = [] + + for image_id, image_data in images.items(): + image_name = image_data.name + image_path = os.path.join(images_dir, image_name) + + if os.path.exists(image_path): + image_files.append(image_path) + + # Get camera parameters + camera = cameras[image_data.camera_id] + + # Convert quaternion to rotation matrix + R = image_data.qvec2rotmat() + t = image_data.tvec + + # Create extrinsic matrix (world to camera) + extrinsic = np.eye(4) + extrinsic[:3, :3] = R + extrinsic[:3, 3] = t + extrinsics.append(extrinsic) + + # Create intrinsic matrix + if camera.model == "PINHOLE": + fx, fy, cx, cy = camera.params + elif camera.model == "SIMPLE_PINHOLE": + f, cx, cy = camera.params + fx = fy = f + else: + # For other models, use basic pinhole approximation + fx = fy = camera.params[0] if len(camera.params) > 0 else 1000 + cx = camera.width / 2 + cy = camera.height / 2 + + intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + intrinsics.append(intrinsic) + + if not image_files: + raise typer.BadParameter("No valid images found in COLMAP data") + + typer.echo(f"Found {len(image_files)} valid images with pose data") + + return image_files, np.array(extrinsics), np.array(intrinsics) + + except Exception as e: + raise typer.BadParameter(f"Failed to load COLMAP data: {e}") + + +class VideoHandler(InputHandler): + """Video handler""" + + @staticmethod + def process(video_path: str, output_dir: str, fps: float = 1.0) -> List[str]: + """Process video, extract frames""" + InputHandler.validate_path(video_path, "Video file") + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise typer.BadParameter(f"Cannot open video: {video_path}") + + # Get video properties + video_fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + duration = total_frames / video_fps + + # Calculate frame interval (ensure at least 1) + frame_interval = max(1, int(video_fps / fps)) + actual_fps = video_fps / frame_interval + + typer.echo(f"Video FPS: {video_fps:.2f}, Duration: {duration:.2f}s") + + # Warn if requested FPS is higher than video FPS + if fps > video_fps: + typer.echo( + f"⚠️ Warning: Requested sampling FPS ({fps:.2f}) exceeds video FPS ({video_fps:.2f})", # noqa: E501 + err=True, + ) + typer.echo( + f"⚠️ Using maximum available FPS: {actual_fps:.2f} (extracting every frame)", + err=True, + ) + + typer.echo(f"Extracting frames at {actual_fps:.2f} FPS (every {frame_interval} frame(s))") + + # Create output directory + frames_dir = os.path.join(output_dir, "input_images") + os.makedirs(frames_dir, exist_ok=True) + + frame_count = 0 + saved_count = 0 + + while True: + ret, frame = cap.read() + if not ret: + break + + if frame_count % frame_interval == 0: + frame_path = os.path.join(frames_dir, f"{saved_count:06d}.png") + cv2.imwrite(frame_path, frame) + saved_count += 1 + + frame_count += 1 + + cap.release() + typer.echo(f"Extracted {saved_count} frames to {frames_dir}") + + # Get frame file list + frame_files = sorted( + [f for f in os.listdir(frames_dir) if f.endswith((".png", ".jpg", ".jpeg"))] + ) + if not frame_files: + raise typer.BadParameter("No frames extracted from video") + + return [os.path.join(frames_dir, f) for f in frame_files] + + +def parse_export_feat(export_feat_str: str) -> List[int]: + """Parse export_feat parameter""" + if not export_feat_str: + return [] + + try: + return [int(x.strip()) for x in export_feat_str.split(",") if x.strip()] + except ValueError: + raise typer.BadParameter( + f"Invalid export_feat format: {export_feat_str}. " + "Use comma-separated integers like '0,1,2'" + ) diff --git a/depth_anything_3/specs.py b/depth_anything_3/specs.py new file mode 100644 index 0000000000000000000000000000000000000000..fe5b30255e9fd48d988ff00c896fb3dbadf197ea --- /dev/null +++ b/depth_anything_3/specs.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional +import numpy as np +import torch + + +@dataclass +class Gaussians: + """3DGS parameters, all in world space""" + + means: torch.Tensor # world points, "batch gaussian dim" + scales: torch.Tensor # scales_std, "batch gaussian 3" + rotations: torch.Tensor # world_quat_wxyz, "batch gaussian 4" + harmonics: torch.Tensor # world SH, "batch gaussian 3 d_sh" + opacities: torch.Tensor # opacity | opacity SH, "batch gaussian" | "batch gaussian 1 d_sh" + + +@dataclass +class Prediction: + depth: np.ndarray # N, H, W + is_metric: int + sky: np.ndarray | None = None # N, H, W + conf: np.ndarray | None = None # N, H, W + extrinsics: np.ndarray | None = None # N, 4, 4 + intrinsics: np.ndarray | None = None # N, 3, 3 + processed_images: np.ndarray | None = None # N, H, W, 3 - processed images for visualization + gaussians: Gaussians | None = None # 3D gaussians + aux: dict[str, Any] = None # + scale_factor: Optional[float] = None # metric scale diff --git a/depth_anything_3/utils/alignment.py b/depth_anything_3/utils/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..ceb8983fbb85fd8cad51fdd82b794aa3a114df78 --- /dev/null +++ b/depth_anything_3/utils/alignment.py @@ -0,0 +1,163 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Alignment utilities for depth estimation and metric scaling. +""" + +from typing import Tuple +import torch + + +def least_squares_scale_scalar( + a: torch.Tensor, b: torch.Tensor, eps: float = 1e-12 +) -> torch.Tensor: + """ + Compute least squares scale factor s such that a ≈ s * b. + + Args: + a: First tensor + b: Second tensor + eps: Small epsilon for numerical stability + + Returns: + Scalar tensor containing the scale factor + + Raises: + ValueError: If tensors have mismatched shapes or devices + TypeError: If tensors are not floating point + """ + if a.shape != b.shape: + raise ValueError(f"Shape mismatch: {a.shape} vs {b.shape}") + if a.device != b.device: + raise ValueError(f"Device mismatch: {a.device} vs {b.device}") + if not a.is_floating_point() or not b.is_floating_point(): + raise TypeError("Tensors must be floating point type") + + # Compute dot products for least squares solution + num = torch.dot(a.reshape(-1), b.reshape(-1)) + den = torch.dot(b.reshape(-1), b.reshape(-1)).clamp_min(eps) + return num / den + + +def compute_sky_mask(sky_prediction: torch.Tensor, threshold: float = 0.3) -> torch.Tensor: + """ + Compute non-sky mask from sky prediction. + + Args: + sky_prediction: Sky prediction tensor + threshold: Threshold for sky classification + + Returns: + Boolean mask where True indicates non-sky regions + """ + return sky_prediction < threshold + + +def compute_alignment_mask( + depth_conf: torch.Tensor, + non_sky_mask: torch.Tensor, + depth: torch.Tensor, + metric_depth: torch.Tensor, + median_conf: torch.Tensor, + min_depth_threshold: float = 1e-3, + min_metric_depth_threshold: float = 1e-2, +) -> torch.Tensor: + """ + Compute mask for depth alignment based on confidence and depth thresholds. + + Args: + depth_conf: Depth confidence tensor + non_sky_mask: Non-sky region mask + depth: Predicted depth tensor + metric_depth: Metric depth tensor + median_conf: Median confidence threshold + min_depth_threshold: Minimum depth threshold + min_metric_depth_threshold: Minimum metric depth threshold + + Returns: + Boolean mask for valid alignment regions + """ + return ( + (depth_conf >= median_conf) + & non_sky_mask + & (metric_depth > min_metric_depth_threshold) + & (depth > min_depth_threshold) + ) + + +def sample_tensor_for_quantile(tensor: torch.Tensor, max_samples: int = 100000) -> torch.Tensor: + """ + Sample tensor elements for quantile computation to reduce memory usage. + + Args: + tensor: Input tensor to sample + max_samples: Maximum number of samples to take + + Returns: + Sampled tensor + """ + if tensor.numel() <= max_samples: + return tensor + + idx = torch.randperm(tensor.numel(), device=tensor.device)[:max_samples] + return tensor.flatten()[idx] + + +def apply_metric_scaling( + depth: torch.Tensor, intrinsics: torch.Tensor, scale_factor: float = 300.0 +) -> torch.Tensor: + """ + Apply metric scaling to depth based on camera intrinsics. + + Args: + depth: Input depth tensor + intrinsics: Camera intrinsics tensor + scale_factor: Scaling factor for metric conversion + + Returns: + Scaled depth tensor + """ + focal_length = (intrinsics[:, :, 0, 0] + intrinsics[:, :, 1, 1]) / 2 + return depth * (focal_length[:, :, None, None] / scale_factor) + + +def set_sky_regions_to_max_depth( + depth: torch.Tensor, + depth_conf: torch.Tensor, + non_sky_mask: torch.Tensor, + max_depth: float = 200.0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Set sky regions to maximum depth and high confidence. + + Args: + depth: Depth tensor + depth_conf: Depth confidence tensor + non_sky_mask: Non-sky region mask + max_depth: Maximum depth value for sky regions + + Returns: + Tuple of (updated_depth, updated_depth_conf) + """ + depth = depth.clone() + + # Set sky regions to max depth and high confidence + depth[~non_sky_mask] = max_depth + if depth_conf is not None: + depth_conf = depth_conf.clone() + depth_conf[~non_sky_mask] = 1.0 + return depth, depth_conf + else: + return depth, None diff --git a/depth_anything_3/utils/api_helpers.py b/depth_anything_3/utils/api_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..b327331d9ec61a3047be3e330f41156eb124adc4 --- /dev/null +++ b/depth_anything_3/utils/api_helpers.py @@ -0,0 +1,58 @@ +import argparse + + +def parse_scalar(s): + if not isinstance(s, str): + return s + t = s.strip() + l = t.lower() + if l == "true": + return True + if l == "false": + return False + if l in ("none", "null"): + return None + try: + return int(t, 10) + except Exception: + pass + try: + return float(t) + except Exception: + return s + + +def fn_kv_csv(s: str) -> dict[str, dict[str, object]]: + """ + Parse a string of comma-separated triplets: fn:key:value + + Returns: + dict[fn_name] -> dict[key] = parsed_value + + Example: + "fn1:width:1920,fn1:height:1080,fn2:quality:0.8" + -> {"fn1": {"width": 1920, "height": 1080}, "fn2": {"quality": 0.8}} + """ + result: dict[str, dict[str, object]] = {} + if not s: + return result + + for item in s.split(","): + if not item: + continue + parts = item.split(":", 2) # allow value to contain ":" beyond first two separators + if len(parts) < 3: + raise argparse.ArgumentTypeError(f"Bad item '{item}', expected FN:KEY:VALUE") + fn, key, raw_val = parts[0], parts[1], parts[2] + # If you need to allow colons in values, join leftover parts: + # fn, key, raw_val = parts[0], parts[1], ":".join(parts[2:]) + + if not fn: + raise argparse.ArgumentTypeError(f"Bad item '{item}': empty function name") + if not key: + raise argparse.ArgumentTypeError(f"Bad item '{item}': empty key") + + val = parse_scalar(raw_val) + bucket = result.setdefault(fn, {}) + bucket[key] = val + return result diff --git a/depth_anything_3/utils/camera_trj_helpers.py b/depth_anything_3/utils/camera_trj_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..83624f359553908abd28af2d59f2066a7f7a7b15 --- /dev/null +++ b/depth_anything_3/utils/camera_trj_helpers.py @@ -0,0 +1,479 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from einops import einsum, rearrange, reduce + +try: + from scipy.spatial.transform import Rotation as R +except ImportError: + from depth_anything_3.utils.logger import logger + + logger.warn("Dependency 'scipy' not found. Required for interpolating camera trajectory.") + +from depth_anything_3.utils.geometry import as_homogeneous + + +@torch.no_grad() +def render_stabilization_path(poses, k_size=45): + """Rendering stabilized camera path. + poses: [batch, 4, 4] or [batch, 3, 4], + return: + smooth path: [batch 4 4]""" + num_frames = poses.shape[0] + device = poses.device + dtype = poses.dtype + + # Early exit for trivial cases + if num_frames <= 1: + return as_homogeneous(poses) + + # Make k_size safe: positive odd and not larger than num_frames + # 1) Ensure odd + if k_size < 1: + k_size = 1 + if k_size % 2 == 0: + k_size += 1 + # 2) Cap to num_frames (keep odd) + max_odd = num_frames if (num_frames % 2 == 1) else (num_frames - 1) + if max_odd < 1: + max_odd = 1 # covers num_frames == 0 theoretically + k_size = min(k_size, max_odd) + # 3) enforce a minimum of 3 when possible (for better smoothing) + if num_frames >= 3 and k_size < 3: + k_size = 3 + + input_poses = [] + for i in range(num_frames): + input_poses.append( + torch.cat([poses[i, :3, 0:1], poses[i, :3, 1:2], poses[i, :3, 3:4]], dim=-1) + ) + input_poses = torch.stack(input_poses) # (num_frames, 3, 3) + + # Prepare Gaussian kernel + gaussian_kernel = cv2.getGaussianKernel(ksize=k_size, sigma=-1).astype(np.float32).squeeze() + gaussian_kernel = torch.tensor(gaussian_kernel, dtype=dtype, device=device).view(1, 1, -1) + pad = k_size // 2 + + output_vectors = [] + for idx in range(3): # For r1, r2, t + vec = ( + input_poses[:, :, idx].T.unsqueeze(0).unsqueeze(0) + ) # (1, 1, 3, num_frames) -> (1, 1, 3, num_frames) + # But actually, we want (batch=3, channel=1, width=num_frames) + # So: + vec = input_poses[:, :, idx].T.unsqueeze(1) # (3, 1, num_frames) + vec_padded = F.pad(vec, (pad, pad), mode="reflect") + filtered = F.conv1d(vec_padded, gaussian_kernel) + output_vectors.append(filtered.squeeze(1).T) # (num_frames, 3) + + output_r1, output_r2, output_t = output_vectors # Each is (num_frames, 3) + + # Normalize r1 and r2 + output_r1 = output_r1 / output_r1.norm(dim=-1, keepdim=True) + output_r2 = output_r2 / output_r2.norm(dim=-1, keepdim=True) + + output_poses = [] + for i in range(num_frames): + output_r3 = torch.linalg.cross(output_r1[i], output_r2[i]) + render_pose = torch.cat( + [ + output_r1[i].unsqueeze(-1), + output_r2[i].unsqueeze(-1), + output_r3.unsqueeze(-1), + output_t[i].unsqueeze(-1), + ], + dim=-1, + ) + output_poses.append(render_pose[:3, :]) + output_poses = as_homogeneous(torch.stack(output_poses, dim=0)) + + return output_poses + + +@torch.no_grad() +def render_wander_path( + cam2world: torch.Tensor, + intrinsic: torch.Tensor, + h: int, + w: int, + num_frames: int = 120, + max_disp: float = 48.0, +): + device, dtype = cam2world.device, cam2world.dtype + fx = intrinsic[0, 0] * w + r = max_disp / fx + th = torch.linspace(0, 2.0 * torch.pi, steps=num_frames, device=device, dtype=dtype) + x = r * torch.sin(th) + yz = r * torch.cos(th) / 3.0 + T = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(num_frames, 1, 1) + T[:, :3, 3] = torch.stack([x, yz, yz], dim=-1) * -1.0 + c2ws = cam2world.unsqueeze(0) @ T + # Start at reference pose and end back at reference pose + c2ws = torch.cat([cam2world.unsqueeze(0), c2ws, cam2world.unsqueeze(0)], dim=0) + Ks = intrinsic.unsqueeze(0).repeat(c2ws.shape[0], 1, 1) + return c2ws, Ks + + +@torch.no_grad() +def render_dolly_zoom_path( + cam2world: torch.Tensor, + intrinsic: torch.Tensor, + h: int, + w: int, + num_frames: int = 120, + max_disp: float = 0.1, + D_focus: float = 10.0, +): + device, dtype = cam2world.device, cam2world.dtype + fx0, fy0 = intrinsic[0, 0] * w, intrinsic[1, 1] * h + t = torch.linspace(0.0, 2.0, steps=num_frames, device=device, dtype=dtype) + z = 0.5 * (1.0 - torch.cos(torch.pi * t)) * max_disp + T = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(num_frames, 1, 1) + T[:, 2, 3] = -z + c2ws = cam2world.unsqueeze(0) @ T + Df = torch.as_tensor(D_focus, device=device, dtype=dtype) + scale = (Df / (Df + z)).clamp(min=1e-6) + Ks = intrinsic.unsqueeze(0).repeat(num_frames, 1, 1) + Ks[:, 0, 0] = (fx0 * scale) / w + Ks[:, 1, 1] = (fy0 * scale) / h + return c2ws, Ks + + +@torch.no_grad() +def interpolate_intrinsics( + initial: torch.Tensor, # "*#batch 3 3" + final: torch.Tensor, # "*#batch 3 3" + t: torch.Tensor, # " time_step" +) -> torch.Tensor: # "*batch time_step 3 3" + initial = rearrange(initial, "... i j -> ... () i j") + final = rearrange(final, "... i j -> ... () i j") + t = rearrange(t, "t -> t () ()") + return initial + (final - initial) * t + + +def intersect_rays( + a_origins: torch.Tensor, # "*#batch dim" + a_directions: torch.Tensor, # "*#batch dim" + b_origins: torch.Tensor, # "*#batch dim" + b_directions: torch.Tensor, # "*#batch dim" +) -> torch.Tensor: # "*batch dim" + """Compute the least-squares intersection of rays. Uses the math from here: + https://math.stackexchange.com/a/1762491/286022 + """ + + # Broadcast and stack the tensors. + a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors( + a_origins, a_directions, b_origins, b_directions + ) + origins = torch.stack((a_origins, b_origins), dim=-2) + directions = torch.stack((a_directions, b_directions), dim=-2) + + # Compute n_i * n_i^T - eye(3) from the equation. + n = einsum(directions, directions, "... n i, ... n j -> ... n i j") + n = n - torch.eye(3, dtype=origins.dtype, device=origins.device) + + # Compute the left-hand side of the equation. + lhs = reduce(n, "... n i j -> ... i j", "sum") + + # Compute the right-hand side of the equation. + rhs = einsum(n, origins, "... n i j, ... n j -> ... n i") + rhs = reduce(rhs, "... n i -> ... i", "sum") + + # Left-matrix-multiply both sides by the inverse of lhs to find p. + return torch.linalg.lstsq(lhs, rhs).solution + + +def normalize(a: torch.Tensor) -> torch.Tensor: # "*#batch dim" -> "*#batch dim" + return a / a.norm(dim=-1, keepdim=True) + + +def generate_coordinate_frame( + y: torch.Tensor, # "*#batch 3" + z: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 3 3" + """Generate a coordinate frame given perpendicular, unit-length Y and Z vectors.""" + y, z = torch.broadcast_tensors(y, z) + return torch.stack([y.cross(z, dim=-1), y, z], dim=-1) + + +def generate_rotation_coordinate_frame( + a: torch.Tensor, # "*#batch 3" + b: torch.Tensor, # "*#batch 3" + eps: float = 1e-4, +) -> torch.Tensor: # "*batch 3 3" + """Generate a coordinate frame where the Y direction is normal to the plane defined + by unit vectors a and b. The other axes are arbitrary.""" + device = a.device + + # Replace every entry in b that's parallel to the corresponding entry in a with an + # arbitrary vector. + b = b.detach().clone() + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 0, 1], dtype=b.dtype, device=device) + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 1, 0], dtype=b.dtype, device=device) + + # Generate the coordinate frame. The initial cross product defines the plane. + return generate_coordinate_frame(normalize(torch.linalg.cross(a, b)), a) + + +def matrix_to_euler( + rotations: torch.Tensor, # "*batch 3 3" + pattern: str, +) -> torch.Tensor: # "*batch 3" + *batch, _, _ = rotations.shape + rotations = rotations.reshape(-1, 3, 3) + angles_np = R.from_matrix(rotations.detach().cpu().numpy()).as_euler(pattern) + rotations = torch.tensor(angles_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3) + + +def euler_to_matrix( + rotations: torch.Tensor, # "*batch 3" + pattern: str, +) -> torch.Tensor: # "*batch 3 3" + *batch, _ = rotations.shape + rotations = rotations.reshape(-1, 3) + matrix_np = R.from_euler(pattern, rotations.detach().cpu().numpy()).as_matrix() + rotations = torch.tensor(matrix_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3, 3) + + +def extrinsics_to_pivot_parameters( + extrinsics: torch.Tensor, # "*#batch 4 4" + pivot_coordinate_frame: torch.Tensor, # "*#batch 3 3" + pivot_point: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 5" + """Convert the extrinsics to a representation with 5 degrees of freedom: + 1. Distance from pivot point in the "X" (look cross pivot axis) direction. + 2. Distance from pivot point in the "Y" (pivot axis) direction. + 3. Distance from pivot point in the Z (look) direction + 4. Angle in plane + 5. Twist (rotation not in plane) + """ + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + # Compute the translation elements of the pivot parametrization. + translation_frame = generate_coordinate_frame(pivot_axis, extrinsics[..., :3, 2]) + origin = extrinsics[..., :3, 3] + delta = pivot_point - origin + translation = einsum(translation_frame, delta, "... i j, ... i -> ... j") + + # Add the rotation elements of the pivot parametrization. + inverted = pivot_coordinate_frame.inverse() @ extrinsics[..., :3, :3] + y, _, z = matrix_to_euler(inverted, "YXZ").unbind(dim=-1) + + return torch.cat([translation, y[..., None], z[..., None]], dim=-1) + + +def pivot_parameters_to_extrinsics( + parameters: torch.Tensor, # "*#batch 5" + pivot_coordinate_frame: torch.Tensor, # "*#batch 3 3" + pivot_point: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 4 4" + translation, y, z = parameters.split((3, 1, 1), dim=-1) + + euler = torch.cat((y, torch.zeros_like(y), z), dim=-1) + rotation = pivot_coordinate_frame @ euler_to_matrix(euler, "YXZ") + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + translation_frame = generate_coordinate_frame(pivot_axis, rotation[..., :3, 2]) + delta = einsum(translation_frame, translation, "... i j, ... j -> ... i") + origin = pivot_point - delta + + *batch, _ = origin.shape + extrinsics = torch.eye(4, dtype=parameters.dtype, device=parameters.device) + extrinsics = extrinsics.broadcast_to((*batch, 4, 4)).clone() + extrinsics[..., 3, 3] = 1 + extrinsics[..., :3, :3] = rotation + extrinsics[..., :3, 3] = origin + return extrinsics + + +def interpolate_circular( + a: torch.Tensor, # "*#batch" + b: torch.Tensor, # "*#batch" + t: torch.Tensor, # "*#batch" +) -> torch.Tensor: # " *batch" + a, b, t = torch.broadcast_tensors(a, b, t) + + tau = 2 * torch.pi + a = a % tau + b = b % tau + + # Consider piecewise edge cases. + d = (b - a).abs() + a_left = a - tau + d_left = (b - a_left).abs() + a_right = a + tau + d_right = (b - a_right).abs() + use_d = (d < d_left) & (d < d_right) + use_d_left = (d_left < d_right) & (~use_d) + use_d_right = (~use_d) & (~use_d_left) + + result = a + (b - a) * t + result[use_d_left] = (a_left + (b - a_left) * t)[use_d_left] + result[use_d_right] = (a_right + (b - a_right) * t)[use_d_right] + + return result + + +def interpolate_pivot_parameters( + initial: torch.Tensor, # "*#batch 5" + final: torch.Tensor, # "*#batch 5" + t: torch.Tensor, # " time_step" +) -> torch.Tensor: # "*batch time_step 5" + initial = rearrange(initial, "... d -> ... () d") + final = rearrange(final, "... d -> ... () d") + t = rearrange(t, "t -> t ()") + ti, ri = initial.split((3, 2), dim=-1) + tf, rf = final.split((3, 2), dim=-1) + + t_lerp = ti + (tf - ti) * t + r_lerp = interpolate_circular(ri, rf, t) + + return torch.cat((t_lerp, r_lerp), dim=-1) + + +@torch.no_grad() +def interpolate_extrinsics( + initial: torch.Tensor, # "*#batch 4 4" + final: torch.Tensor, # "*#batch 4 4" + t: torch.Tensor, # " time_step" + eps: float = 1e-4, +) -> torch.Tensor: # "*batch time_step 4 4" + """Interpolate extrinsics by rotating around their "focus point," which is the + least-squares intersection between the look vectors of the initial and final + extrinsics. + """ + + initial = initial.type(torch.float64) + final = final.type(torch.float64) + t = t.type(torch.float64) + + # Based on the dot product between the look vectors, pick from one of two cases: + # 1. Look vectors are parallel: interpolate about their origins' midpoint. + # 3. Look vectors aren't parallel: interpolate about their focus point. + initial_look = initial[..., :3, 2] + final_look = final[..., :3, 2] + dot_products = einsum(initial_look, final_look, "... i, ... i -> ...") + parallel_mask = (dot_products.abs() - 1).abs() < eps + + # Pick focus points. + initial_origin = initial[..., :3, 3] + final_origin = final[..., :3, 3] + pivot_point = 0.5 * (initial_origin + final_origin) + pivot_point[~parallel_mask] = intersect_rays( + initial_origin[~parallel_mask], + initial_look[~parallel_mask], + final_origin[~parallel_mask], + final_look[~parallel_mask], + ) + + # Convert to pivot parameters. + pivot_frame = generate_rotation_coordinate_frame(initial_look, final_look, eps=eps) + initial_params = extrinsics_to_pivot_parameters(initial, pivot_frame, pivot_point) + final_params = extrinsics_to_pivot_parameters(final, pivot_frame, pivot_point) + + # Interpolate the pivot parameters. + interpolated_params = interpolate_pivot_parameters(initial_params, final_params, t) + + # Convert back. + return pivot_parameters_to_extrinsics( + interpolated_params.type(torch.float32), + rearrange(pivot_frame, "... i j -> ... () i j").type(torch.float32), + rearrange(pivot_point, "... xyz -> ... () xyz").type(torch.float32), + ) + + +@torch.no_grad() +def generate_wobble_transformation( + radius: torch.Tensor, # "*#batch" + t: torch.Tensor, # " time_step" + num_rotations: int = 1, + scale_radius_with_t: bool = True, +) -> torch.Tensor: # "*batch time_step 4 4"]: + # Generate a translation in the image plane. + tf = torch.eye(4, dtype=torch.float32, device=t.device) + tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone() + radius = radius[..., None] + if scale_radius_with_t: + radius = radius * t + tf[..., 0, 3] = torch.sin(2 * torch.pi * num_rotations * t) * radius + tf[..., 1, 3] = -torch.cos(2 * torch.pi * num_rotations * t) * radius + return tf + + +@torch.no_grad() +def render_wobble_inter_path( + cam2world: torch.Tensor, intr_normed: torch.Tensor, inter_len: int, n_skip: int = 3 +): + """ + cam2world: [batch, 4, 4], + intr_normed: [batch, 3, 3] + """ + frame_per_round = n_skip * inter_len + num_rotations = 1 + + t = torch.linspace(0, 1, frame_per_round, dtype=torch.float32, device=cam2world.device) + # t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + tgt_c2w_b = [] + tgt_intr_b = [] + for b_idx in range(cam2world.shape[0]): + tgt_c2w = [] + tgt_intr = [] + for cur_idx in range(0, cam2world.shape[1] - n_skip, n_skip): + origin_a = cam2world[b_idx, cur_idx, :3, 3] + origin_b = cam2world[b_idx, cur_idx + n_skip, :3, 3] + delta = (origin_a - origin_b).norm(dim=-1) + if cur_idx == 0: + delta_prev = delta + else: + delta = (delta_prev + delta) / 2 + delta_prev = delta + tf = generate_wobble_transformation( + radius=delta * 0.5, + t=t, + num_rotations=num_rotations, + scale_radius_with_t=False, + ) + cur_extrs = ( + interpolate_extrinsics( + cam2world[b_idx, cur_idx], + cam2world[b_idx, cur_idx + n_skip], + t, + ) + @ tf + ) + tgt_c2w.append(cur_extrs[(0 if cur_idx == 0 else 1) :]) + tgt_intr.append( + interpolate_intrinsics( + intr_normed[b_idx, cur_idx], + intr_normed[b_idx, cur_idx + n_skip], + t, + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_c2w_b.append(torch.cat(tgt_c2w)) + tgt_intr_b.append(torch.cat(tgt_intr)) + tgt_c2w = torch.stack(tgt_c2w_b) # b v 4 4 + tgt_intr = torch.stack(tgt_intr_b) # b v 3 3 + return tgt_c2w, tgt_intr diff --git a/depth_anything_3/utils/constants.py b/depth_anything_3/utils/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..25c330ee00820f44df05aa624f0aa4763afe1164 --- /dev/null +++ b/depth_anything_3/utils/constants.py @@ -0,0 +1,270 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +DEFAULT_MODEL = "depth-anything/DA3NESTED-GIANT-LARGE-1.1" +DEFAULT_EXPORT_DIR = "workspace/gallery/scene" +DEFAULT_GALLERY_DIR = "workspace/gallery" +DEFAULT_GRADIO_DIR = "workspace/gradio" +THRESH_FOR_REF_SELECTION = 3 + +# ============================================================================= +# Benchmark Evaluation Constants +# ============================================================================= + +# Default evaluation workspace directory +DEFAULT_EVAL_WORKSPACE = "workspace/evaluation" + +# Default reference view selection strategy for evaluation +# Use "first" for consistent and reproducible evaluation results +# Other options: "saddle_balanced", "auto", "mid" +EVAL_REF_VIEW_STRATEGY = "first" + +# ----------------------------------------------------------------------------- +# DTU Dataset Configuration +# Reference: https://roboimagedata.compute.dtu.dk/ +# Note: DepthAnything3 was never trained on any images from DTU. +# ----------------------------------------------------------------------------- + +# Root directory for DTU evaluation data (MVSNet format) +# Download from: https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view +DTU_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu" + +# List of DTU evaluation scenes +DTU_SCENES = [ + "scan1", + "scan4", + "scan9", + "scan10", + "scan11", + "scan12", + "scan13", + "scan15", + "scan23", + "scan24", + "scan29", + "scan32", + "scan33", + "scan34", + "scan48", + "scan49", + "scan62", + "scan75", + "scan77", + "scan110", + "scan114", + "scan118", +] + +# Point cloud fusion hyperparameters +DTU_DIST_THRESH = 0.2 # Distance threshold for geometric consistency (mm) +DTU_NUM_CONSIST = 4 # Minimum number of consistent views for a point +DTU_MAX_POINTS = 4_000_000 # Maximum points in fused point cloud + +# 3D reconstruction evaluation hyperparameters +DTU_DOWN_DENSE = 0.2 # Downsample density for evaluation (mm) +DTU_PATCH_SIZE = 60 # Patch size for boundary handling +DTU_MAX_DIST = 20 # Outlier threshold for accuracy/completeness (mm) + +# ----------------------------------------------------------------------------- +# DTU-64 Dataset Configuration (Pose Evaluation Only) +# This is a subset of DTU with 64 images per scene for pose evaluation. +# Note: This dataset is ONLY for pose evaluation, not 3D reconstruction. +# ----------------------------------------------------------------------------- + +# Root directory for DTU-64 evaluation data +DTU64_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu64" +DTU64_CAMERA_ROOT = "workspace/benchmark_dataset/dtu64/Cameras" + +# List of DTU-64 evaluation scenes (13 scenes) +DTU64_SCENES = [ + "scan105", + "scan114", + "scan118", + "scan122", + "scan24", + "scan37", + "scan40", + "scan55", + "scan63", + "scan65", + "scan69", + "scan83", + "scan97", +] + +# ----------------------------------------------------------------------------- +# ETH3D Dataset Configuration +# Reference: https://www.eth3d.net/ +# High-resolution multi-view stereo benchmark with laser-scanned ground truth. +# Note: DepthAnything3 was never trained on any images from ETH3D. +# ----------------------------------------------------------------------------- + +# Root directory for ETH3D evaluation data +ETH3D_EVAL_DATA_ROOT = "workspace/benchmark_dataset/eth3d" + +# List of ETH3D evaluation scenes (indoor and outdoor) +ETH3D_SCENES = [ + "courtyard", + "electro", + "kicker", + "pipes", + "relief", + # "terrace", # Excluded: known issues + "delivery_area", + "facade", + # "meadow", # Excluded: known issues + "office", + "playground", + "relief_2", + "terrains", +] + +# Images to filter out (known problematic views per scene) +ETH3D_FILTER_KEYS = { + "delivery_area": ["711.JPG", "712.JPG", "713.JPG", "714.JPG"], + "electro": ["9289.JPG", "9290.JPG", "9291.JPG", "9292.JPG", "9293.JPG", "9298.JPG"], + "playground": ["587.JPG", "588.JPG", "589.JPG", "590.JPG", "591.JPG", "592.JPG"], + "relief": [ + "427.JPG", "428.JPG", "429.JPG", "430.JPG", "431.JPG", "432.JPG", + "433.JPG", "434.JPG", "435.JPG", "436.JPG", "437.JPG", "438.JPG", + ], + "relief_2": [ + "458.JPG", "459.JPG", "460.JPG", "461.JPG", "462.JPG", "463.JPG", + "464.JPG", "465.JPG", "466.JPG", "467.JPG", "468.JPG", + ], +} + +# TSDF fusion hyperparameters (scaled for outdoor scenes) +ETH3D_VOXEL_LENGTH = 4.0 / 512.0 * 5 # Voxel size for TSDF (meters) +ETH3D_SDF_TRUNC = 0.04 * 5 # SDF truncation distance (meters) +ETH3D_MAX_DEPTH = 100000.0 # Maximum depth for integration (effectively no truncation) + +# Point cloud sampling +ETH3D_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +ETH3D_EVAL_THRESHOLD = 0.05 * 5 # Distance threshold for precision/recall (meters) +ETH3D_DOWN_SAMPLE = 4.0 / 512.0 * 5 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# 7Scenes Dataset Configuration +# ============================================================================== +# Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/ +# Note: Indoor RGB-D dataset with ground truth poses and meshes. + +# Root directory for 7Scenes evaluation data +SEVENSCENES_EVAL_DATA_ROOT = "workspace/benchmark_dataset/7scenes" + +# List of 7Scenes evaluation scenes +SEVENSCENES_SCENES = [ + "chess", + "fire", + "heads", + "office", + "pumpkin", + "redkitchen", + "stairs", +] + +# Fixed camera intrinsics for 7Scenes (all images share same intrinsics) +SEVENSCENES_FX = 585.0 +SEVENSCENES_FY = 585.0 +SEVENSCENES_CX = 320.0 +SEVENSCENES_CY = 240.0 + +# TSDF fusion hyperparameters (indoor scenes, smaller voxels) +SEVENSCENES_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) +SEVENSCENES_SDF_TRUNC = 0.04 # SDF truncation distance (meters) +SEVENSCENES_MAX_DEPTH = 1000000.0 # Maximum depth for integration (no truncation) + +# Point cloud sampling +SEVENSCENES_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +SEVENSCENES_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +SEVENSCENES_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# ScanNet++ Dataset Configuration +# ============================================================================== +# Reference: https://kaldir.vc.in.tum.de/scannetpp/ +# Note: High-quality indoor RGB-D dataset with iPhone and DSLR images. + +# Root directory for ScanNet++ evaluation data +SCANNETPP_EVAL_DATA_ROOT = "workspace/benchmark_dataset/scannetpp" + +# List of ScanNet++ evaluation scenes +SCANNETPP_SCENES = [ + "09c1414f1b", + "1ada7a0617", + "40aec5fffa", + "3e8bba0176", + "acd95847c5", + "578511c8a9", + "5f99900f09", + "c4c04e6d6c", + "f3d64c30f8", + "7bc286c1b6", + "c5439f4607", + "286b55a2bf", + "fb5a96b1a2", + "7831862f02", + "38d58a7a31", + "bde1e479ad", + "9071e139d9", + "21d970d8de", + "bcd2436daf", + "cc5237fd77", +] + +# Input resolution for ScanNet++ (after undistortion and resize) +SCANNETPP_INPUT_H = 768 +SCANNETPP_INPUT_W = 1024 + +# TSDF fusion hyperparameters (indoor scenes) +SCANNETPP_VOXEL_LENGTH = 0.02 # Voxel size for TSDF (meters) +SCANNETPP_SDF_TRUNC = 0.15 # SDF truncation distance (meters) +SCANNETPP_MAX_DEPTH = 5.0 # Maximum depth for integration (meters) + +# Point cloud sampling +SCANNETPP_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +SCANNETPP_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +SCANNETPP_DOWN_SAMPLE = 0.02 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# HiRoom Dataset Configuration +# ============================================================================== +# Note: Indoor RGB-D dataset. + +# Root directory for HiRoom evaluation data +HIROOM_EVAL_DATA_ROOT = "workspace/benchmark_dataset/hiroom/data" +HIROOM_GT_ROOT_PATH = "workspace/benchmark_dataset/hiroom/fused_pcd" +HIROOM_SCENE_LIST_PATH = "workspace/benchmark_dataset/hiroom/selected_scene_list_val.txt" + +# TSDF fusion hyperparameters (indoor scenes) +HIROOM_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) +HIROOM_SDF_TRUNC = 0.04 # SDF truncation distance (meters) +HIROOM_MAX_DEPTH = 10000.0 # Maximum depth for integration (no truncation) + +# Point cloud sampling +HIROOM_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +HIROOM_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +HIROOM_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) diff --git a/depth_anything_3/utils/export/__init__.py b/depth_anything_3/utils/export/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7073989b5048b6e42a6246e9dde2025ca1e130d --- /dev/null +++ b/depth_anything_3/utils/export/__init__.py @@ -0,0 +1,75 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.export.gs import export_to_gs_ply, export_to_gs_video + +from .colmap import export_to_colmap +from .depth_vis import export_to_depth_vis +from .feat_vis import export_to_feat_vis +from .glb import export_to_glb +from .npz import export_to_mini_npz, export_to_npz + +# Canonical set of export formats this dispatcher understands. Treat this as the +# server-trusted enum for any request-facing validation (e.g. the HTTP backend) instead +# of letting a client-supplied export_format string reach the dispatcher unchecked. +SUPPORTED_EXPORT_FORMATS = frozenset( + { + "glb", + "mini_npz", + "npz", + "feat_vis", + "depth_vis", + "gs_ply", + "gs_video", + "colmap", + } +) + + +def export( + prediction: Prediction, + export_format: str, + export_dir: str, + **kwargs, +): + if "-" in export_format: + export_formats = export_format.split("-") + for export_format in export_formats: + export(prediction, export_format, export_dir, **kwargs) + return # Prevent falling through to single-format handling + + if export_format == "glb": + export_to_glb(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "mini_npz": + export_to_mini_npz(prediction, export_dir) + elif export_format == "npz": + export_to_npz(prediction, export_dir) + elif export_format == "feat_vis": + export_to_feat_vis(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "depth_vis": + export_to_depth_vis(prediction, export_dir) + elif export_format == "gs_ply": + export_to_gs_ply(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "gs_video": + export_to_gs_video(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "colmap": + export_to_colmap(prediction, export_dir, **kwargs.get(export_format, {})) + else: + raise ValueError(f"Unsupported export format: {export_format}") + + +__all__ = [ + export, +] diff --git a/depth_anything_3/utils/export/colmap.py b/depth_anything_3/utils/export/colmap.py new file mode 100644 index 0000000000000000000000000000000000000000..e57964cc3302d3f348d91d0b0ecb5d4fe8c0a54b --- /dev/null +++ b/depth_anything_3/utils/export/colmap.py @@ -0,0 +1,150 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import pycolmap +import cv2 as cv +import numpy as np + +from PIL import Image + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.logger import logger + +from .glb import _depths_to_world_points_with_colors + + +def export_to_colmap( + prediction: Prediction, + export_dir: str, + image_paths: list[str], + conf_thresh_percentile: float = 40.0, + process_res_method: str = "upper_bound_resize", +) -> None: + # 1. Data preparation + conf_thresh = np.percentile(prediction.conf, conf_thresh_percentile) + points, colors = _depths_to_world_points_with_colors( + prediction.depth, + prediction.intrinsics, + prediction.extrinsics, # w2c + prediction.processed_images, + prediction.conf, + conf_thresh, + ) + num_points = len(points) + logger.info(f"Exporting to COLMAP with {num_points} points") + num_frames = len(prediction.processed_images) + h, w = prediction.processed_images.shape[1:3] + points_xyf = _create_xyf(num_frames, h, w) + points_xyf = points_xyf[prediction.conf >= conf_thresh] + + # 2. Set Reconstruction + reconstruction = pycolmap.Reconstruction() + + point3d_ids = [] + for vidx in range(num_points): + point3d_id = reconstruction.add_point3D(points[vidx], pycolmap.Track(), colors[vidx]) + point3d_ids.append(point3d_id) + + for fidx in range(num_frames): + orig_w, orig_h = Image.open(image_paths[fidx]).size + + intrinsic = prediction.intrinsics[fidx] + if process_res_method.endswith("resize"): + intrinsic[:1] *= orig_w / w + intrinsic[1:2] *= orig_h / h + elif process_res_method == "crop": + raise NotImplementedError("COLMAP export for crop method is not implemented") + else: + raise ValueError(f"Unknown process_res_method: {process_res_method}") + + pycolmap_intri = np.array( + [intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2]] + ) + + extrinsic = prediction.extrinsics[fidx] + cam_from_world = pycolmap.Rigid3d(pycolmap.Rotation3d(extrinsic[:3, :3]), extrinsic[:3, 3]) + + # set and add camera + camera = pycolmap.Camera() + camera.camera_id = fidx + 1 + camera.model = pycolmap.CameraModelId.PINHOLE + camera.width = orig_w + camera.height = orig_h + camera.params = pycolmap_intri + reconstruction.add_camera(camera) + + # set and add rig (from camera) + rig = pycolmap.Rig() + rig.rig_id = camera.camera_id + rig.add_ref_sensor(camera.sensor_id) + reconstruction.add_rig(rig) + + # set image + image = pycolmap.Image() + image.image_id = fidx + 1 + image.camera_id = camera.camera_id + + # set and add frame (from image) + frame = pycolmap.Frame() + frame.frame_id = image.image_id + frame.rig_id = camera.camera_id + frame.add_data_id(image.data_id) + frame.rig_from_world = cam_from_world + reconstruction.add_frame(frame) + + # set point2d and update track + point2d_list = [] + points_in_frame = points_xyf[:, 2].astype(np.int32) == fidx + for vidx in np.where(points_in_frame)[0]: + point2d = points_xyf[vidx][:2] + point2d[0] *= orig_w / w + point2d[1] *= orig_h / h + point3d_id = point3d_ids[vidx] + point2d_list.append(pycolmap.Point2D(point2d, point3d_id)) + reconstruction.point3D(point3d_id).track.add_element( + image.image_id, len(point2d_list) - 1 + ) + + # set and add image + image.frame_id = image.image_id + image.name = os.path.basename(image_paths[fidx]) + image.points2D = pycolmap.Point2DList(point2d_list) + reconstruction.add_image(image) + + # 3. Export + reconstruction.write(export_dir) + + +def _create_xyf(num_frames, height, width): + """ + Creates a grid of pixel coordinates and frame indices (fidx) for all frames. + """ + # Create coordinate grids for a single frame + y_grid, x_grid = np.indices((height, width), dtype=np.int32) + x_grid = x_grid[np.newaxis, :, :] + y_grid = y_grid[np.newaxis, :, :] + + # Broadcast to all frames + x_coords = np.broadcast_to(x_grid, (num_frames, height, width)) + y_coords = np.broadcast_to(y_grid, (num_frames, height, width)) + + # Create frame indices and broadcast + f_idx = np.arange(num_frames, dtype=np.int32)[:, np.newaxis, np.newaxis] + f_coords = np.broadcast_to(f_idx, (num_frames, height, width)) + + # Stack coordinates and frame indices + points_xyf = np.stack((x_coords, y_coords, f_coords), axis=-1) + + return points_xyf diff --git a/depth_anything_3/utils/export/depth_vis.py b/depth_anything_3/utils/export/depth_vis.py new file mode 100644 index 0000000000000000000000000000000000000000..8accc04e92985e26b8d78a56db80989be515aac7 --- /dev/null +++ b/depth_anything_3/utils/export/depth_vis.py @@ -0,0 +1,41 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import imageio +import numpy as np + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.visualize import visualize_depth + + +def export_to_depth_vis( + prediction: Prediction, + export_dir: str, +): + # Use prediction.processed_images, which is already processed image data + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + images_u8 = prediction.processed_images # (N,H,W,3) uint8 + + os.makedirs(os.path.join(export_dir, "depth_vis"), exist_ok=True) + for idx in range(prediction.depth.shape[0]): + depth_vis = visualize_depth(prediction.depth[idx]) + image_vis = images_u8[idx] + depth_vis = depth_vis.astype(np.uint8) + image_vis = image_vis.astype(np.uint8) + vis_image = np.concatenate([image_vis, depth_vis], axis=1) + save_path = os.path.join(export_dir, f"depth_vis/{idx:04d}.jpg") + imageio.imwrite(save_path, vis_image, quality=95) diff --git a/depth_anything_3/utils/export/feat_vis.py b/depth_anything_3/utils/export/feat_vis.py new file mode 100644 index 0000000000000000000000000000000000000000..c2507f6e675c0a546311ae292b9f6859d218d8e2 --- /dev/null +++ b/depth_anything_3/utils/export/feat_vis.py @@ -0,0 +1,75 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import subprocess +import cv2 +import imageio +import numpy as np +from tqdm.auto import tqdm + +from depth_anything_3.utils.logger import logger +from depth_anything_3.utils.parallel_utils import async_call +from depth_anything_3.utils.pca_utils import PCARGBVisualizer + + +@async_call +def export_to_feat_vis( + prediction, + export_dir, + fps=15, +): + """Export feature visualization with PCA. + + Args: + prediction: Model prediction containing feature maps + export_dir: Directory to export results + fps: Frame rate for output video (default: 15) + """ + out_dir = os.path.join(export_dir, "feat_vis") + os.makedirs(out_dir, exist_ok=True) + + images = prediction.processed_images + for k, v in prediction.aux.items(): + if not k.startswith("feat_layer_"): + continue + os.makedirs(os.path.join(out_dir, k), exist_ok=True) + viz = PCARGBVisualizer(basis_mode="fixed", percentile_mode="global", clip_percent=10.0) + viz.fit_reference(v) + feats_vis = viz.transform_video(v) + for idx in tqdm(range(len(feats_vis))): + img = images[idx] + feat_vis = (feats_vis[idx] * 255).astype(np.uint8) + feat_vis = cv2.resize( + feat_vis, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST + ) + save_path = os.path.join(out_dir, f"{k}/{idx:06d}.jpg") + save = np.concatenate([img, feat_vis], axis=1) + imageio.imwrite(save_path, save, quality=95) + cmd = [ + "ffmpeg", + "-loglevel", "error", + "-hide_banner", + "-y", + "-framerate", str(fps), + "-start_number", "0", + "-i", os.path.join(out_dir, k, "%06d.jpg"), + "-c:v", "libx264", + "-pix_fmt", "yuv420p", + os.path.join(out_dir, f"{k}.mp4"), + ] + try: + subprocess.run(cmd, shell=False, check=False) + except OSError as e: + logger.warn(f"Failed to run ffmpeg for feature visualization export: {e}") diff --git a/depth_anything_3/utils/export/glb.py b/depth_anything_3/utils/export/glb.py new file mode 100644 index 0000000000000000000000000000000000000000..782106a0689cf763eefaec3c327a625c98e342cb --- /dev/null +++ b/depth_anything_3/utils/export/glb.py @@ -0,0 +1,437 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os +import shutil +import numpy as np +import trimesh + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.logger import logger + +from .depth_vis import export_to_depth_vis + + +def set_sky_depth(prediction: Prediction, sky_mask: np.ndarray, sky_depth_def: float = 98.0): + non_sky_mask = ~sky_mask + valid_depth = prediction.depth[non_sky_mask] + if valid_depth.size > 0: + max_depth = np.percentile(valid_depth, sky_depth_def) + prediction.depth[sky_mask] = max_depth + + +def get_conf_thresh( + prediction: Prediction, + sky_mask: np.ndarray, + conf_thresh: float, + conf_thresh_percentile: float = 10.0, + ensure_thresh_percentile: float = 90.0, +): + if sky_mask is not None and (~sky_mask).sum() > 10: + conf_pixels = prediction.conf[~sky_mask] + else: + conf_pixels = prediction.conf + lower = np.percentile(conf_pixels, conf_thresh_percentile) + upper = np.percentile(conf_pixels, ensure_thresh_percentile) + conf_thresh = min(max(conf_thresh, lower), upper) + return conf_thresh + + +def export_to_glb( + prediction: Prediction, + export_dir: str, + num_max_points: int = 1_000_000, + conf_thresh: float = 1.05, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + conf_thresh_percentile: float = 40.0, + ensure_thresh_percentile: float = 90.0, + sky_depth_def: float = 98.0, + show_cameras: bool = True, + camera_size: float = 0.03, + export_depth_vis: bool = True, +) -> str: + """Generate a 3D point cloud and camera wireframes and export them as a ``.glb`` file. + + The function builds a point cloud from the predicted depth maps, aligns it to the + first camera in glTF coordinates (X-right, Y-up, Z-backward), optionally draws + camera wireframes, and writes the result to ``scene.glb``. Auxiliary assets such as + depth visualizations can also be generated alongside the main export. + + Args: + prediction: Model prediction containing depth, confidence, intrinsics, extrinsics, + and pre-processed images. + export_dir: Output directory where the glTF assets will be written. + num_max_points: Maximum number of points retained after downsampling. + conf_thresh: Base confidence threshold used before percentile adjustments. + filter_black_bg: Mark near-black background pixels for removal during confidence filtering. + filter_white_bg: Mark near-white background pixels for removal during confidence filtering. + conf_thresh_percentile: Lower percentile used when adapting the confidence threshold. + ensure_thresh_percentile: Upper percentile clamp for the adaptive threshold. + sky_depth_def: Percentile used to fill sky pixels with plausible depth values. + show_cameras: Whether to render camera wireframes in the exported scene. + camera_size: Relative camera wireframe scale as a fraction of the scene diagonal. + export_depth_vis: Whether to export raster depth visualisations alongside the glTF. + + Returns: + Path to the exported ``scene.glb`` file. + """ + # 1) Use prediction.processed_images, which is already processed image data + assert ( + prediction.processed_images is not None + ), "Export to GLB: prediction.processed_images is required but not available" + assert ( + prediction.depth is not None + ), "Export to GLB: prediction.depth is required but not available" + assert ( + prediction.intrinsics is not None + ), "Export to GLB: prediction.intrinsics is required but not available" + assert ( + prediction.extrinsics is not None + ), "Export to GLB: prediction.extrinsics is required but not available" + assert ( + prediction.conf is not None + ), "Export to GLB: prediction.conf is required but not available" + logger.info(f"conf_thresh_percentile: {conf_thresh_percentile}") + logger.info(f"num max points: {num_max_points}") + logger.info(f"Exporting to GLB with num_max_points: {num_max_points}") + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + images_u8 = prediction.processed_images # (N,H,W,3) uint8 + + # 2) Sky processing (if sky_mask is provided) + if getattr(prediction, "sky_mask", None) is not None: + set_sky_depth(prediction, prediction.sky_mask, sky_depth_def) + + # 3) Confidence threshold (if no conf, then no filtering) + if filter_black_bg: + prediction.conf[(prediction.processed_images < 16).all(axis=-1)] = 1.0 + if filter_white_bg: + prediction.conf[(prediction.processed_images >= 240).all(axis=-1)] = 1.0 + conf_thr = get_conf_thresh( + prediction, + getattr(prediction, "sky_mask", None), + conf_thresh, + conf_thresh_percentile, + ensure_thresh_percentile, + ) + + # 4) Back-project to world coordinates and get colors (world frame) + points, colors = _depths_to_world_points_with_colors( + prediction.depth, + prediction.intrinsics, + prediction.extrinsics, # w2c + images_u8, + prediction.conf, + conf_thr, + ) + + # 5) Based on first camera orientation + glTF axis system, center by point cloud, + # construct alignment transform, and apply to point cloud + A = _compute_alignment_transform_first_cam_glTF_center_by_points( + prediction.extrinsics[0], points + ) # (4,4) + + if points.shape[0] > 0: + points = trimesh.transform_points(points, A) + + # 6) Clean + downsample + points, colors = _filter_and_downsample(points, colors, num_max_points) + + # 7) Assemble scene (add point cloud first) + scene = trimesh.Scene() + if scene.metadata is None: + scene.metadata = {} + scene.metadata["hf_alignment"] = A # For camera wireframes and external reuse + + if points.shape[0] > 0: + pc = trimesh.points.PointCloud(vertices=points, colors=colors) + scene.add_geometry(pc) + + # 8) Draw cameras (wireframe pyramids), using the same transform A + if show_cameras and prediction.intrinsics is not None and prediction.extrinsics is not None: + scene_scale = _estimate_scene_scale(points, fallback=1.0) + H, W = prediction.depth.shape[1:] + _add_cameras_to_scene( + scene=scene, + K=prediction.intrinsics, + ext_w2c=prediction.extrinsics, + image_sizes=[(H, W)] * prediction.depth.shape[0], + scale=scene_scale * camera_size, + ) + + # 9) Export + os.makedirs(export_dir, exist_ok=True) + out_path = os.path.join(export_dir, "scene.glb") + scene.export(out_path) + + if export_depth_vis: + export_to_depth_vis(prediction, export_dir) + depth_vis_thumbnail = os.path.join(export_dir, "depth_vis", "0000.jpg") + if os.path.isfile(depth_vis_thumbnail): + shutil.copy2(depth_vis_thumbnail, os.path.join(export_dir, "scene.jpg")) + else: + logger.warn(f"Depth visualization thumbnail not found at {depth_vis_thumbnail}") + return out_path + + +# ========================= +# utilities +# ========================= + + +def _as_homogeneous44(ext: np.ndarray) -> np.ndarray: + """ + Accept (4,4) or (3,4) extrinsic parameters, return (4,4) homogeneous matrix. + """ + if ext.shape == (4, 4): + return ext + if ext.shape == (3, 4): + H = np.eye(4, dtype=ext.dtype) + H[:3, :4] = ext + return H + raise ValueError(f"extrinsic must be (4,4) or (3,4), got {ext.shape}") + + +def _depths_to_world_points_with_colors( + depth: np.ndarray, + K: np.ndarray, + ext_w2c: np.ndarray, + images_u8: np.ndarray, + conf: np.ndarray | None, + conf_thr: float, +) -> tuple[np.ndarray, np.ndarray]: + """ + For each frame, transform (u,v,1) through K^{-1} to get rays, + multiply by depth to camera frame, then use (w2c)^{-1} to transform to world frame. + Simultaneously extract colors. + """ + N, H, W = depth.shape + us, vs = np.meshgrid(np.arange(W), np.arange(H)) + ones = np.ones_like(us) + pix = np.stack([us, vs, ones], axis=-1).reshape(-1, 3) # (H*W,3) + + pts_all, col_all = [], [] + + for i in range(N): + d = depth[i] # (H,W) + valid = np.isfinite(d) & (d > 0) + if conf is not None: + valid &= conf[i] >= conf_thr + if not np.any(valid): + continue + + d_flat = d.reshape(-1) + vidx = np.flatnonzero(valid.reshape(-1)) + + K_inv = np.linalg.inv(K[i]) # (3,3) + c2w = np.linalg.inv(_as_homogeneous44(ext_w2c[i])) # (4,4) + + rays = K_inv @ pix[vidx].T # (3,M) + Xc = rays * d_flat[vidx][None, :] # (3,M) + Xc_h = np.vstack([Xc, np.ones((1, Xc.shape[1]))]) + Xw = (c2w @ Xc_h)[:3].T.astype(np.float32) # (M,3) + + cols = images_u8[i].reshape(-1, 3)[vidx].astype(np.uint8) # (M,3) + + pts_all.append(Xw) + col_all.append(cols) + + if len(pts_all) == 0: + return np.zeros((0, 3), dtype=np.float32), np.zeros((0, 3), dtype=np.uint8) + + return np.concatenate(pts_all, 0), np.concatenate(col_all, 0) + + +def _filter_and_downsample(points: np.ndarray, colors: np.ndarray, num_max: int): + if points.shape[0] == 0: + return points, colors + finite = np.isfinite(points).all(axis=1) + points, colors = points[finite], colors[finite] + if points.shape[0] > num_max: + idx = np.random.choice(points.shape[0], num_max, replace=False) + points, colors = points[idx], colors[idx] + return points, colors + + +def _estimate_scene_scale(points: np.ndarray, fallback: float = 1.0) -> float: + if points.shape[0] < 2: + return fallback + lo = np.percentile(points, 5, axis=0) + hi = np.percentile(points, 95, axis=0) + diag = np.linalg.norm(hi - lo) + return float(diag if np.isfinite(diag) and diag > 0 else fallback) + + +def _compute_alignment_transform_first_cam_glTF_center_by_points( + ext_w2c0: np.ndarray, + points_world: np.ndarray, +) -> np.ndarray: + """Computes the transformation matrix to align the scene with glTF standards. + + This function calculates a 4x4 homogeneous matrix that centers the scene's + point cloud and transforms its coordinate system from the computer vision (CV) + standard to the glTF standard. + + The transformation process involves three main steps: + 1. **Initial Alignment**: Orients the world coordinate system to match the + first camera's view (x-right, y-down, z-forward). + 2. **Coordinate System Conversion**: Converts the CV camera frame to the + glTF frame (x-right, y-up, z-backward) by flipping the Y and Z axes. + 3. **Centering**: Translates the entire scene so that the median of the + point cloud becomes the new origin (0,0,0). + + Returns: + A 4x4 homogeneous transformation matrix (torch.Tensor or np.ndarray) + that applies these transformations. A: X' = A @ [X;1] + """ + + w2c0 = _as_homogeneous44(ext_w2c0).astype(np.float64) + + # CV -> glTF axis transformation + M = np.eye(4, dtype=np.float64) + M[1, 1] = -1.0 # flip Y + M[2, 2] = -1.0 # flip Z + + # Don't center first + A_no_center = M @ w2c0 + + # Calculate point cloud center in new coordinate system (use median to resist outliers) + if points_world.shape[0] > 0: + pts_tmp = trimesh.transform_points(points_world, A_no_center) + center = np.median(pts_tmp, axis=0) + else: + center = np.zeros(3, dtype=np.float64) + + T_center = np.eye(4, dtype=np.float64) + T_center[:3, 3] = -center + + A = T_center @ A_no_center + return A + + +def _add_cameras_to_scene( + scene: trimesh.Scene, + K: np.ndarray, + ext_w2c: np.ndarray, + image_sizes: list[tuple[int, int]], + scale: float, +) -> None: + """Draws camera frustums to visualize their position and orientation. + + This function renders each camera as a wireframe pyramid, originating from + the camera's center and extending to the corners of its imaging plane. + + It reads the 'hf_alignment' metadata from the scene to ensure the + wireframes are correctly aligned with the 3D point cloud. + """ + N = K.shape[0] + if N == 0: + return + + # Alignment matrix consistent with point cloud (use identity matrix if missing) + A = None + try: + A = scene.metadata.get("hf_alignment", None) if scene.metadata else None + except Exception: + A = None + if A is None: + A = np.eye(4, dtype=np.float64) + + for i in range(N): + H, W = image_sizes[i] + segs = _camera_frustum_lines(K[i], ext_w2c[i], W, H, scale) # (8,2,3) world frame + # Apply unified transformation + segs = trimesh.transform_points(segs.reshape(-1, 3), A).reshape(-1, 2, 3) + path = trimesh.load_path(segs) + color = _index_color_rgb(i, N) + if hasattr(path, "colors"): + path.colors = np.tile(color, (len(path.entities), 1)) + scene.add_geometry(path) + + +def _camera_frustum_lines( + K: np.ndarray, ext_w2c: np.ndarray, W: int, H: int, scale: float +) -> np.ndarray: + corners = np.array( + [ + [0, 0, 1.0], + [W - 1, 0, 1.0], + [W - 1, H - 1, 1.0], + [0, H - 1, 1.0], + ], + dtype=float, + ) # (4,3) + + K_inv = np.linalg.inv(K) + c2w = np.linalg.inv(_as_homogeneous44(ext_w2c)) + + # camera center in world + Cw = (c2w @ np.array([0, 0, 0, 1.0]))[:3] + + # rays -> z=1 plane points (camera frame) + rays = (K_inv @ corners.T).T + z = rays[:, 2:3] + z[z == 0] = 1.0 + plane_cam = (rays / z) * scale # (4,3) + + # to world + plane_w = [] + for p in plane_cam: + pw = (c2w @ np.array([p[0], p[1], p[2], 1.0]))[:3] + plane_w.append(pw) + plane_w = np.stack(plane_w, 0) # (4,3) + + segs = [] + # center to corners + for k in range(4): + segs.append(np.stack([Cw, plane_w[k]], 0)) + # rectangle edges + order = [0, 1, 2, 3, 0] + for a, b in zip(order[:-1], order[1:]): + segs.append(np.stack([plane_w[a], plane_w[b]], 0)) + + return np.stack(segs, 0) # (8,2,3) + + +def _index_color_rgb(i: int, n: int) -> np.ndarray: + h = (i + 0.5) / max(n, 1) + s, v = 0.85, 0.95 + r, g, b = _hsv_to_rgb(h, s, v) + return (np.array([r, g, b]) * 255).astype(np.uint8) + + +def _hsv_to_rgb(h: float, s: float, v: float) -> tuple[float, float, float]: + i = int(h * 6.0) + f = h * 6.0 - i + p = v * (1.0 - s) + q = v * (1.0 - f * s) + t = v * (1.0 - (1.0 - f) * s) + i = i % 6 + if i == 0: + r, g, b = v, t, p + elif i == 1: + r, g, b = q, v, p + elif i == 2: + r, g, b = p, v, t + elif i == 3: + r, g, b = p, q, v + elif i == 4: + r, g, b = t, p, v + else: + r, g, b = v, p, q + return r, g, b diff --git a/depth_anything_3/utils/export/gs.py b/depth_anything_3/utils/export/gs.py new file mode 100644 index 0000000000000000000000000000000000000000..90077cf25651c7977c1c1da320b7c0931fdb18ba --- /dev/null +++ b/depth_anything_3/utils/export/gs.py @@ -0,0 +1,154 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import Literal, Optional +import moviepy.editor as mpy +import torch + +from depth_anything_3.model.utils.gs_renderer import run_renderer_in_chunk_w_trj_mode +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.gsply_helpers import save_gaussian_ply +from depth_anything_3.utils.layout_helpers import hcat, vcat +from depth_anything_3.utils.visualize import vis_depth_map_tensor + +VIDEO_QUALITY_MAP = { + "low": {"crf": "28", "preset": "veryfast"}, + "medium": {"crf": "23", "preset": "medium"}, + "high": {"crf": "18", "preset": "slow"}, +} + + +def export_to_gs_ply( + prediction: Prediction, + export_dir: str, + gs_views_interval: Optional[ + int + ] = 1, # export GS every N views, useful for extremely dense inputs +): + gs_world = prediction.gaussians + pred_depth = torch.from_numpy(prediction.depth).unsqueeze(-1).to(gs_world.means) # v h w 1 + idx = 0 + os.makedirs(os.path.join(export_dir, "gs_ply"), exist_ok=True) + save_path = os.path.join(export_dir, f"gs_ply/{idx:04d}.ply") + if gs_views_interval is None: # select around 12 views in total + gs_views_interval = max(pred_depth.shape[0] // 12, 1) + save_gaussian_ply( + gaussians=gs_world, + save_path=save_path, + ctx_depth=pred_depth, + shift_and_scale=False, + save_sh_dc_only=True, + gs_views_interval=gs_views_interval, + inv_opacity=True, + prune_by_depth_percent=0.9, + prune_border_gs=True, + match_3dgs_mcmc_dev=False, + ) + + +def export_to_gs_video( + prediction: Prediction, + export_dir: str, + extrinsics: Optional[torch.Tensor] = None, # render views' world2cam, "b v 4 4" + intrinsics: Optional[torch.Tensor] = None, # render views' unnormed intrinsics, "b v 3 3" + out_image_hw: Optional[tuple[int, int]] = None, # render views' resolution, (h, w) + chunk_size: Optional[int] = 4, + trj_mode: Literal[ + "original", + "smooth", + "interpolate", + "interpolate_smooth", + "wander", + "dolly_zoom", + "extend", + "wobble_inter", + ] = "extend", + color_mode: Literal["RGB+D", "RGB+ED"] = "RGB+ED", + vis_depth: Optional[Literal["hcat", "vcat"]] = "hcat", + enable_tqdm: Optional[bool] = True, + output_name: Optional[str] = None, + video_quality: Literal["low", "medium", "high"] = "high", +) -> None: + gs_world = prediction.gaussians + # if target poses are not provided, render the (smooth/interpolate) input poses + if extrinsics is not None: + tgt_extrs = extrinsics + else: + tgt_extrs = torch.from_numpy(prediction.extrinsics).unsqueeze(0).to(gs_world.means) + if prediction.is_metric: + scale_factor = prediction.scale_factor + if scale_factor is not None: + tgt_extrs[:, :, :3, 3] /= scale_factor + tgt_intrs = ( + intrinsics + if intrinsics is not None + else torch.from_numpy(prediction.intrinsics).unsqueeze(0).to(gs_world.means) + ) + # if render resolution is not provided, render the input ones + if out_image_hw is not None: + H, W = out_image_hw + else: + H, W = prediction.depth.shape[-2:] + # if single views, render wander trj + if tgt_extrs.shape[1] <= 1: + trj_mode = "wander" + # trj_mode = "dolly_zoom" + + color, depth = run_renderer_in_chunk_w_trj_mode( + gaussians=gs_world, + extrinsics=tgt_extrs, + intrinsics=tgt_intrs, + image_shape=(H, W), + chunk_size=chunk_size, + trj_mode=trj_mode, + use_sh=True, + color_mode=color_mode, + enable_tqdm=enable_tqdm, + ) + + # save as video + ffmpeg_params = [ + "-crf", + VIDEO_QUALITY_MAP[video_quality]["crf"], + "-preset", + VIDEO_QUALITY_MAP[video_quality]["preset"], + "-pix_fmt", + "yuv420p", + ] # best compatibility + + os.makedirs(os.path.join(export_dir, "gs_video"), exist_ok=True) + for idx in range(color.shape[0]): + video_i = color[idx] + if vis_depth is not None: + depth_i = vis_depth_map_tensor(depth[0]) + cat_fn = hcat if vis_depth == "hcat" else vcat + video_i = torch.stack([cat_fn(c, d) for c, d in zip(video_i, depth_i)]) + frames = list( + (video_i.clamp(0, 1) * 255).byte().permute(0, 2, 3, 1).cpu().numpy() + ) # T x H x W x C, uint8, numpy() + + fps = 24 + clip = mpy.ImageSequenceClip(frames, fps=fps) + output_name = f"{idx:04d}_{trj_mode}" if output_name is None else output_name + save_path = os.path.join(export_dir, f"gs_video/{output_name}.mp4") + # clip.write_videofile(save_path, codec="libx264", audio=False, bitrate="4000k") + clip.write_videofile( + save_path, + codec="libx264", + audio=False, + fps=fps, + ffmpeg_params=ffmpeg_params, + ) + return diff --git a/depth_anything_3/utils/export/npz.py b/depth_anything_3/utils/export/npz.py new file mode 100644 index 0000000000000000000000000000000000000000..75df4f5194066c76a177f679e54721cb39739dd7 --- /dev/null +++ b/depth_anything_3/utils/export/npz.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import numpy as np + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.parallel_utils import async_call + + +@async_call +def export_to_npz( + prediction: Prediction, + export_dir: str, +): + output_file = os.path.join(export_dir, "exports", "npz", "results.npz") + os.makedirs(os.path.dirname(output_file), exist_ok=True) + + # Use prediction.processed_images, which is already processed image data + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + image = prediction.processed_images # (N,H,W,3) uint8 + + # Build save dict with only non-None values + save_dict = { + "image": image, + "depth": np.round(prediction.depth, 6), + } + + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + # aux = {k: np.round(v, 4) for k, v in prediction.aux.items()} + np.savez_compressed(output_file, **save_dict) + + +@async_call +def export_to_mini_npz( + prediction: Prediction, + export_dir: str, +): + output_file = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + os.makedirs(os.path.dirname(output_file), exist_ok=True) + + # Build save dict with only non-None values + save_dict = { + "depth": np.round(prediction.depth, 8), + } + + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + np.savez_compressed(output_file, **save_dict) diff --git a/depth_anything_3/utils/export/utils.py b/depth_anything_3/utils/export/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..81f45fb563ce595bf547bebe829c9b83eb175f1c --- /dev/null +++ b/depth_anything_3/utils/export/utils.py @@ -0,0 +1,30 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + + +def _denorm_and_to_uint8(image_tensor: torch.Tensor) -> np.ndarray: + """Denormalize to [0,255] and output (N, H, W, 3) uint8.""" + resnet_mean = torch.tensor( + [0.485, 0.456, 0.406], dtype=image_tensor.dtype, device=image_tensor.device + ) + resnet_std = torch.tensor( + [0.229, 0.224, 0.225], dtype=image_tensor.dtype, device=image_tensor.device + ) + img = image_tensor * resnet_std[None, :, None, None] + resnet_mean[None, :, None, None] + img = torch.clamp(img, 0.0, 1.0) + img = (img.permute(0, 2, 3, 1).cpu().numpy() * 255.0).round().astype(np.uint8) # (N,H,W,3) + return img diff --git a/depth_anything_3/utils/geometry.py b/depth_anything_3/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..41a42190d637ef25db3599cc90f9a4f40be3b8d7 --- /dev/null +++ b/depth_anything_3/utils/geometry.py @@ -0,0 +1,498 @@ +# flake8: noqa: F722 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from types import SimpleNamespace +from typing import Optional +import numpy as np +import torch +import torch.nn.functional as F +from einops import einsum + + +def as_homogeneous(ext): + """ + Accept (..., 3,4) or (..., 4,4) extrinsics, return (...,4,4) homogeneous matrix. + Supports torch.Tensor or np.ndarray. + """ + if isinstance(ext, torch.Tensor): + # If already in homogeneous form + if ext.shape[-2:] == (4, 4): + return ext + elif ext.shape[-2:] == (3, 4): + # Create a new homogeneous matrix + ones = torch.zeros_like(ext[..., :1, :4]) + ones[..., 0, 3] = 1.0 + return torch.cat([ext, ones], dim=-2) + else: + raise ValueError(f"Invalid shape for torch.Tensor: {ext.shape}") + + elif isinstance(ext, np.ndarray): + if ext.shape[-2:] == (4, 4): + return ext + elif ext.shape[-2:] == (3, 4): + ones = np.zeros_like(ext[..., :1, :4]) + ones[..., 0, 3] = 1.0 + return np.concatenate([ext, ones], axis=-2) + else: + raise ValueError(f"Invalid shape for np.ndarray: {ext.shape}") + + else: + raise TypeError("Input must be a torch.Tensor or np.ndarray.") + + +@torch.jit.script +def affine_inverse(A: torch.Tensor): + R = A[..., :3, :3] # ..., 3, 3 + T = A[..., :3, 3:] # ..., 3, 1 + P = A[..., 3:, :] # ..., 1, 4 + return torch.cat([torch.cat([R.mT, -R.mT @ T], dim=-1), P], dim=-2) + + +def transpose_last_two_axes(arr): + """ + for np < 2 + """ + if arr.ndim < 2: + return arr + axes = list(range(arr.ndim)) + # swap the last two + axes[-2], axes[-1] = axes[-1], axes[-2] + return arr.transpose(axes) + + +def affine_inverse_np(A: np.ndarray): + R = A[..., :3, :3] + T = A[..., :3, 3:] + P = A[..., 3:, :] + return np.concatenate( + [ + np.concatenate([transpose_last_two_axes(R), -transpose_last_two_axes(R) @ T], axis=-1), + P, + ], + axis=-2, + ) + + +def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor: + """ + Quaternion Order: XYZW or say ijkr, scalar-last + + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + i, j, k, r = torch.unbind(quaternions, -1) + # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape( + batch_dim + (4,) + ) + + # Convert from rijk to ijkr + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def sample_image_grid( + shape: tuple[int, ...], + device: torch.device = torch.device("cpu"), +) -> tuple[ + torch.Tensor, # float coordinates (xy indexing), "*shape dim" + torch.Tensor, # integer indices (ij indexing), "*shape dim" +]: + """Get normalized (range 0 to 1) coordinates and integer indices for an image.""" + + # Each entry is a pixel-wise integer coordinate. In the 2D case, each entry is a + # (row, col) coordinate. + indices = [torch.arange(length, device=device) for length in shape] + stacked_indices = torch.stack(torch.meshgrid(*indices, indexing="ij"), dim=-1) + + # Each entry is a floating-point coordinate in the range (0, 1). In the 2D case, + # each entry is an (x, y) coordinate. + coordinates = [(idx + 0.5) / length for idx, length in zip(indices, shape)] + coordinates = reversed(coordinates) + coordinates = torch.stack(torch.meshgrid(*coordinates, indexing="xy"), dim=-1) + + return coordinates, stacked_indices + + +def homogenize_points(points: torch.Tensor) -> torch.Tensor: # "*batch dim" # "*batch dim+1" + """Convert batched points (xyz) to (xyz1).""" + return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + +def homogenize_vectors(vectors: torch.Tensor) -> torch.Tensor: # "*batch dim" # "*batch dim+1" + """Convert batched vectors (xyz) to (xyz0).""" + return torch.cat([vectors, torch.zeros_like(vectors[..., :1])], dim=-1) + + +def transform_rigid( + homogeneous_coordinates: torch.Tensor, # "*#batch dim" + transformation: torch.Tensor, # "*#batch dim dim" +) -> torch.Tensor: # "*batch dim" + """Apply a rigid-body transformation to points or vectors.""" + return einsum( + transformation, + homogeneous_coordinates.to(transformation.dtype), + "... i j, ... j -> ... i", + ) + + +def transform_cam2world( + homogeneous_coordinates: torch.Tensor, # "*#batch dim" + extrinsics: torch.Tensor, # "*#batch dim dim" +) -> torch.Tensor: # "*batch dim" + """Transform points from 3D camera coordinates to 3D world coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics) + + +def unproject( + coordinates: torch.Tensor, # "*#batch dim" + z: torch.Tensor, # "*#batch" + intrinsics: torch.Tensor, # "*#batch dim+1 dim+1" +) -> torch.Tensor: # "*batch dim+1" + """Unproject 2D camera coordinates with the given Z values.""" + + # Apply the inverse intrinsics to the coordinates. + coordinates = homogenize_points(coordinates) + ray_directions = einsum( + intrinsics.float().inverse().to(intrinsics), + coordinates.to(intrinsics.dtype), + "... i j, ... j -> ... i", + ) + + # Apply the supplied depth values. + return ray_directions * z[..., None] + + +def get_world_rays( + coordinates: torch.Tensor, # "*#batch dim" + extrinsics: torch.Tensor, # "*#batch dim+2 dim+2" + intrinsics: torch.Tensor, # "*#batch dim+1 dim+1" +) -> tuple[ + torch.Tensor, # origins, "*batch dim+1" + torch.Tensor, # directions, "*batch dim+1" +]: + # Get camera-space ray directions. + directions = unproject( + coordinates, + torch.ones_like(coordinates[..., 0]), + intrinsics, + ) + directions = directions / directions.norm(dim=-1, keepdim=True) + + # Transform ray directions to world coordinates. + directions = homogenize_vectors(directions) + directions = transform_cam2world(directions, extrinsics)[..., :-1] + + # Tile the ray origins to have the same shape as the ray directions. + origins = extrinsics[..., :-1, -1].broadcast_to(directions.shape) + + return origins, directions + + +def get_fov(intrinsics: torch.Tensor) -> torch.Tensor: # "batch 3 3" -> "batch 2" + intrinsics_inv = intrinsics.float().inverse().to(intrinsics) + + def process_vector(vector): + vector = torch.tensor(vector, dtype=intrinsics.dtype, device=intrinsics.device) + vector = einsum(intrinsics_inv, vector, "b i j, j -> b i") + return vector / vector.norm(dim=-1, keepdim=True) + + left = process_vector([0, 0.5, 1]) + right = process_vector([1, 0.5, 1]) + top = process_vector([0.5, 0, 1]) + bottom = process_vector([0.5, 1, 1]) + fov_x = (left * right).sum(dim=-1).acos() + fov_y = (top * bottom).sum(dim=-1).acos() + return torch.stack((fov_x, fov_y), dim=-1) + + +def map_pdf_to_opacity( + pdf: torch.Tensor, # " *batch" + global_step: int = 0, + opacity_mapping: Optional[dict] = None, +) -> torch.Tensor: # " *batch" + # https://www.desmos.com/calculator/opvwti3ba9 + + # Figure out the exponent. + if opacity_mapping is not None: + cfg = SimpleNamespace(**opacity_mapping) + x = cfg.initial + min(global_step / cfg.warm_up, 1) * (cfg.final - cfg.initial) + else: + x = 0.0 + exponent = 2**x + + # Map the probability density to an opacity. + return 0.5 * (1 - (1 - pdf) ** exponent + pdf ** (1 / exponent)) + +def normalize_homogenous_points(points): + """Normalize the point vectors""" + return points / points[..., -1:] + +def inverse_intrinsic_matrix(ixts): + """ """ + return torch.inverse(ixts) + +def pixel_space_to_camera_space(pixel_space_points, depth, intrinsics): + """ + Convert pixel space points to camera space points. + + Args: + pixel_space_points (torch.Tensor): Pixel space points with shape (h, w, 2) + depth (torch.Tensor): Depth map with shape (b, v, h, w, 1) + intrinsics (torch.Tensor): Camera intrinsics with shape (b, v, 3, 3) + + Returns: + torch.Tensor: Camera space points with shape (b, v, h, w, 3). + """ + pixel_space_points = homogenize_points(pixel_space_points) + # camera_space_points = torch.einsum( + # "b v i j , h w j -> b v h w i", intrinsics.inverse(), pixel_space_points + # ) + camera_space_points = torch.einsum( + "b v i j , h w j -> b v h w i", inverse_intrinsic_matrix(intrinsics), pixel_space_points + ) + camera_space_points = camera_space_points * depth + return camera_space_points + + +def camera_space_to_world_space(camera_space_points, c2w): + """ + Convert camera space points to world space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + camera_space_points = homogenize_points(camera_space_points) + world_space_points = torch.einsum("b v i j , b v h w j -> b v h w i", c2w, camera_space_points) + return world_space_points[..., :3] + + +def camera_space_to_pixel_space(camera_space_points, intrinsics): + """ + Convert camera space points to pixel space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v1, v2, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 3, 3) + + Returns: + torch.Tensor: World space points with shape (b, v1, v2, h, w, 2). + """ + camera_space_points = normalize_homogenous_points(camera_space_points) + pixel_space_points = torch.einsum( + "b u i j , b v u h w j -> b v u h w i", intrinsics, camera_space_points + ) + return pixel_space_points[..., :2] + + +def world_space_to_camera_space(world_space_points, c2w): + """ + Convert world space points to pixel space points. + + Args: + world_space_points (torch.Tensor): World space points with shape (b, v1, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 4, 4) + + Returns: + torch.Tensor: Camera space points with shape (b, v1, v2, h, w, 3). + """ + world_space_points = homogenize_points(world_space_points) + camera_space_points = torch.einsum( + "b u i j , b v h w j -> b v u h w i", c2w.inverse(), world_space_points + ) + return camera_space_points[..., :3] + + +def unproject_depth( + depth, intrinsics, c2w=None, ixt_normalized=False, num_patches_x=None, num_patches_y=None +): + """ + Turn the depth map into a 3D point cloud in world space + + Args: + depth: (b, v, h, w, 1) + intrinsics: (b, v, 3, 3) + c2w: (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + if c2w is None: + c2w = torch.eye(4, device=depth.device, dtype=depth.dtype) + c2w = c2w[None, None].repeat(depth.shape[0], depth.shape[1], 1, 1) + + if not ixt_normalized: + # Compute indices of pixels + h, w = depth.shape[-3], depth.shape[-2] + x_grid, y_grid = torch.meshgrid( + torch.arange(w, device=depth.device, dtype=depth.dtype), + torch.arange(h, device=depth.device, dtype=depth.dtype), + indexing="xy", + ) # (h, w), (h, w) + else: + # ixt_normalized: h=w=2.0. cx, cy, fx, fy are normalized according to h=w=2.0 + assert num_patches_x is not None and num_patches_y is not None + dx = 1 / num_patches_x + dy = 1 / num_patches_y + max_y = 1 - dy + min_y = -max_y + max_x = 1 - dx + min_x = -max_x + + grid_shift = 1.0 + y_grid, x_grid = torch.meshgrid( + torch.linspace( + min_y + grid_shift, + max_y + grid_shift, + num_patches_y, + dtype=torch.float32, + device=depth.device, + ), + torch.linspace( + min_x + grid_shift, + max_x + grid_shift, + num_patches_x, + dtype=torch.float32, + device=depth.device, + ), + indexing="ij", + ) + + # Compute coordinates of pixels in camera space + pixel_space_points = torch.stack((x_grid, y_grid), dim=-1) # (..., h, w, 2) + camera_points = pixel_space_to_camera_space( + pixel_space_points, depth, intrinsics + ) # (..., h, w, 3) + + # Convert points to world space + world_points = camera_space_to_world_space(camera_points, c2w) # (..., h, w, 3) + + return world_points \ No newline at end of file diff --git a/depth_anything_3/utils/gsply_helpers.py b/depth_anything_3/utils/gsply_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..5733009e4e91ad80ab59179c80e2df8a0430ab5f --- /dev/null +++ b/depth_anything_3/utils/gsply_helpers.py @@ -0,0 +1,173 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from pathlib import Path +from typing import Optional +import numpy as np +import torch +from einops import rearrange, repeat +from plyfile import PlyData, PlyElement +from torch import Tensor + +from depth_anything_3.specs import Gaussians + + +def construct_list_of_attributes(num_rest: int) -> list[str]: + attributes = ["x", "y", "z", "nx", "ny", "nz"] + for i in range(3): + attributes.append(f"f_dc_{i}") + for i in range(num_rest): + attributes.append(f"f_rest_{i}") + attributes.append("opacity") + for i in range(3): + attributes.append(f"scale_{i}") + for i in range(4): + attributes.append(f"rot_{i}") + return attributes + + +def export_ply( + means: Tensor, # "gaussian 3" + scales: Tensor, # "gaussian 3" + rotations: Tensor, # "gaussian 4" + harmonics: Tensor, # "gaussian 3 d_sh" + opacities: Tensor, # "gaussian" + path: Path, + shift_and_scale: bool = False, + save_sh_dc_only: bool = True, + match_3dgs_mcmc_dev: Optional[bool] = False, +): + if shift_and_scale: + # Shift the scene so that the median Gaussian is at the origin. + means = means - means.median(dim=0).values + + # Rescale the scene so that most Gaussians are within range [-1, 1]. + scale_factor = means.abs().quantile(0.95, dim=0).max() + means = means / scale_factor + scales = scales / scale_factor + + rotations = rotations.detach().cpu().numpy() + + # Since current model use SH_degree = 4, + # which require large memory to store, we can only save the DC band to save memory. + f_dc = harmonics[..., 0] + f_rest = harmonics[..., 1:].flatten(start_dim=1) + + if match_3dgs_mcmc_dev: + sh_degree = 3 + n_rest = 3 * (sh_degree + 1) ** 2 - 3 + f_rest = repeat( + torch.zeros_like(harmonics[..., :1]), "... i -> ... (n i)", n=(n_rest // 3) + ).flatten(start_dim=1) + dtype_full = [ + (attribute, "f4") + for attribute in construct_list_of_attributes(num_rest=n_rest) + if attribute not in ("nx", "ny", "nz") + ] + else: + dtype_full = [ + (attribute, "f4") + for attribute in construct_list_of_attributes( + 0 if save_sh_dc_only else f_rest.shape[1] + ) + ] + elements = np.empty(means.shape[0], dtype=dtype_full) + attributes = [ + means.detach().cpu().numpy(), + torch.zeros_like(means).detach().cpu().numpy(), + f_dc.detach().cpu().contiguous().numpy(), + f_rest.detach().cpu().contiguous().numpy(), + opacities[..., None].detach().cpu().numpy(), + scales.log().detach().cpu().numpy(), + rotations, + ] + if match_3dgs_mcmc_dev: + attributes.pop(1) # dummy normal is not needed + elif save_sh_dc_only: + attributes.pop(3) # remove f_rest from attributes + + attributes = np.concatenate(attributes, axis=1) + elements[:] = list(map(tuple, attributes)) + path.parent.mkdir(exist_ok=True, parents=True) + PlyData([PlyElement.describe(elements, "vertex")]).write(path) + + +def inverse_sigmoid(x): + return torch.log(x / (1 - x)) + + +def save_gaussian_ply( + gaussians: Gaussians, + save_path: str, + ctx_depth: torch.Tensor, # depth of input views; for getting shape and filtering, "v h w 1" + shift_and_scale: bool = False, + save_sh_dc_only: bool = True, + gs_views_interval: int = 1, + inv_opacity: Optional[bool] = True, + prune_by_depth_percent: Optional[float] = 1.0, + prune_border_gs: Optional[bool] = True, + match_3dgs_mcmc_dev: Optional[bool] = False, +): + b = gaussians.means.shape[0] + assert b == 1, "must set batch_size=1 when exporting 3D gaussians" + src_v, out_h, out_w, _ = ctx_depth.shape + + # extract gs params + world_means = gaussians.means + world_shs = gaussians.harmonics + world_rotations = gaussians.rotations + gs_scales = gaussians.scales + gs_opacities = inverse_sigmoid(gaussians.opacities) if inv_opacity else gaussians.opacities + + # Create a mask to filter the Gaussians. + + # TODO: prune the sky region here + + # throw away Gaussians at the borders, since they're generally of lower quality. + if prune_border_gs: + mask = torch.zeros_like(ctx_depth, dtype=torch.bool) + gstrim_h = int(8 / 256 * out_h) + gstrim_w = int(8 / 256 * out_w) + mask[:, gstrim_h:-gstrim_h, gstrim_w:-gstrim_w, :] = 1 + else: + mask = torch.ones_like(ctx_depth, dtype=torch.bool) + + # trim the far away point based on depth; + if prune_by_depth_percent is not None and prune_by_depth_percent < 1: + in_depths = ctx_depth + d_percentile = torch.quantile( + in_depths.view(in_depths.shape[0], -1), q=prune_by_depth_percent, dim=1 + ).view(-1, 1, 1) + d_mask = (in_depths[..., 0] <= d_percentile).unsqueeze(-1) + mask = mask & d_mask + mask = mask.squeeze(-1) # v h w + + # helper fn, must place after mask + def trim_select_reshape(element): + selected_element = rearrange( + element[0], "(v h w) ... -> v h w ...", v=src_v, h=out_h, w=out_w + ) + selected_element = selected_element[::gs_views_interval][mask[::gs_views_interval]] + return selected_element + + export_ply( + means=trim_select_reshape(world_means), + scales=trim_select_reshape(gs_scales), + rotations=trim_select_reshape(world_rotations), + harmonics=trim_select_reshape(world_shs), + opacities=trim_select_reshape(gs_opacities), + path=Path(save_path), + shift_and_scale=shift_and_scale, + save_sh_dc_only=save_sh_dc_only, + match_3dgs_mcmc_dev=match_3dgs_mcmc_dev, + ) diff --git a/depth_anything_3/utils/io/input_processor.py b/depth_anything_3/utils/io/input_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..fa60194123c22c84a0cdeb577cbf10e2704278bb --- /dev/null +++ b/depth_anything_3/utils/io/input_processor.py @@ -0,0 +1,501 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Input processor for Depth Anything 3 (parallelized). + +This version removes the square center-crop step for "*crop" methods (same as your note). +In addition, it parallelizes per-image preprocessing using the provided `parallel_execution`. +""" + +from __future__ import annotations + +from typing import Sequence +import cv2 +import numpy as np +import torch +import torchvision.transforms as T +from PIL import Image + +from depth_anything_3.utils.logger import logger +from depth_anything_3.utils.parallel_utils import parallel_execution + + +class InputProcessor: + """Prepares a batch of images for model inference. + This processor converts a list of image file paths into a single, model-ready + tensor. The processing pipeline is executed in parallel across multiple workers + for efficiency. + + Pipeline: + 1) Load image and convert to RGB + 2) Boundary resize (upper/lower bound, preserving aspect ratio) + 3) Enforce divisibility by PATCH_SIZE: + - "*resize" methods: each dimension is rounded to nearest multiple + (may up/downscale a few px) + - "*crop" methods: each dimension is floored to nearest multiple via center crop + 4) Convert to tensor and apply ImageNet normalization + 5) Stack into (1, N, 3, H, W) + + Parallelization: + - Each image is processed independently in a worker. + - Order of outputs matches the input order. + """ + + NORMALIZE = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + PATCH_SIZE = 14 + + def __init__(self): + pass + + # ----------------------------- + # Public API + # ----------------------------- + def __call__( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + *, + num_workers: int = 8, + print_progress: bool = False, + sequential: bool | None = None, + desc: str | None = "Preprocess", + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Returns: + (tensor, extrinsics_list, intrinsics_list) + tensor shape: (1, N, 3, H, W) + """ + sequential = self._resolve_sequential(sequential, num_workers) + exts_list, ixts_list = self._validate_and_pack_meta(image, extrinsics, intrinsics) + + results = self._run_parallel( + image=image, + exts_list=exts_list, + ixts_list=ixts_list, + process_res=process_res, + process_res_method=process_res_method, + num_workers=num_workers, + print_progress=print_progress, + sequential=sequential, + desc=desc, + ) + + proc_imgs, out_sizes, out_ixts, out_exts = self._unpack_results(results) + proc_imgs, out_sizes, out_ixts = self._unify_batch_shapes(proc_imgs, out_sizes, out_ixts) + + batch_tensor = self._stack_batch(proc_imgs) + out_exts = ( + torch.from_numpy(np.asarray(out_exts)).float() + if out_exts is not None and out_exts[0] is not None + else None + ) + out_ixts = ( + torch.from_numpy(np.asarray(out_ixts)).float() + if out_ixts is not None and out_ixts[0] is not None + else None + ) + return (batch_tensor, out_exts, out_ixts) + + # ----------------------------- + # __call__ helpers + # ----------------------------- + def _resolve_sequential(self, sequential: bool | None, num_workers: int) -> bool: + return (num_workers <= 1) if sequential is None else sequential + + def _validate_and_pack_meta( + self, + images: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None, + intrinsics: np.ndarray | None, + ) -> tuple[list[np.ndarray | None] | None, list[np.ndarray | None] | None]: + if extrinsics is not None and len(extrinsics) != len(images): + raise ValueError("Length of extrinsics must match images when provided.") + if intrinsics is not None and len(intrinsics) != len(images): + raise ValueError("Length of intrinsics must match images when provided.") + exts_list = [e for e in extrinsics] if extrinsics is not None else None + ixts_list = [k for k in intrinsics] if intrinsics is not None else None + return exts_list, ixts_list + + def _run_parallel( + self, + *, + image: list[np.ndarray | Image.Image | str], + exts_list: list[np.ndarray | None] | None, + ixts_list: list[np.ndarray | None] | None, + process_res: int, + process_res_method: str, + num_workers: int, + print_progress: bool, + sequential: bool, + desc: str | None, + ): + results = parallel_execution( + image, + exts_list, + ixts_list, + action=self._process_one, # (img, extrinsic, intrinsic, ...) + num_processes=num_workers, + print_progress=print_progress, + sequential=sequential, + desc=desc, + process_res=process_res, + process_res_method=process_res_method, + ) + if not results: + raise RuntimeError( + "No preprocessing results returned. Check inputs and parallel_execution." + ) + return results + + def _unpack_results(self, results): + """ + results: List[Tuple[torch.Tensor, Tuple[H, W], Optional[np.ndarray], Optional[np.ndarray]]] + -> processed_images, out_sizes, out_intrinsics, out_extrinsics + """ + try: + processed_images, out_sizes, out_intrinsics, out_extrinsics = zip(*results) + except Exception as e: + raise RuntimeError( + "Unexpected results structure from parallel_execution: " + f"{type(results)} / sample: {results[0]}" + ) from e + + return list(processed_images), list(out_sizes), list(out_intrinsics), list(out_extrinsics) + + def _unify_batch_shapes( + self, + processed_images: list[torch.Tensor], + out_sizes: list[tuple[int, int]], + out_intrinsics: list[np.ndarray | None], + ) -> tuple[list[torch.Tensor], list[tuple[int, int]], list[np.ndarray | None]]: + """Center-crop all tensors to the smallest H, W; adjust intrinsics' cx, cy accordingly.""" + if len(set(out_sizes)) <= 1: + return processed_images, out_sizes, out_intrinsics + + min_h = min(h for h, _ in out_sizes) + min_w = min(w for _, w in out_sizes) + logger.warn( + f"Images in batch have different sizes {out_sizes}; " + f"center-cropping all to smallest ({min_h},{min_w})" + ) + + center_crop = T.CenterCrop((min_h, min_w)) + new_imgs, new_sizes, new_ixts = [], [], [] + for img_t, (H, W), K in zip(processed_images, out_sizes, out_intrinsics): + crop_top = max(0, (H - min_h) // 2) + crop_left = max(0, (W - min_w) // 2) + new_imgs.append(center_crop(img_t)) + new_sizes.append((min_h, min_w)) + if K is None: + new_ixts.append(None) + else: + K_adj = K.copy() + K_adj[0, 2] -= crop_left + K_adj[1, 2] -= crop_top + new_ixts.append(K_adj) + return new_imgs, new_sizes, new_ixts + + def _stack_batch(self, processed_images: list[torch.Tensor]) -> torch.Tensor: + return torch.stack(processed_images) + + # ----------------------------- + # Per-item worker + # ----------------------------- + def _process_one( + self, + img: np.ndarray | Image.Image | str, + extrinsic: np.ndarray | None = None, + intrinsic: np.ndarray | None = None, + *, + process_res: int, + process_res_method: str, + ) -> tuple[torch.Tensor, tuple[int, int], np.ndarray | None, np.ndarray | None]: + # Load & remember original size + pil_img = self._load_image(img) + orig_w, orig_h = pil_img.size + + # Boundary resize + pil_img = self._resize_image(pil_img, process_res, process_res_method) + w, h = pil_img.size + intrinsic = self._resize_ixt(intrinsic, orig_w, orig_h, w, h) + + # Enforce divisibility by PATCH_SIZE + if process_res_method.endswith("resize"): + pil_img = self._make_divisible_by_resize(pil_img, self.PATCH_SIZE) + new_w, new_h = pil_img.size + intrinsic = self._resize_ixt(intrinsic, w, h, new_w, new_h) + w, h = new_w, new_h + elif process_res_method.endswith("crop"): + pil_img = self._make_divisible_by_crop(pil_img, self.PATCH_SIZE) + new_w, new_h = pil_img.size + intrinsic = self._crop_ixt(intrinsic, w, h, new_w, new_h) + w, h = new_w, new_h + else: + raise ValueError(f"Unsupported process_res_method: {process_res_method}") + + # Convert to tensor & normalize + img_tensor = self._normalize_image(pil_img) + _, H, W = img_tensor.shape + assert (W, H) == (w, h), "Tensor size mismatch with PIL image size after processing." + + # Return: (img_tensor, (H, W), intrinsic, extrinsic) + return img_tensor, (H, W), intrinsic, extrinsic + + # ----------------------------- + # Intrinsics transforms + # ----------------------------- + def _resize_ixt( + self, + intrinsic: np.ndarray | None, + orig_w: int, + orig_h: int, + w: int, + h: int, + ) -> np.ndarray | None: + if intrinsic is None: + return None + K = intrinsic.copy() + # scale fx, cx by w ratio; fy, cy by h ratio + K[:1] *= w / float(orig_w) + K[1:2] *= h / float(orig_h) + return K + + def _crop_ixt( + self, + intrinsic: np.ndarray | None, + orig_w: int, + orig_h: int, + w: int, + h: int, + ) -> np.ndarray | None: + if intrinsic is None: + return None + K = intrinsic.copy() + crop_h = (orig_h - h) // 2 + crop_w = (orig_w - w) // 2 + K[0, 2] -= crop_w + K[1, 2] -= crop_h + return K + + # ----------------------------- + # I/O & normalization + # ----------------------------- + def _load_image(self, img: np.ndarray | Image.Image | str) -> Image.Image: + if isinstance(img, str): + return Image.open(img).convert("RGB") + elif isinstance(img, np.ndarray): + # Assume HxWxC uint8/RGB + return Image.fromarray(img).convert("RGB") + elif isinstance(img, Image.Image): + return img.convert("RGB") + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + def _normalize_image(self, img: Image.Image) -> torch.Tensor: + img_tensor = T.ToTensor()(img) + return self.NORMALIZE(img_tensor) + + # ----------------------------- + # Boundary resizing + # ----------------------------- + def _resize_image(self, img: Image.Image, target_size: int, method: str) -> Image.Image: + if method in ("upper_bound_resize", "upper_bound_crop"): + return self._resize_longest_side(img, target_size) + elif method in ("lower_bound_resize", "lower_bound_crop"): + return self._resize_shortest_side(img, target_size) + else: + raise ValueError(f"Unsupported resize method: {method}") + + def _resize_longest_side(self, img: Image.Image, target_size: int) -> Image.Image: + w, h = img.size + longest = max(w, h) + if longest == target_size: + return img + scale = target_size / float(longest) + new_w = max(1, int(round(w * scale))) + new_h = max(1, int(round(h * scale))) + interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + def _resize_shortest_side(self, img: Image.Image, target_size: int) -> Image.Image: + w, h = img.size + shortest = min(w, h) + if shortest == target_size: + return img + scale = target_size / float(shortest) + new_w = max(1, int(round(w * scale))) + new_h = max(1, int(round(h * scale))) + interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + # ----------------------------- + # Make divisible by PATCH_SIZE + # ----------------------------- + def _make_divisible_by_crop(self, img: Image.Image, patch: int) -> Image.Image: + """ + Floor each dimension to the nearest multiple of PATCH_SIZE via center crop. + Example: 504x377 -> 504x364 + """ + w, h = img.size + new_w = (w // patch) * patch + new_h = (h // patch) * patch + if new_w == w and new_h == h: + return img + left = (w - new_w) // 2 + top = (h - new_h) // 2 + return img.crop((left, top, left + new_w, top + new_h)) + + def _make_divisible_by_resize(self, img: Image.Image, patch: int) -> Image.Image: + """ + Round each dimension to nearest multiple of PATCH_SIZE via small resize. + """ + w, h = img.size + + def nearest_multiple(x: int, p: int) -> int: + down = (x // p) * p + up = down + p + return up if abs(up - x) <= abs(x - down) else down + + new_w = max(1, nearest_multiple(w, patch)) + new_h = max(1, nearest_multiple(h, patch)) + if new_w == w and new_h == h: + return img + upscale = (new_w > w) or (new_h > h) + interpolation = cv2.INTER_CUBIC if upscale else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + +# Backward compatibility alias +InputAdapter = InputProcessor + + +# =========================== +# Minimal test runner (parallel execution) +# =========================== +if __name__ == "__main__": + """ + Minimal test suite: + - Creates pairs of images so batch shapes match. + - Tests all four process_res_methods. + - Prints fx fy cx cy IN->OUT per image. + - Includes cases with K/E provided and with None. + """ + + def fmt_k_line(K: np.ndarray | None) -> str: + if K is None: + return "None" + fx, fy, cx, cy = float(K[0, 0]), float(K[1, 1]), float(K[0, 2]), float(K[1, 2]) + return f"fx={fx:.3f} fy={fy:.3f} cx={cx:.3f} cy={cy:.3f}" + + def show_result( + tag: str, + tensor: torch.Tensor, + Ks_in: Sequence[np.ndarray | None] | None = None, + Ks_out: Sequence[np.ndarray | None] | None = None, + ): + B, N, C, H, W = tensor.shape + print(f"[{tag}] shape={tuple(tensor.shape)} HxW=({H},{W}) div14=({H%14==0},{W%14==0})") + assert H % 14 == 0 and W % 14 == 0, f"{tag}: output size not divisible by 14!" + if Ks_in is not None or Ks_out is not None: + Ks_in = Ks_in or [None] * N + Ks_out = Ks_out or [None] * N + for i in range(N): + print(f" K[{i}]: {fmt_k_line(Ks_in[i])} -> {fmt_k_line(Ks_out[i])}") + + proc = InputProcessor() + process_res = 504 + methods = ["upper_bound_resize", "upper_bound_crop", "lower_bound_resize", "lower_bound_crop"] + + # Example sizes (two orientations) + small_sizes = [(680, 1208), (1208, 680)] + large_sizes = [(1208, 680), (680, 1208)] + + def make_K(w, h, fx=1200.0, fy=1100.0): + cx, cy = w / 2.0, h / 2.0 + K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32) + return K + + def run_suite(suite_name: str, sizes: list[tuple[int, int]]): + print(f"\n===== {suite_name} =====") + for w, h in sizes: + img = Image.new("RGB", (w, h), color=(123, 222, 100)) + batch_imgs = [img, img] + + # intrinsics / extrinsics examples + Ks_in = [make_K(w, h), make_K(w, h)] + Es_in = [np.eye(4, dtype=np.float32), np.eye(4, dtype=np.float32)] + + for m in methods: + tensor, Es_out, Ks_out = proc( + image=batch_imgs, + process_res=process_res, + process_res_method=m, + num_workers=8, + print_progress=False, + intrinsics=Ks_in, # test with non-None + extrinsics=Es_in, + ) + show_result(f"{suite_name} size=({w},{h}) | {m}", tensor, Ks_in, Ks_out) + + # Also test None path + tensor2, Es_out2, Ks_out2 = proc( + image=batch_imgs, + process_res=process_res, + process_res_method="upper_bound_resize", + num_workers=8, + intrinsics=None, + extrinsics=None, + ) + show_result( + f"{suite_name} size=({w},{h}) | upper_bound_resize | no K/E", + tensor2, + None, + Ks_out2, + ) + + run_suite("SMALL", small_sizes) + run_suite("LARGE", large_sizes) + + # Extra sanity for 504x376 + print("\n===== EXTRA sanity for 504x376 =====") + img_example = Image.new("RGB", (504, 376), color=(10, 20, 30)) + Ks_in_extra = [make_K(504, 376, fx=900.0, fy=900.0), make_K(504, 376, fx=900.0, fy=900.0)] + + out_r, _, Ks_out_r = proc( + image=[img_example, img_example], + process_res=504, + process_res_method="upper_bound_resize", + num_workers=8, + intrinsics=Ks_in_extra, + ) + out_c, _, Ks_out_c = proc( + image=[img_example, img_example], + process_res=504, + process_res_method="upper_bound_crop", + num_workers=8, + intrinsics=Ks_in_extra, + ) + _, _, _, Hr, Wr = out_r.shape + _, _, _, Hc, Wc = out_c.shape + print(f"upper_bound_resize -> ({Hr},{Wr}) (rounded to nearest multiple of 14)") + show_result("Ks after upper_bound_resize", out_r, Ks_in_extra, Ks_out_r) + print(f"upper_bound_crop -> ({Hc},{Wc}) (floored to multiple of 14)") + show_result("Ks after upper_bound_crop", out_c, Ks_in_extra, Ks_out_c) diff --git a/depth_anything_3/utils/io/output_processor.py b/depth_anything_3/utils/io/output_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..c317eb9d596c1687b5281891a035993868cc5f8c --- /dev/null +++ b/depth_anything_3/utils/io/output_processor.py @@ -0,0 +1,172 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Output processor for Depth Anything 3. + +This module handles model output processing, including tensor-to-numpy conversion, +batch dimension removal, and Prediction object creation. +""" + +from __future__ import annotations + +import numpy as np +import torch +from addict import Dict as AddictDict + +from depth_anything_3.specs import Prediction + + +class OutputProcessor: + """ + Output processor for converting model outputs to Prediction objects. + + Handles tensor-to-numpy conversion, batch dimension removal, + and creates structured Prediction objects with proper data types. + """ + + def __init__(self) -> None: + """Initialize the output processor.""" + + def __call__(self, model_output: dict[str, torch.Tensor]) -> Prediction: + """ + Convert model output to Prediction object. + + Args: + model_output: Model output dictionary containing depth, conf, extrinsics, intrinsics + Expected shapes: depth (B, N, 1, H, W), conf (B, N, 1, H, W), + extrinsics (B, N, 4, 4), intrinsics (B, N, 3, 3) + + Returns: + Prediction: Object containing depth estimation results with shapes: + depth (N, H, W), conf (N, H, W), extrinsics (N, 4, 4), intrinsics (N, 3, 3) + """ + # Extract data from batch dimension (B=1, N=number of images) + depth = self._extract_depth(model_output) + conf = self._extract_conf(model_output) + extrinsics = self._extract_extrinsics(model_output) + intrinsics = self._extract_intrinsics(model_output) + sky = self._extract_sky(model_output) + aux = self._extract_aux(model_output) + gaussians = model_output.get("gaussians", None) + scale_factor = model_output.get("scale_factor", None) + + return Prediction( + depth=depth, + sky=sky, + conf=conf, + extrinsics=extrinsics, + intrinsics=intrinsics, + is_metric=getattr(model_output, "is_metric", 0), + gaussians=gaussians, + aux=aux, + scale_factor=scale_factor, + ) + + def _extract_depth(self, model_output: dict[str, torch.Tensor]) -> np.ndarray: + """ + Extract depth tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Depth array with shape (N, H, W) + """ + depth = model_output["depth"].squeeze(0).squeeze(-1).cpu().numpy() # (N, H, W) + return depth + + def _extract_conf(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract confidence tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Confidence array with shape (N, H, W) or None + """ + conf = model_output.get("depth_conf", None) + if conf is not None: + conf = conf.squeeze(0).cpu().numpy() # (N, H, W) + return conf + + def _extract_extrinsics(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract extrinsics tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Extrinsics array with shape (N, 4, 4) or None + """ + extrinsics = model_output.get("extrinsics", None) + if extrinsics is not None: + extrinsics = extrinsics.squeeze(0).cpu().numpy() # (N, 4, 4) + return extrinsics + + def _extract_intrinsics(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract intrinsics tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Intrinsics array with shape (N, 3, 3) or None + """ + intrinsics = model_output.get("intrinsics", None) + if intrinsics is not None: + intrinsics = intrinsics.squeeze(0).cpu().numpy() # (N, 3, 3) + return intrinsics + + def _extract_sky(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract sky tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Sky mask array with shape (N, H, W) or None + """ + sky = model_output.get("sky", None) + if sky is not None: + sky = sky.squeeze(0).cpu().numpy() >= 0.5 # (N, H, W) + return sky + + def _extract_aux(self, model_output: dict[str, torch.Tensor]) -> AddictDict: + """ + Extract auxiliary data from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Dictionary containing auxiliary data + """ + aux = model_output.get("aux", None) + ret = AddictDict() + if aux is not None: + for k in aux.keys(): + if isinstance(aux[k], torch.Tensor): + ret[k] = aux[k].squeeze(0).cpu().numpy() + else: + ret[k] = aux[k] + return ret + + +# Backward compatibility alias +OutputAdapter = OutputProcessor diff --git a/depth_anything_3/utils/layout_helpers.py b/depth_anything_3/utils/layout_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..189c170b2007c979e580b69ca929638560923fb2 --- /dev/null +++ b/depth_anything_3/utils/layout_helpers.py @@ -0,0 +1,216 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""This file contains useful layout utilities for images. They are: + +- add_border: Add a border to an image. +- cat/hcat/vcat: Join images by arranging them in a line. If the images have different + sizes, they are aligned as specified (start, end, center). Allows you to specify a gap + between images. + +Images are assumed to be float32 tensors with shape (channel, height, width). +""" + +from typing import Any, Generator, Iterable, Literal, Union +import torch +from torch import Tensor + +Alignment = Literal["start", "center", "end"] +Axis = Literal["horizontal", "vertical"] +Color = Union[ + int, + float, + Iterable[int], + Iterable[float], + Tensor, + Tensor, +] + + +def _sanitize_color(color: Color) -> Tensor: # "#channel" + # Convert tensor to list (or individual item). + if isinstance(color, torch.Tensor): + color = color.tolist() + + # Turn iterators and individual items into lists. + if isinstance(color, Iterable): + color = list(color) + else: + color = [color] + + return torch.tensor(color, dtype=torch.float32) + + +def _intersperse(iterable: Iterable, delimiter: Any) -> Generator[Any, None, None]: + it = iter(iterable) + yield next(it) + for item in it: + yield delimiter + yield item + + +def _get_main_dim(main_axis: Axis) -> int: + return { + "horizontal": 2, + "vertical": 1, + }[main_axis] + + +def _get_cross_dim(main_axis: Axis) -> int: + return { + "horizontal": 1, + "vertical": 2, + }[main_axis] + + +def _compute_offset(base: int, overlay: int, align: Alignment) -> slice: + assert base >= overlay + offset = { + "start": 0, + "center": (base - overlay) // 2, + "end": base - overlay, + }[align] + return slice(offset, offset + overlay) + + +def overlay( + base: Tensor, # "channel base_height base_width" + overlay: Tensor, # "channel overlay_height overlay_width" + main_axis: Axis, + main_axis_alignment: Alignment, + cross_axis_alignment: Alignment, +) -> Tensor: # "channel base_height base_width" + # The overlay must be smaller than the base. + _, base_height, base_width = base.shape + _, overlay_height, overlay_width = overlay.shape + assert base_height >= overlay_height and base_width >= overlay_width + + # Compute spacing on the main dimension. + main_dim = _get_main_dim(main_axis) + main_slice = _compute_offset( + base.shape[main_dim], overlay.shape[main_dim], main_axis_alignment + ) + + # Compute spacing on the cross dimension. + cross_dim = _get_cross_dim(main_axis) + cross_slice = _compute_offset( + base.shape[cross_dim], overlay.shape[cross_dim], cross_axis_alignment + ) + + # Combine the slices and paste the overlay onto the base accordingly. + selector = [..., None, None] + selector[main_dim] = main_slice + selector[cross_dim] = cross_slice + result = base.clone() + result[selector] = overlay + return result + + +def cat( + main_axis: Axis, + *images: Iterable[Tensor], # "channel _ _" + align: Alignment = "center", + gap: int = 8, + gap_color: Color = 1, +) -> Tensor: # "channel height width" + """Arrange images in a line. The interface resembles a CSS div with flexbox.""" + device = images[0].device + gap_color = _sanitize_color(gap_color).to(device) + + # Find the maximum image side length in the cross axis dimension. + cross_dim = _get_cross_dim(main_axis) + cross_axis_length = max(image.shape[cross_dim] for image in images) + + # Pad the images. + padded_images = [] + for image in images: + # Create an empty image with the correct size. + padded_shape = list(image.shape) + padded_shape[cross_dim] = cross_axis_length + base = torch.ones(padded_shape, dtype=torch.float32, device=device) + base = base * gap_color[:, None, None] + padded_images.append(overlay(base, image, main_axis, "start", align)) + + # Intersperse separators if necessary. + if gap > 0: + # Generate a separator. + c, _, _ = images[0].shape + separator_size = [gap, gap] + separator_size[cross_dim - 1] = cross_axis_length + separator = torch.ones((c, *separator_size), dtype=torch.float32, device=device) + separator = separator * gap_color[:, None, None] + + # Intersperse the separator between the images. + padded_images = list(_intersperse(padded_images, separator)) + + return torch.cat(padded_images, dim=_get_main_dim(main_axis)) + + +def hcat( + *images: Iterable[Tensor], # "channel _ _" + align: Literal["start", "center", "end", "top", "bottom"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "horizontal", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "top": "start", + "bottom": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def vcat( + *images: Iterable[Tensor], # "channel _ _" + align: Literal["start", "center", "end", "left", "right"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "vertical", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "left": "start", + "right": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def add_border( + image: Tensor, # "channel height width" + border: int = 8, + color: Color = 1, +) -> Tensor: # "channel new_height new_width" + color = _sanitize_color(color).to(image) + c, h, w = image.shape + result = torch.empty( + (c, h + 2 * border, w + 2 * border), dtype=torch.float32, device=image.device + ) + result[:] = color[:, None, None] + result[:, border : h + border, border : w + border] = image + return result diff --git a/depth_anything_3/utils/logger.py b/depth_anything_3/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..0eb4f60696a085001cf4866ccfe1654170702a2d --- /dev/null +++ b/depth_anything_3/utils/logger.py @@ -0,0 +1,82 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys + + +class Color: + RED = "\033[91m" + YELLOW = "\033[93m" + WHITE = "\033[97m" + GREEN = "\033[92m" + RESET = "\033[0m" + + +LOG_LEVELS = {"ERROR": 0, "WARN": 1, "INFO": 2, "DEBUG": 3} + +COLOR_MAP = {"ERROR": Color.RED, "WARN": Color.YELLOW, "INFO": Color.WHITE, "DEBUG": Color.GREEN} + + +def get_env_log_level(): + level = os.environ.get("DA3_LOG_LEVEL", "INFO").upper() + return LOG_LEVELS.get(level, LOG_LEVELS["INFO"]) + + +class Logger: + def __init__(self): + self.level = get_env_log_level() + + def log(self, level_str, *args, **kwargs): + level_key = level_str.split(":")[0].strip() + level_val = LOG_LEVELS.get(level_key) + if level_val is None: + raise ValueError(f"Unknown log level: {level_str}") + if self.level >= level_val: + color = COLOR_MAP[level_key] + msg = " ".join(str(arg) for arg in args) + + # Align log level output in square brackets + # ERROR and DEBUG are 5 characters, INFO and WARN have an extra space for alignment + tag = level_key + if tag in ("INFO", "WARN"): + tag += " " + print( + f"{color}[{tag}] {msg}{Color.RESET}", + file=sys.stderr if level_key == "ERROR" else sys.stdout, + **kwargs, + ) + + def error(self, *args, **kwargs): + self.log("ERROR:", *args, **kwargs) + + def warn(self, *args, **kwargs): + self.log("WARN:", *args, **kwargs) + + def info(self, *args, **kwargs): + self.log("INFO:", *args, **kwargs) + + def debug(self, *args, **kwargs): + self.log("DEBUG:", *args, **kwargs) + + +logger = Logger() + +__all__ = ["logger"] + +if __name__ == "__main__": + logger.info("This is an info message") + logger.warn("This is a warning message") + logger.error("This is an error message") + logger.debug("This is a debug message") diff --git a/depth_anything_3/utils/memory.py b/depth_anything_3/utils/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..682dad76b2fccd606246144788cec0198eed2a62 --- /dev/null +++ b/depth_anything_3/utils/memory.py @@ -0,0 +1,127 @@ +""" +GPU memory utility helpers. + +Shared cleanup and memory checking logic used by both the backend API and +the Gradio UI to keep memory-management behavior consistent. +""" +from __future__ import annotations + +import gc + +from typing import Any, Dict, Optional + +import torch + + +def get_gpu_memory_info() -> Optional[Dict[str, Any]]: + """Return a snapshot of current GPU memory usage or None if CUDA not available. + + Keys in returned dict: total_gb, allocated_gb, reserved_gb, free_gb, utilization + """ + if not torch.cuda.is_available(): + return None + + try: + device = torch.cuda.current_device() + total_memory = torch.cuda.get_device_properties(device).total_memory + allocated_memory = torch.cuda.memory_allocated(device) + reserved_memory = torch.cuda.memory_reserved(device) + free_memory = total_memory - reserved_memory + + return { + "total_gb": total_memory / 1024 ** 3, + "allocated_gb": allocated_memory / 1024 ** 3, + "reserved_gb": reserved_memory / 1024 ** 3, + "free_gb": free_memory / 1024 ** 3, + "utilization": (reserved_memory / total_memory) * 100, + } + except Exception: + return None + + +def cleanup_cuda_memory() -> None: + """Perform a robust GPU cleanup sequence. + + This includes synchronizing, emptying caches, collecting IPC handles and + running the Python garbage collector. Use this instead of a raw + ``torch.cuda.empty_cache()`` where you need reliable freeing of GPU memory + between model loads or in error handling paths. + """ + try: + if torch.cuda.is_available(): + mem_before = get_gpu_memory_info() + + torch.cuda.synchronize() + torch.cuda.empty_cache() + # Collect cross-process cuda resources + try: + torch.cuda.ipc_collect() + except Exception: + # Older PyTorch versions or non-cuda devices may not support + # ipc_collect (no-op if not available) + pass + gc.collect() + + mem_after = get_gpu_memory_info() + if mem_before and mem_after: + freed = mem_before["reserved_gb"] - mem_after["reserved_gb"] + print( + f"CUDA cleanup: freed {freed:.2f}GB, " + f"available: {mem_after['free_gb']:.2f}GB/{mem_after['total_gb']:.2f}GB" + ) + else: + print("CUDA memory cleanup completed") + except Exception as e: + print(f"Warning: CUDA cleanup failed: {e}") + + +def check_memory_availability(required_gb: float = 2.0) -> tuple[bool, str]: + """Return whether at least ``required_gb`` seems available on the current GPU. + + The returned tuple is (is_available, message) with a human-friendly message. + """ + try: + if not torch.cuda.is_available(): + return False, "CUDA is not available" + + mem_info = get_gpu_memory_info() + if mem_info is None: + return True, "Cannot check memory, proceeding anyway" + + if mem_info["free_gb"] < required_gb: + return ( + False, + ( + f"Insufficient GPU memory: {mem_info['free_gb']:.2f}GB available, " + f"{required_gb:.2f}GB required. Total: {mem_info['total_gb']:.2f}GB, " + f"Used: {mem_info['reserved_gb']:.2f}GB ({mem_info['utilization']:.1f}%)" + ), + ) + + return ( + True, + ( + f"Memory check passed: {mem_info['free_gb']:.2f}GB available, " + f"{required_gb:.2f}GB required" + ), + ) + except Exception as e: + return True, f"Memory check failed: {e}, proceeding anyway" +def estimate_memory_requirement(num_images: int, process_res: int) -> float: + """Heuristic estimate for memory usage (GB) based on image count and resolution. + + This mirrors the simple policy used by the backend service so other code + (e.g., Gradio UI) can make consistent decisions when checking available + memory before loading a model or running inference. + + Args: + num_images: Number of images to process. + process_res: Processing resolution. + + Returns: + Estimated memory requirement in GB. + """ + base_memory = 2.0 + per_image_memory = (process_res / 504) ** 2 * 0.5 + total_memory = base_memory + (num_images * per_image_memory * 0.1) + return total_memory diff --git a/depth_anything_3/utils/model_loading.py b/depth_anything_3/utils/model_loading.py new file mode 100644 index 0000000000000000000000000000000000000000..b6d43b5bbab5a0989eae272422192103cd4d5bee --- /dev/null +++ b/depth_anything_3/utils/model_loading.py @@ -0,0 +1,149 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model loading and state dict conversion utilities. +""" + +from typing import Dict, Tuple +import torch + +from depth_anything_3.utils.logger import logger + + +def convert_general_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert general model state dict to match current model architecture. + + Args: + state_dict: Original state dictionary + + Returns: + Converted state dictionary + """ + # Replace module prefixes + state_dict = {k.replace("module.", "model."): v for k, v in state_dict.items()} + state_dict = {k.replace(".net.", ".backbone."): v for k, v in state_dict.items()} + + # Remove camera token if present + if "model.backbone.pretrained.camera_token" in state_dict: + del state_dict["model.backbone.pretrained.camera_token"] + + # Replace camera token naming + state_dict = { + k.replace(".camera_token_extra", ".camera_token"): v for k, v in state_dict.items() + } + + # Replace head naming + state_dict = { + k.replace("model.all_heads.camera_cond_head", "model.cam_enc"): v + for k, v in state_dict.items() + } + state_dict = { + k.replace("model.all_heads.camera_head", "model.cam_dec"): v for k, v in state_dict.items() + } + state_dict = {k.replace(".more_mlps.", ".backbone."): v for k, v in state_dict.items()} + state_dict = {k.replace(".fc_rot.", ".fc_qvec."): v for k, v in state_dict.items()} + state_dict = { + k.replace("model.all_heads.head", "model.head"): v for k, v in state_dict.items() + } + + # Replace output naming + state_dict = { + k.replace("output_conv2_additional.sky_mask", "sky_output_conv2"): v + for k, v in state_dict.items() + } + state_dict = {k.replace("_ray.", "_aux."): v for k, v in state_dict.items()} + + # Update GS-DPT head naming and value + state_dict = {k.replace("gaussian_param_head.", "gs_head."): v for k, v in state_dict.items()} + + return state_dict + + +def convert_metric_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert metric model state dict to match current model architecture. + + Args: + state_dict: Original metric state dictionary + + Returns: + Converted state dictionary + """ + # Add module prefix for metric models + state_dict = {"module." + k: v for k, v in state_dict.items()} + return convert_general_state_dict(state_dict) + + +def load_pretrained_weights(model, model_path: str, is_metric: bool = False) -> Tuple[list, list]: + """ + Load pretrained weights for a single model. + + Args: + model: Model instance to load weights into + model_path: Path to the pretrained weights + is_metric: Whether this is a metric model + + Returns: + Tuple of (missed_keys, unexpected_keys) + """ + state_dict = torch.load(model_path, map_location="cpu") + + if is_metric: + state_dict = convert_metric_state_dict(state_dict) + else: + state_dict = convert_general_state_dict(state_dict) + + missed, unexpected = model.load_state_dict(state_dict, strict=False) + logger.info("Missed keys:", missed) + logger.info("Unexpected keys:", unexpected) + + return missed, unexpected + + +def load_pretrained_nested_weights( + model, main_model_path: str, metric_model_path: str +) -> Tuple[list, list]: + """ + Load pretrained weights for a nested model with both main and metric branches. + + Args: + model: Nested model instance + main_model_path: Path to main model weights + metric_model_path: Path to metric model weights + + Returns: + Tuple of (missed_keys, unexpected_keys) + """ + # Load main model weights + state_dict0 = torch.load(main_model_path, map_location="cpu") + state_dict0 = convert_general_state_dict(state_dict0) + state_dict0 = {k.replace("model.", "model.da3."): v for k, v in state_dict0.items()} + + # Load metric model weights + state_dict1 = torch.load(metric_model_path, map_location="cpu") + state_dict1 = convert_metric_state_dict(state_dict1) + state_dict1 = {k.replace("model.", "model.da3_metric."): v for k, v in state_dict1.items()} + + # Combine state dictionaries + combined_state_dict = state_dict0.copy() + combined_state_dict.update(state_dict1) + + missed, unexpected = model.load_state_dict(combined_state_dict, strict=False) + + print("Missed keys:", missed) + print("Unexpected keys:", unexpected) + + return missed, unexpected diff --git a/depth_anything_3/utils/parallel_utils.py b/depth_anything_3/utils/parallel_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff108e95d205097f9f2012ab87ad9e265d58d8d --- /dev/null +++ b/depth_anything_3/utils/parallel_utils.py @@ -0,0 +1,133 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +from functools import wraps +from multiprocessing.pool import ThreadPool +from threading import Thread +from typing import Callable, Dict, List +import imageio +from tqdm import tqdm + + +def async_call_func(func): + @wraps(func) + async def wrapper(*args, **kwargs): + loop = asyncio.get_event_loop() + # Use run_in_executor to run the blocking function in a separate thread + return await loop.run_in_executor(None, func, *args, **kwargs) + + return wrapper + + +slice_func = lambda chunk_index, chunk_dim, chunk_size: [slice(None)] * chunk_dim + [ + slice(chunk_index, chunk_index + chunk_size) +] + + +def async_call(fn): + def wrapper(*args, **kwargs): + Thread(target=fn, args=args, kwargs=kwargs).start() + + return wrapper + + +def _save_image_impl(save_img, save_path): + """Common implementation for saving images synchronously or asynchronously""" + os.makedirs(os.path.dirname(save_path), exist_ok=True) + imageio.imwrite(save_path, save_img) + + +@async_call +def save_image_async(save_img, save_path): + """Save image asynchronously""" + _save_image_impl(save_img, save_path) + + +def save_image(save_img, save_path): + """Save image synchronously""" + _save_image_impl(save_img, save_path) + + +def parallel_execution( + *args, + action: Callable, + num_processes=32, + print_progress=False, + sequential=False, + async_return=False, + desc=None, + **kwargs, +): + # Partially copy from EasyVolumetricVideo (parallel_execution) + # NOTE: we expect first arg / or kwargs to be distributed + # NOTE: print_progress arg is reserved. + # `*args` packs all positional arguments passed to the function into a tuple + args = list(args) + + def get_length(args: List, kwargs: Dict): + for a in args: + if isinstance(a, list): + return len(a) + for v in kwargs.values(): + if isinstance(v, list): + return len(v) + raise NotImplementedError + + def get_action_args(length: int, args: List, kwargs: Dict, i: int): + action_args = [ + (arg[i] if isinstance(arg, list) and len(arg) == length else arg) for arg in args + ] + # TODO: Support all types of iterable + action_kwargs = { + key: ( + kwargs[key][i] + if isinstance(kwargs[key], list) and len(kwargs[key]) == length + else kwargs[key] + ) + for key in kwargs + } + return action_args, action_kwargs + + if not sequential: + # Create ThreadPool + pool = ThreadPool(processes=num_processes) + + # Spawn threads + results = [] + asyncs = [] + length = get_length(args, kwargs) + for i in range(length): + action_args, action_kwargs = get_action_args(length, args, kwargs, i) + async_result = pool.apply_async(action, action_args, action_kwargs) + asyncs.append(async_result) + + # Join threads and get return values + if not async_return: + for async_result in tqdm(asyncs, desc=desc, disable=not print_progress): + results.append(async_result.get()) # will sync the corresponding thread + pool.close() + pool.join() + return results + else: + return pool + else: + results = [] + length = get_length(args, kwargs) + for i in tqdm(range(length), desc=desc, disable=not print_progress): + action_args, action_kwargs = get_action_args(length, args, kwargs, i) + async_result = action(*action_args, **action_kwargs) + results.append(async_result) + return results diff --git a/depth_anything_3/utils/pca_utils.py b/depth_anything_3/utils/pca_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2b9eee268cd8692d885bf093700a5752077b9d7d --- /dev/null +++ b/depth_anything_3/utils/pca_utils.py @@ -0,0 +1,284 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +PCA utilities for feature visualization and dimensionality reduction (video-friendly). +- Support frame-by-frame: transform_frame / transform_video +- Support one-time global PCA fitting and reuse (mean, V3) for stable colors +- Support Procrustes alignment (solving principal component order/sign/rotation jumps) +- Support global fixed or temporal EMA for percentiles (time dimension only, no spatial) +""" + +import numpy as np +import torch + + +def pca_to_rgb_4d_bf16_percentile( + x_np: np.ndarray, + device=None, + q_oversample: int = 6, + clip_percent: float = 10.0, # Percentage to clip from top and bottom (0~49.9) + return_uint8: bool = False, + enable_autocast_bf16: bool = True, +): + """ + Reduce numpy array of shape (49, 27, 36, 3072) to 3D via PCA and visualize as (49, 27, 36, 3). + - PCA uses torch.pca_lowrank (randomized SVD), defaults to GPU. + - Uses CUDA bf16 autocast in computation (if available), + then per-channel percentile clipping and normalization. + - Default removes 5% outliers from top and bottom (adjustable via clip_percent) to + improve visualization contrast. + + Parameters + ---------- + x_np : np.ndarray + Shape must be (49, 27, 36, 3072). dtype recommended float32/float64. + device : str | None + Specify 'cuda' or 'cpu'. Auto-select if None (prefer cuda). + q_oversample : int + Oversampling q for pca_lowrank, must be >= 3. + Slightly larger than target dim (3) is more stable, default 6. + clip_percent : float + Percentage to clip from top and bottom (0~49.9), + e.g. 5.0 means clip lowest 5% and highest 5% per channel. + return_uint8 : bool + True returns uint8(0~255), otherwise returns float32(0~1). + enable_autocast_bf16 : bool + Enable bf16 autocast on CUDA. + + Returns + ------- + np.ndarray + Array of shape (49, 27, 36, 3), float32[0,1] or uint8[0,255]. + """ + assert ( + x_np.ndim == 4 + ) # and x_np.shape[-1] == 3072, f"expect (49,27,36,3072), got {x_np.shape}" + B1, B2, B3, D = x_np.shape + N = B1 * B2 * B3 + + # Device selection + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Convert input to torch, unified float32 + X = torch.from_numpy(x_np.reshape(N, D)).to(device=device, dtype=torch.float32) + + # Parameter and safety checks + k = 3 + q = max(int(q_oversample), k) + clip_percent = float(clip_percent) + if not (0.0 <= clip_percent < 50.0): + raise ValueError( + "clip_percent must be in [0, 50), e.g. 5.0 means clip 5% from top and bottom" + ) + low = clip_percent / 100.0 + high = 1.0 - low + + with torch.no_grad(): + # Zero mean + mean = X.mean(dim=0, keepdim=True) + Xc = X - mean + + # Main computation: PCA + projection, try to use bf16 + # (auto-fallback if operator not supported) + device.startswith("cuda") and enable_autocast_bf16 + U, S, V = torch.pca_lowrank(Xc, q=q, center=False) # V: (D, q) + V3 = V[:, :k] # (3072, 3) + PCs = Xc @ V3 # (N, 3) + + # === Per-channel percentile clipping and normalization to [0,1] === + # Vectorized one-time calculation of low/high percentiles for each channel + qs = torch.tensor([low, high], device=PCs.device, dtype=PCs.dtype) + qvals = torch.quantile(PCs, q=qs, dim=0) # Shape (2, 3) + lo = qvals[0] # (3,) + hi = qvals[1] # (3,) + + # Avoid degenerate case where hi==lo + denom = torch.clamp(hi - lo, min=1e-8) + + # Broadcast clipping + normalization + PCs = torch.clamp(PCs, lo, hi) + PCs = (PCs - lo) / denom # (N, 3) in [0,1] + + # Restore 4D + PCs = PCs.reshape(B1, B2, B3, k) + + # Output + if return_uint8: + out = (PCs * 255.0).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + else: + out = PCs.clamp(0, 1).to(torch.float32).cpu().numpy() + + return out + + +class PCARGBVisualizer: + """ + Stable PCA→RGB for video features shaped (T, H, W, D) or a single frame (H, W, D). + - Global mean/V3 reference for stable colors + - Per-frame PCA with Procrustes alignment to V3_ref (basis_mode='procrustes') + - Percentile normalization with global or EMA stats (time-only, no spatial smoothing) + """ + + def __init__( + self, + device=None, + q_oversample: int = 16, + clip_percent: float = 10.0, + return_uint8: bool = False, + enable_autocast_bf16: bool = True, + basis_mode: str = "procrustes", # 'fixed' | 'procrustes' + percentile_mode: str = "ema", # 'global' | 'ema' + ema_alpha: float = 0.1, + denom_eps: float = 1e-4, + ): + assert 0.0 <= clip_percent < 50.0 + assert basis_mode in ("fixed", "procrustes") + assert percentile_mode in ("global", "ema") + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.q = max(int(q_oversample), 6) + self.clip_percent = float(clip_percent) + self.return_uint8 = return_uint8 + self.enable_autocast_bf16 = enable_autocast_bf16 + self.basis_mode = basis_mode + self.percentile_mode = percentile_mode + self.ema_alpha = float(ema_alpha) + self.denom_eps = float(denom_eps) + + # reference state + self.mean_ref = None # (1, D) + self.V3_ref = None # (D, 3) + self.lo_ref = None # (3,) + self.hi_ref = None # (3,) + + @torch.no_grad() + def fit_reference(self, frames): + """ + Fit global mean/V3 and initialize percentiles from a reference set. + frames: ndarray (T,H,W,D) or list of (H,W,D) + """ + if isinstance(frames, np.ndarray): + if frames.ndim != 4: + raise ValueError("fit_reference expects (T,H,W,D) ndarray.") + T, H, W, D = frames.shape + X = torch.from_numpy(frames.reshape(T * H * W, D)) + else: # list of (H,W,D) + xs = [torch.from_numpy(x.reshape(-1, x.shape[-1])) for x in frames] + D = xs[0].shape[-1] + X = torch.cat(xs, dim=0) + + X = X.to(self.device, dtype=torch.float32) + X = torch.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6) + + mean = X.mean(0, keepdim=True) + Xc = X - mean + + U, S, V = torch.pca_lowrank(Xc, q=max(self.q, 8), center=False) + V3 = V[:, :3] # (D,3) + + PCs = Xc @ V3 + low = self.clip_percent / 100.0 + high = 1.0 - low + qs = torch.tensor([low, high], device=PCs.device, dtype=PCs.dtype) + qvals = torch.quantile(PCs, q=qs, dim=0) + lo, hi = qvals[0], qvals[1] + + self.mean_ref = mean + self.V3_ref = V3 + if self.percentile_mode == "global": + self.lo_ref, self.hi_ref = lo, hi + else: + self.lo_ref = lo.clone() + self.hi_ref = hi.clone() + + @torch.no_grad() + def _project_with_stable_colors(self, X: torch.Tensor) -> torch.Tensor: + """ + X: (N,D) where N = H*W + Returns PCs_raw: (N,3) using stable basis (fixed or Procrustes-aligned) + """ + assert self.mean_ref is not None and self.V3_ref is not None, "Call fit_reference() first." + X = torch.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6) + Xc = X - self.mean_ref + + if self.basis_mode == "fixed": + V3_used = self.V3_ref + else: + U, S, V = torch.pca_lowrank(Xc, q=max(self.q, 6), center=False) + V3 = V[:, :3] # (D,3) + M = V3.T @ self.V3_ref + Uo, So, Vh = torch.linalg.svd(M) + R = Uo @ Vh + V3_used = V3 @ R + # Optional polarity fix via anchor + a = self.V3_ref.mean(0, keepdim=True) + sign = torch.sign((V3_used * a).sum(0, keepdim=True)).clamp(min=-1) + V3_used = V3_used * sign + + return Xc @ V3_used + + @torch.no_grad() + def _normalize_rgb(self, PCs_raw: torch.Tensor) -> torch.Tensor: + assert self.lo_ref is not None and self.hi_ref is not None + if self.percentile_mode == "global": + lo, hi = self.lo_ref, self.hi_ref + else: + low = self.clip_percent / 100.0 + high = 1.0 - low + qs = torch.tensor([low, high], device=PCs_raw.device, dtype=PCs_raw.dtype) + qvals = torch.quantile(PCs_raw, q=qs, dim=0) + lo_now, hi_now = qvals[0], qvals[1] + a = self.ema_alpha + self.lo_ref = (1 - a) * self.lo_ref + a * lo_now + self.hi_ref = (1 - a) * self.hi_ref + a * hi_now + lo, hi = self.lo_ref, self.hi_ref + + denom = torch.clamp(hi - lo, min=self.denom_eps) + PCs = torch.clamp(PCs_raw, lo, hi) + PCs = (PCs - lo) / denom + return PCs.clamp_(0, 1) + + @torch.no_grad() + def transform_frame(self, frame: np.ndarray) -> np.ndarray: + """ + frame: (H,W,D) -> (H,W,3) + """ + if frame.ndim != 3: + raise ValueError("transform_frame expects (H,W,D).") + H, W, D = frame.shape + X = torch.from_numpy(frame.reshape(H * W, D)).to(self.device, dtype=torch.float32) + PCs_raw = self._project_with_stable_colors(X) + PCs = self._normalize_rgb(PCs_raw).reshape(H, W, 3) + if self.return_uint8: + return (PCs * 255.0).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + return PCs.to(torch.float32).cpu().numpy() + + @torch.no_grad() + def transform_video(self, frames) -> np.ndarray: + """ + frames: (T,H,W,D) or list of (H,W,D) + returns: (T,H,W,3) + """ + outs = [] + if isinstance(frames, np.ndarray): + if frames.ndim != 4: + raise ValueError("transform_video expects (T,H,W,D).") + T, H, W, D = frames.shape + for t in range(T): + outs.append(self.transform_frame(frames[t])) + else: + for f in frames: + outs.append(self.transform_frame(f)) + return np.stack(outs, axis=0) diff --git a/depth_anything_3/utils/pose_align.py b/depth_anything_3/utils/pose_align.py new file mode 100644 index 0000000000000000000000000000000000000000..695d07fc1210e3cc3614c9b22c56d765b1106500 --- /dev/null +++ b/depth_anything_3/utils/pose_align.py @@ -0,0 +1,347 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List +import numpy as np +import torch +from evo.core.trajectory import PosePath3D + +from depth_anything_3.utils.geometry import affine_inverse, affine_inverse_np + + +def batch_apply_alignment_to_enc( + rots: torch.Tensor, trans: torch.Tensor, scales: torch.Tensor, enc_list: List[torch.Tensor] +): + pass + + +def batch_apply_alignment_to_ext( + rots: torch.Tensor, trans: torch.Tensor, scales: torch.Tensor, ext: torch.Tensor +): + device, _ = ext.device, ext.dtype + if ext.shape[-2:] == (3, 4): + pad = torch.zeros((*ext.shape[:-2], 4, 4), dtype=ext.dtype, device=device) + pad[..., :3, :4] = ext + pad[..., 3, 3] = 1.0 + ext = pad + pose_est = affine_inverse(ext) + pose_new_align_rot = rots[:, None] @ pose_est[..., :3, :3] + pose_new_align_trans = ( + scales[:, None, None] * (rots[:, None] @ pose_est[..., :3, 3:])[..., 0] + trans[:, None] + ) + pose_new_align = torch.zeros_like(ext) + pose_new_align[..., :3, :3] = pose_new_align_rot + pose_new_align[..., :3, 3] = pose_new_align_trans + pose_new_align[..., 3, 3] = 1.0 + return affine_inverse(pose_new_align)[:, :3] + + +def batch_align_poses_umeyama(ext_ref: torch.Tensor, ext_est: torch.Tensor): + device, dtype = ext_ref.device, ext_ref.dtype + assert ext_ref.dtype in [torch.float32, torch.float64] + assert ext_est.dtype in [torch.float32, torch.float64] + assert ext_ref.requires_grad is False + assert ext_est.requires_grad is False + rots, trans, scales = [], [], [] + for b in range(ext_ref.shape[0]): + r, t, s = align_poses_umeyama(ext_ref[b].cpu().numpy(), ext_est[b].cpu().numpy()) + rots.append(torch.from_numpy(r).to(device=device, dtype=dtype)) + trans.append(torch.from_numpy(t).to(device=device, dtype=dtype)) + scales.append(torch.tensor(s, device=device, dtype=dtype)) + return torch.stack(rots), torch.stack(trans), torch.stack(scales) + + +# Dependencies: affine_inverse_np, PosePath3D (maintain consistency with your existing project) + + +def _to44(ext): + if ext.shape[1] == 3: + out = np.eye(4)[None].repeat(len(ext), 0) + out[:, :3, :4] = ext + return out + return ext + + +def _poses_from_ext(ext_ref, ext_est): + ext_ref = _to44(ext_ref) + ext_est = _to44(ext_est) + pose_ref = affine_inverse_np(ext_ref) + pose_est = affine_inverse_np(ext_est) + return pose_ref, pose_est + + +def _umeyama_sim3_from_paths(pose_ref, pose_est): + path_ref = PosePath3D(poses_se3=pose_ref.copy()) + path_est = PosePath3D(poses_se3=pose_est.copy()) + r, t, s = path_est.align(path_ref, correct_scale=True) + pose_est_aligned = np.stack(path_est.poses_se3) + return r, t, s, pose_est_aligned + + +def _apply_sim3_to_poses(poses, r, t, s): + out = poses.copy() + Ri = poses[:, :3, :3] + ti = poses[:, :3, 3] + out[:, :3, :3] = r @ Ri + out[:, :3, 3] = (r @ (s * ti.T)).T + t + return out + + +def _median_nn_thresh(pose_ref, pose_est_aligned): + P_ref = pose_ref[:, :3, 3] + P_est = pose_est_aligned[:, :3, 3] + dists = [] + for p in P_est: + dd = np.linalg.norm(P_ref - p[None, :], axis=1) + dists.append(dd.min()) + return float(np.median(dists)) if dists else 0.0 + + +def _ransac_align_sim3( + pose_ref, pose_est, sub_n=None, inlier_thresh=None, max_iters=10, random_state=None +): + rng = np.random.default_rng(random_state) + N = pose_ref.shape[0] + idx_all = np.arange(N) + if sub_n is None: + sub_n = max(3, (N + 1) // 2) + else: + sub_n = max(3, min(sub_n, N)) + + # Pre-alignment + default threshold + r0, t0, s0, pose_est0 = _umeyama_sim3_from_paths(pose_ref, pose_est) + if inlier_thresh is None: + inlier_thresh = _median_nn_thresh(pose_ref, pose_est0) + + P_ref_all = pose_ref[:, :3, 3] + + best_model = (r0, t0, s0) + best_inliers = None + best_score = (-1, np.inf) # (num_inliers, mean_err) + + for _ in range(max_iters): + sample = rng.choice(idx_all, size=sub_n, replace=False) + try: + r, t, s, _ = _umeyama_sim3_from_paths(pose_ref[sample], pose_est[sample]) + except Exception: + continue + pose_h = _apply_sim3_to_poses(pose_est, r, t, s) + P_h = pose_h[:, :3, 3] + errs = np.linalg.norm(P_h - P_ref_all, axis=1) # Match by same index + inliers = errs <= inlier_thresh + k = int(inliers.sum()) + mean_err = float(errs[inliers].mean()) if k > 0 else np.inf + if (k > best_score[0]) or (k == best_score[0] and mean_err < best_score[1]): + best_score = (k, mean_err) + best_model = (r, t, s) + best_inliers = inliers + + # Fit again with best inliers + if best_inliers is not None and best_inliers.sum() >= 3: + r, t, s, _ = _umeyama_sim3_from_paths(pose_ref[best_inliers], pose_est[best_inliers]) + else: + r, t, s = best_model + return r, t, s + + +def align_poses_umeyama( + ext_ref: np.ndarray, + ext_est: np.ndarray, + return_aligned=False, + ransac=False, + sub_n=None, + inlier_thresh=None, + ransac_max_iters=10, + random_state=None, +): + """ + Align estimated trajectory to reference using Umeyama Sim(3). + Default no RANSAC; if ransac=True, use RANSAC (max iterations default 10). + - sub_n defaults to half the number of frames (rounded up, at least 3) + - inlier_thresh defaults to median of "distance from each estimated pose to + nearest reference pose after pre-alignment" + Returns rotation (3x3), translation (3,), scale; optionally returns aligned extrinsics (4x4). + """ + pose_ref, pose_est = _poses_from_ext(ext_ref, ext_est) + + if not ransac: + r, t, s, pose_est_aligned = _umeyama_sim3_from_paths(pose_ref, pose_est) + else: + r, t, s = _ransac_align_sim3( + pose_ref, + pose_est, + sub_n=sub_n, + inlier_thresh=inlier_thresh, + max_iters=ransac_max_iters, + random_state=random_state, + ) + pose_est_aligned = _apply_sim3_to_poses(pose_est, r, t, s) + + if return_aligned: + ext_est_aligned = affine_inverse_np(pose_est_aligned) + return r, t, s, ext_est_aligned + return r, t, s + + +# def align_poses_umeyama(ext_ref: np.ndarray, ext_est: np.ndarray, return_aligned=False): +# """ +# Align estimated trajectory to reference trajectory using Umeyama Sim(3) +# alignment (via evo PosePath3D). # noqa +# Returns rotation, translation, and scale. +# """ +# # If input extrinsics are 3x4, convert to 4x4 by padding +# if ext_ref.shape[1] == 3: +# ext_ref_ = np.eye(4)[None].repeat(len(ext_ref), 0) +# ext_ref_[:, :3] = ext_ref +# ext_ref = ext_ref_ +# if ext_est.shape[1] == 3: +# ext_est_ = np.eye(4)[None].repeat(len(ext_est), 0) +# ext_est_[:, :3] = ext_est +# ext_est = ext_est_ + +# # Convert to camera poses (inverse extrinsics) +# pose_ref = affine_inverse_np(ext_ref) +# pose_est = affine_inverse_np(ext_est) + +# # Create evo PosePath3D objects +# path_ref = PosePath3D(poses_se3=pose_ref) +# path_est = PosePath3D(poses_se3=pose_est) +# r, t, s = path_est.align(path_ref, correct_scale=True) +# if return_aligned: +# return r, t, s, affine_inverse_np(np.stack(path_est.poses_se3)) +# else: +# return r, t, s + + +def apply_umeyama_alignment_to_ext( + rot: np.ndarray, # (3,3) + trans: np.ndarray, # (3,) or (1,3) + scale: float, + ext_est: np.ndarray, # (...,4,4) or (...,3,4) +) -> np.ndarray: + """ + Apply Sim(3) (R, t, s) to a batch of world-to-camera extrinsics ext_est. + Returns the aligned extrinsics, with the same shape as input. + """ + + # Allow 3x4 extrinsics: pad to 4x4 + if ext_est.shape[-2:] == (3, 4): + pad = np.zeros((*ext_est.shape[:-2], 4, 4), dtype=ext_est.dtype) + pad[..., :3, :4] = ext_est + pad[..., 3, 3] = 1.0 + ext_est = pad + + # Convert world-to-camera to camera-to-world + pose_est = affine_inverse_np(ext_est) # (...,4,4) + R_e = pose_est[..., :3, :3] # (...,3,3) + t_e = pose_est[..., :3, 3] # (...,3) + + # Apply Sim(3) transformation + R_a = np.einsum("ij,...jk->...ik", rot, R_e) # (...,3,3) + t_a = scale * np.einsum("ij,...j->...i", rot, t_e) + trans # (...,3) + + # Assemble the transformed pose + pose_a = np.zeros_like(pose_est) + pose_a[..., :3, :3] = R_a + pose_a[..., :3, 3] = t_a + pose_a[..., 3, 3] = 1.0 + + # Convert back to world-to-camera + return affine_inverse_np(pose_a) + + +def transform_points_sim3(points, rot, trans, scale, inverse=False): + """ + Sim(3) transform point cloud + points: (N, 3) + rot: (3, 3) + trans: (3,) or (1, 3) + scale: float + inverse: Whether to do inverse transform (ref->est) + Returns: (N, 3) + """ + if not inverse: + # Forward: est -> ref + return scale * (points @ rot.T) + trans + else: + # Inverse: ref -> est + return ((points - trans) @ rot) / scale + + +def _rand_rot(): + u1, u2, u3 = np.random.rand(3) + q = np.array( + [ + np.sqrt(1 - u1) * np.sin(2 * np.math.pi * u2), + np.sqrt(1 - u1) * np.cos(2 * np.math.pi * u2), + np.sqrt(u1) * np.sin(2 * np.math.pi * u3), + np.sqrt(u1) * np.cos(2 * np.math.pi * u3), + ] + ) + w, x, y, z = q + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ] + ) + + +def _rand_pose(): + R, t = _rand_rot(), np.random.randn(3) + P = np.eye(4) + P[:3, :3] = R + P[:3, 3] = t + return P + + +if __name__ == "__main__": + np.random.seed(42) + # 1. Randomly generate reference trajectory and Sim(3) + N = 8 + pose_ref = np.stack([_rand_pose() for _ in range(N)]) # (N,4,4) cam→world + rot_gt = _rand_rot() + scale_gt = 2.3 + trans_gt = np.random.randn(3) + # 2. Generate estimated trajectory (apply Sim(3)) + pose_est = np.zeros_like(pose_ref) + for i in range(N): + R = pose_ref[i][:3, :3] + t = pose_ref[i][:3, 3] + pose_est[i][:3, :3] = rot_gt @ R + pose_est[i][:3, 3] = scale_gt * (rot_gt @ t) + trans_gt + pose_est[i][3, 3] = 1.0 + # 3. Get extrinsics (world->cam) + ext_ref = affine_inverse_np(pose_ref) + ext_est = affine_inverse_np(pose_est) + # 4. Use umeyama alignment, estimate Sim(3) + r_est, t_est, s_est = align_poses_umeyama(ext_ref, ext_est) + print("GT scale:", scale_gt, "Estimated:", s_est) + print("GT trans:", trans_gt, "Estimated:", t_est) + print("GT rot:\n", rot_gt, "\nEstimated:\n", r_est) + # 5. Random point cloud, in ref frame + num_points = 100 + points_ref = np.random.randn(num_points, 3) + # 6. Use GT Sim(3) inverse transform to est frame + points_est = transform_points_sim3(points_ref, rot_gt, trans_gt, scale_gt, inverse=True) + # 7. Use estimated Sim(3) forward transform back to ref frame + points_ref_recovered = transform_points_sim3(points_est, r_est, t_est, s_est, inverse=False) + # 8. Check error + err = np.abs(points_ref_recovered - points_ref) + print("Point cloud sim3 transform error (mean abs):", err.mean()) + print("Point cloud sim3 transform error (max abs):", err.max()) + assert err.mean() < 1e-6, "Mean sim3 transform error too large!" + assert err.max() < 1e-5, "Max sim3 transform error too large!" + print("Sim(3) point cloud transform & alignment test passed!") diff --git a/depth_anything_3/utils/ray_utils.py b/depth_anything_3/utils/ray_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6244dc40089a7bfda7889ddc153d48c44f717da7 --- /dev/null +++ b/depth_anything_3/utils/ray_utils.py @@ -0,0 +1,523 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from einops import repeat +from .geometry import unproject_depth + + +def compute_optimal_rotation_intrinsics_batch( + rays_origin, rays_target, z_threshold=1e-4, reproj_threshold=0.2, weights=None, + n_sample = None, + n_iter=100, + num_sample_for_ransac=8, + rand_sample_iters_idx=None, +): + """ + Args: + rays_origin (torch.Tensor): (B, N, 3) + rays_target (torch.Tensor): (B, N, 3) + z_threshold (float): Threshold for z value to be considered valid. + + Returns: + R (torch.tensor): (3, 3) + focal_length (torch.tensor): (2,) + principal_point (torch.tensor): (2,) + """ + device = rays_origin.device + B, N, _ = rays_origin.shape + z_mask = torch.logical_and( + torch.abs(rays_target[:, :, 2]) > z_threshold, torch.abs(rays_origin[:, :, 2]) > z_threshold + ) # (B, N, 1) + rays_origin = rays_origin.clone() + rays_target = rays_target.clone() + rays_origin[:, :, 0][z_mask] /= rays_origin[:, :, 2][z_mask] + rays_origin[:, :, 1][z_mask] /= rays_origin[:, :, 2][z_mask] + rays_target[:, :, 0][z_mask] /= rays_target[:, :, 2][z_mask] + rays_target[:, :, 1][z_mask] /= rays_target[:, :, 2][z_mask] + + rays_origin = rays_origin[:, :, :2] + rays_target = rays_target[:, :, :2] + assert weights is not None, "weights must be provided" + weights[~z_mask] = 0 + + A_list = [] + max_chunk_size = 2 + for i in range(0, rays_origin.shape[0], max_chunk_size): + A = ransac_find_homography_weighted_fast_batch( + rays_origin[i:i+max_chunk_size], + rays_target[i:i+max_chunk_size], + weights[i:i+max_chunk_size], + n_iter=n_iter, + n_sample = n_sample, + num_sample_for_ransac=num_sample_for_ransac, + reproj_threshold=reproj_threshold, + rand_sample_iters_idx=rand_sample_iters_idx, + max_inlier_num=8000, + ) + A = A.to(device) + A_need_inv_mask = torch.linalg.det(A) < 0 + A[A_need_inv_mask] = -A[A_need_inv_mask] + A_list.append(A) + + A = torch.cat(A_list, dim=0) + + R_list = [] + f_list = [] + pp_list = [] + for i in range(A.shape[0]): + R, L = ql_decomposition(A[i]) + L = L / L[2][2] + + f = torch.stack((L[0][0], L[1][1])) + pp = torch.stack((L[2][0], L[2][1])) + R_list.append(R) + f_list.append(f) + pp_list.append(pp) + + R = torch.stack(R_list) + f = torch.stack(f_list) + pp = torch.stack(pp_list) + + return R, f, pp + + +# https://www.reddit.com/r/learnmath/comments/v1crd7/linear_algebra_qr_to_ql_decomposition/ +def ql_decomposition(A): + P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=A.device).float() + A_tilde = torch.matmul(A, P) + Q_tilde, R_tilde = torch.linalg.qr(A_tilde) + Q = torch.matmul(Q_tilde, P) + L = torch.matmul(torch.matmul(P, R_tilde), P) + d = torch.diag(L) + Q[:, 0] *= torch.sign(d[0]) + Q[:, 1] *= torch.sign(d[1]) + Q[:, 2] *= torch.sign(d[2]) + L[0] *= torch.sign(d[0]) + L[1] *= torch.sign(d[1]) + L[2] *= torch.sign(d[2]) + return Q, L + +def find_homography_least_squares_weighted_torch(src_pts, dst_pts, confident_weight): + """ + src_pts: (N,2) source points (torch.Tensor, float32/float64) + dst_pts: (N,2) target points (torch.Tensor, float32/float64) + confident_weight: (N,) weights (torch.Tensor) + Returns: (3,3) homography matrix H (torch.Tensor) + """ + assert src_pts.shape == dst_pts.shape + N = src_pts.shape[0] + if N < 4: + raise ValueError("At least 4 points are required to compute homography.") + assert confident_weight.shape == (N,) + + w = confident_weight.sqrt().unsqueeze(1) # (N,1) + + x = src_pts[:, 0:1] # (N,1) + y = src_pts[:, 1:2] # (N,1) + u = dst_pts[:, 0:1] + v = dst_pts[:, 1:2] + + zeros = torch.zeros_like(x) + + # Construct A matrix (2N, 9) + A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=1) + A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=1) + A = torch.cat([A1, A2], dim=0) # (2N, 9) + + # SVD + # Note: torch.linalg.svd returns U, S, Vh, where Vh is the transpose of V + _, _, Vh = torch.linalg.svd(A) + H = Vh[-1].reshape(3, 3) + H = H / H[-1, -1] + return H + + +def ransac_find_homography_weighted( + src_pts, + dst_pts, + confident_weight, + n_iter=100, + sample_ratio=0.2, + reproj_threshold=3.0, + num_sample_for_ransac=16, + random_seed=None, +): + """ + RANSAC version of weighted Homography estimation. + Sample 4 points from the top 50% weighted points each time. + reproj_threshold: points with reprojection error less than this value are inliers + Returns: best_H + """ + if random_seed is not None: + torch.manual_seed(random_seed) + N = src_pts.shape[0] + assert N >= 4 + # 1. Select top 50% weighted points + sorted_idx = torch.argsort(confident_weight, descending=True) + n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) + candidate_idx = sorted_idx[:n_sample] + best_inlier_mask = None + best_score = 0 + for _ in range(n_iter): + # 2. Randomly sample 4 points + idx = candidate_idx[torch.randperm(n_sample)[:num_sample_for_ransac]] + # 3. Compute Homography + try: + H = find_homography_least_squares_weighted_torch( + src_pts[idx], dst_pts[idx], confident_weight[idx] + ) + except Exception: + H = torch.eye(3, dtype=src_pts.dtype, device=src_pts.device) + # 4. Compute reprojection error for all points + src_homo = torch.cat( + [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 + ) + proj = (H @ src_homo.T).T + proj = proj[:, :2] / proj[:, 2:3] + error = ((proj - dst_pts) ** 2).sum(dim=1).sqrt() # Euclidean distance + inlier_mask = error < reproj_threshold + total_score = (inlier_mask * confident_weight).sum().item() + n_inlier = inlier_mask.sum().item() + if n_inlier < 4: + continue # At least 4 inliers required for fitting + + if total_score > best_score: + best_score = total_score + best_inlier_mask = inlier_mask + + # 5. Refit Homography using inliers + H_inlier = find_homography_least_squares_weighted_torch( + src_pts[best_inlier_mask], dst_pts[best_inlier_mask], confident_weight[best_inlier_mask] + ) + + return H_inlier + + +def find_homography_least_squares_weighted_torch_batch( + src_pts_batch, dst_pts_batch, confident_weight_batch +): + """ + Batch version of weighted least squares Homography + src_pts_batch: (B, K, 2) + dst_pts_batch: (B, K, 2) + confident_weight_batch: (B, K) + Returns: (B, 3, 3) + """ + B, K, _ = src_pts_batch.shape + w = confident_weight_batch.sqrt().unsqueeze(2) # (B,K,1) + x = src_pts_batch[:, :, 0:1] + y = src_pts_batch[:, :, 1:2] + u = dst_pts_batch[:, :, 0:1] + v = dst_pts_batch[:, :, 1:2] + zeros = torch.zeros_like(x) + A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=2) + A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=2) + A = torch.cat([A1, A2], dim=1) # (B, 2K, 9) + # SVD: torch.linalg.svd supports batch + _, _, Vh = torch.linalg.svd(A) + H = Vh[:, -1].reshape(B, 3, 3) + H = H / H[:, 2:3, 2:3] + return H + + +def ransac_find_homography_weighted_fast( + src_pts, + dst_pts, + confident_weight, + n_sample, + n_iter=100, + reproj_threshold=3.0, + num_sample_for_ransac=8, + random_seed=None, + rand_sample_iters_idx=None, +): + """ + Batch version of RANSAC weighted Homography estimation. + Returns: H_inlier + """ + if random_seed is not None: + torch.manual_seed(random_seed) + N = src_pts.shape[0] + device = src_pts.device + assert N >= 4 + # 1. Select top weighted points by sample_ratio + sorted_idx = torch.argsort(confident_weight, descending=True) + candidate_idx = sorted_idx[:n_sample] # (n_sample,) + if rand_sample_iters_idx is None: + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + # 2. Generate all sampling groups at once + # shape: (n_iter, num_sample_for_ransac) + rand_idx = candidate_idx[rand_sample_iters_idx] # (n_iter, num_sample_for_ransac) + # 3. Construct batch input + src_pts_batch = src_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) + dst_pts_batch = dst_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) + confident_weight_batch = confident_weight[rand_idx] # (n_iter, num_sample_for_ransac) + # 4. Batch fit Homography + H_batch = find_homography_least_squares_weighted_torch_batch( + src_pts_batch, dst_pts_batch, confident_weight_batch + ) # (n_iter, 3, 3) + # 5. Batch evaluate inliers for all H + src_homo = torch.cat( + [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 + ) # (N,3) + src_homo_expand = src_homo.unsqueeze(0).expand(n_iter, N, 3) # (n_iter, N, 3) + dst_pts_expand = dst_pts.unsqueeze(0).expand(n_iter, N, 2) # (n_iter, N, 2) + confident_weight_expand = confident_weight.unsqueeze(0).expand(n_iter, N) # (n_iter, N) + # H_batch: (n_iter, 3, 3) + proj = torch.bmm(src_homo_expand, H_batch.transpose(1, 2)) # (n_iter, N, 3) + proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (n_iter, N, 2) + error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=2).sqrt() # (n_iter, N) + inlier_mask = error < reproj_threshold # (n_iter, N) + total_score = (inlier_mask * confident_weight_expand).sum(dim=1) # (n_iter,) + # 6. Select the sampling group with the highest score + best_idx = torch.argmax(total_score) + best_inlier_mask = inlier_mask[best_idx] # (N,) + inlier_src_pts = src_pts[best_inlier_mask] + inlier_dst_pts = dst_pts[best_inlier_mask] + inlier_confident_weight = confident_weight[best_inlier_mask] + + max_inlier_num = 10000 + sorted_idx = torch.argsort(inlier_confident_weight, descending=True) + + # method 1: sort according to confident_weight, and only keep max_inlier_num pts + # sorted_idx = sorted_idx[:max_inlier_num] + + # method 2: random choose max_inlier_num pts + sorted_idx = sorted_idx[torch.randperm(len(sorted_idx))[:max_inlier_num]] + + inlier_src_pts = inlier_src_pts[sorted_idx] + inlier_dst_pts = inlier_dst_pts[sorted_idx] + inlier_confident_weight = inlier_confident_weight[sorted_idx] + # 7. Refit Homography using inliers + H_inlier = find_homography_least_squares_weighted_torch( + inlier_src_pts, inlier_dst_pts, inlier_confident_weight + ) + return H_inlier + + +def ransac_find_homography_weighted_fast_batch( + src_pts, # (B, N, 3) + dst_pts, # (B, N, 2) + confident_weight, # (B, N) + n_sample, + n_iter=100, + reproj_threshold=3.0, + num_sample_for_ransac=8, + max_inlier_num=10000, + random_seed=None, + rand_sample_iters_idx=None, +): + """ + Batch version of RANSAC weighted Homography estimation (supports batch). + Input: + src_pts: (B, N, 2) + dst_pts: (B, N, 2) + confident_weight: (B, N) + Returns: + H_inlier: (B, 3, 3) + """ + if random_seed is not None: + torch.manual_seed(random_seed) + B, N, _ = src_pts.shape + assert N >= 4 + + device = src_pts.device + + # 1. Select top weighted points by sample_ratio + sorted_idx = torch.argsort(confident_weight, descending=True, dim=1) # (B, N) + candidate_idx = sorted_idx[:, :n_sample] # (B, n_sample) + + # 2. Generate all sampling groups at once + # rand_idx: (B, n_iter, num_sample_for_ransac) + if rand_sample_iters_idx is None: + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + + rand_idx = candidate_idx[:, rand_sample_iters_idx] # (B, n_iter, num_sample_for_ransac) + + # 3. Construct batch input + # Indexing method below: (B, n_iter, num_sample_for_ransac, ...) + b_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, n_iter, num_sample_for_ransac) + src_pts_batch = src_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) + dst_pts_batch = dst_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) + confident_weight_batch = confident_weight[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac) + + # 4. Batch fit Homography + # Need to implement batch version that supports (B, n_iter, num_sample_for_ransac, ...) input + # Output H_batch: (B, n_iter, 3, 3) + cB, cN = src_pts_batch.shape[:2] + H_batch = find_homography_least_squares_weighted_torch_batch( + src_pts_batch.flatten(0, 1), dst_pts_batch.flatten(0, 1), confident_weight_batch.flatten(0, 1) + ) # (B, n_iter, 3, 3) + H_batch = H_batch.unflatten(0, (cB, cN)) + + # 5. Batch evaluate inliers for all H + src_homo = torch.cat( + [src_pts, torch.ones(B, N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=2 + ) # (B, N, 3) + src_homo_expand = src_homo.unsqueeze(1).expand(B, n_iter, N, 3) # (B, n_iter, N, 3) + dst_pts_expand = dst_pts.unsqueeze(1).expand(B, n_iter, N, 2) # (B, n_iter, N, 2) + confident_weight_expand = confident_weight.unsqueeze(1).expand(B, n_iter, N) # (B, n_iter, N) + + # H_batch: (B, n_iter, 3, 3) + # Need to reshape H_batch to (B*n_iter, 3, 3), src_homo_expand to (B*n_iter, N, 3) + H_batch_flat = H_batch.reshape(-1, 3, 3) + src_homo_expand_flat = src_homo_expand.reshape(-1, N, 3) + proj = torch.bmm(src_homo_expand_flat, H_batch_flat.transpose(1, 2)) # (B*n_iter, N, 3) + proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (B*n_iter, N, 2) + proj_xy = proj_xy.reshape(B, n_iter, N, 2) + error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=3).sqrt() # (B, n_iter, N) + inlier_mask = error < reproj_threshold # (B, n_iter, N) + total_score = (inlier_mask * confident_weight_expand).sum(dim=2) # (B, n_iter) + + # 6. Select the sampling group with the highest score + best_idx = torch.argmax(total_score, dim=1) # (B,) + best_inlier_mask = inlier_mask[torch.arange(B, device=device), best_idx] # (B, N) + + # 7. Refit Homography using inliers + H_inlier_list = [] + for b in range(B): + mask = best_inlier_mask[b] + inlier_src_pts = src_pts[b][mask] # (?, 3) + inlier_dst_pts = dst_pts[b][mask] # (?, 2) + inlier_confident_weight = confident_weight[b][mask] # (?) + + sorted_idx = torch.argsort(inlier_confident_weight, descending=True) + # # method 1: sort according to confident_weight, and only keep max_inlier_num pts + # sorted_idx = sorted_idx[:max_inlier_num] + # method 2: random choose max_inlier_num pts + if len(sorted_idx) > max_inlier_num: + # random choose from first 95% confident pts + keep_len = max(int(len(sorted_idx) * 0.95), max_inlier_num) + sorted_idx = sorted_idx[:keep_len] + perm = torch.randperm(len(sorted_idx), device=device)[:max_inlier_num] + sorted_idx = sorted_idx[perm] + inlier_src_pts = inlier_src_pts[sorted_idx] + inlier_dst_pts = inlier_dst_pts[sorted_idx] + inlier_confident_weight = inlier_confident_weight[sorted_idx] + + H_inlier = find_homography_least_squares_weighted_torch( + inlier_src_pts, inlier_dst_pts, inlier_confident_weight + ) # (3, 3) + H_inlier_list.append(H_inlier) + H_inlier = torch.stack(H_inlier_list, dim=0) # (B, 3, 3) + return H_inlier + +def get_params_for_ransac(N, device): + n_iter=100 + sample_ratio=0.3 + num_sample_for_ransac=8 + n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + return n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx + + +def camray_to_caminfo(camray, confidence=None, reproj_threshold=0.2, training=False): + """ + Args: + camray: (B, S, num_patches_y, num_patches_x, 6) + confidence: (B, S, num_patches_y, num_patches_x) + Returns: + R: (B, S, 3, 3) + T: (B, S, 3) + focal_lengths: (B, S, 2) + principal_points: (B, S, 2) + """ + if confidence is None: + confidence = torch.ones_like(camray[:, :, :, :, 0]) + B, S, num_patches_y, num_patches_x, _ = camray.shape + # identity K, assume imw=imh=2.0 + I_K = torch.eye(3, dtype=camray.dtype, device=camray.device) + I_K[0, 2] = 1.0 + I_K[1, 2] = 1.0 + # repeat I_K to match camray + I_K = I_K.unsqueeze(0).unsqueeze(0).expand(B, S, -1, -1) + + cam_plane_depth = torch.ones( + B, S, num_patches_y, num_patches_x, 1, dtype=camray.dtype, device=camray.device + ) + I_cam_plane_unproj = unproject_depth( + cam_plane_depth, + I_K, + c2w=None, + ixt_normalized=True, + num_patches_x=num_patches_x, + num_patches_y=num_patches_y, + ) # (B, S, num_patches_y, num_patches_x, 3) + + camray = camray.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x, 6) + I_cam_plane_unproj = I_cam_plane_unproj.flatten(0, 1).flatten( + 1, 2 + ) # (B*S, num_patches_y*num_patches_x, 3) + confidence = confidence.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x) + + # Compute optimal rotation to align rays + N = camray.shape[-2] + device = camray.device + n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx = get_params_for_ransac(N, device) + + # Use batch processing (confidence is guaranteed to be not None at this point) + if training: + camray = camray.clone().detach() + I_cam_plane_unproj = I_cam_plane_unproj.clone().detach() + confidence = confidence.clone().detach() + R, focal_lengths, principal_points = compute_optimal_rotation_intrinsics_batch( + I_cam_plane_unproj, + camray[:, :, :3], + reproj_threshold=reproj_threshold, + weights=confidence, + n_sample = n_sample, + n_iter=n_iter, + num_sample_for_ransac=num_sample_for_ransac, + rand_sample_iters_idx=rand_sample_iters_idx, + ) + + T = torch.sum(camray[:, :, 3:] * confidence.unsqueeze(-1), dim=1) / torch.sum( + confidence, dim=-1, keepdim=True + ) + + R = R.reshape(B, S, 3, 3) + T = T.reshape(B, S, 3) + focal_lengths = focal_lengths.reshape(B, S, 2) + principal_points = principal_points.reshape(B, S, 2) + + return R, T, 1.0 / focal_lengths, principal_points + 1.0 + +def get_extrinsic_from_camray(camray, conf, patch_size_y, patch_size_x, training=False): + pred_R, pred_T, pred_focal_lengths, pred_principal_points = camray_to_caminfo( + camray, confidence=conf.squeeze(-1), training=training + ) + + pred_extrinsic = torch.cat( + [ + torch.cat([pred_R, pred_T.unsqueeze(-1)], dim=-1), + repeat( + torch.tensor([0, 0, 0, 1], dtype=pred_R.dtype, device=pred_R.device), + "c -> b s 1 c", + b=pred_R.shape[0], + s=pred_R.shape[1], + ), + ], + dim=-2, + ) # B, S, 4, 4 + return pred_extrinsic, pred_focal_lengths, pred_principal_points \ No newline at end of file diff --git a/depth_anything_3/utils/read_write_model.py b/depth_anything_3/utils/read_write_model.py new file mode 100644 index 0000000000000000000000000000000000000000..4b4cf197f609665e46f14e6b3e8b6657efc5660c --- /dev/null +++ b/depth_anything_3/utils/read_write_model.py @@ -0,0 +1,585 @@ +# Copyright (c), ETH Zurich and UNC Chapel Hill. +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# All rights reserved. +# +# This file has been modified by ByteDance Ltd. and/or its affiliates. on 11/05/2025 +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of +# its contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + + +import argparse +import collections +import os +import struct +import numpy as np + +CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"]) +Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) +BaseImage = collections.namedtuple( + "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"] +) +Point3D = collections.namedtuple( + "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"] +) + + +class Image(BaseImage): + def qvec2rotmat(self): + return qvec2rotmat(self.qvec) + + +CAMERA_MODELS = { + CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), + CameraModel(model_id=1, model_name="PINHOLE", num_params=4), + CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), + CameraModel(model_id=3, model_name="RADIAL", num_params=5), + CameraModel(model_id=4, model_name="OPENCV", num_params=8), + CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), + CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), + CameraModel(model_id=7, model_name="FOV", num_params=5), + CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), + CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), + CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), +} +CAMERA_MODEL_IDS = {camera_model.model_id: camera_model for camera_model in CAMERA_MODELS} +CAMERA_MODEL_NAMES = {camera_model.model_name: camera_model for camera_model in CAMERA_MODELS} + + +def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): + """Read and unpack the next bytes from a binary file. + :param fid: + :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + :param endian_character: Any of {@, =, <, >, !} + :return: Tuple of read and unpacked values. + """ + data = fid.read(num_bytes) + return struct.unpack(endian_character + format_char_sequence, data) + + +def write_next_bytes(fid, data, format_char_sequence, endian_character="<"): + """pack and write to a binary file. + :param fid: + :param data: data to send, if multiple elements are sent at the same time, + they should be encapsuled either in a list or a tuple + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + should be the same length as the data list or tuple + :param endian_character: Any of {@, =, <, >, !} + """ + if isinstance(data, (list, tuple)): + bytes = struct.pack(endian_character + format_char_sequence, *data) + else: + bytes = struct.pack(endian_character + format_char_sequence, data) + fid.write(bytes) + + +def read_cameras_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasText(const std::string& path) + void Reconstruction::ReadCamerasText(const std::string& path) + """ + cameras = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + camera_id = int(elems[0]) + model = elems[1] + width = int(elems[2]) + height = int(elems[3]) + params = np.array(tuple(map(float, elems[4:]))) + cameras[camera_id] = Camera( + id=camera_id, + model=model, + width=width, + height=height, + params=params, + ) + return cameras + + +def read_cameras_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + cameras = {} + with open(path_to_model_file, "rb") as fid: + num_cameras = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_cameras): + camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ") + camera_id = camera_properties[0] + model_id = camera_properties[1] + model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name + width = camera_properties[2] + height = camera_properties[3] + num_params = CAMERA_MODEL_IDS[model_id].num_params + params = read_next_bytes( + fid, + num_bytes=8 * num_params, + format_char_sequence="d" * num_params, + ) + cameras[camera_id] = Camera( + id=camera_id, + model=model_name, + width=width, + height=height, + params=np.array(params), + ) + assert len(cameras) == num_cameras + return cameras + + +def write_cameras_text(cameras, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasText(const std::string& path) + void Reconstruction::ReadCamerasText(const std::string& path) + """ + HEADER = ( + "# Camera list with one line of data per camera:\n" + + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + + f"# Number of cameras: {len(cameras)}\n" + ) + with open(path, "w") as fid: + fid.write(HEADER) + for _, cam in cameras.items(): + to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params] + line = " ".join([str(elem) for elem in to_write]) + fid.write(line + "\n") + + +def write_cameras_binary(cameras, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(cameras), "Q") + for _, cam in cameras.items(): + model_id = CAMERA_MODEL_NAMES[cam.model].model_id + camera_properties = [cam.id, model_id, cam.width, cam.height] + write_next_bytes(fid, camera_properties, "iiQQ") + for p in cam.params: + write_next_bytes(fid, float(p), "d") + return cameras + + +def read_images_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesText(const std::string& path) + void Reconstruction::WriteImagesText(const std::string& path) + """ + images = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + image_id = int(elems[0]) + qvec = np.array(tuple(map(float, elems[1:5]))) + tvec = np.array(tuple(map(float, elems[5:8]))) + camera_id = int(elems[8]) + image_name = elems[9] + elems = fid.readline().split() + xys = np.column_stack( + [ + tuple(map(float, elems[0::3])), + tuple(map(float, elems[1::3])), + ] + ) + point3D_ids = np.array(tuple(map(int, elems[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def read_images_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + images = {} + with open(path_to_model_file, "rb") as fid: + num_reg_images = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_reg_images): + binary_image_properties = read_next_bytes( + fid, num_bytes=64, format_char_sequence="idddddddi" + ) + image_id = binary_image_properties[0] + qvec = np.array(binary_image_properties[1:5]) + tvec = np.array(binary_image_properties[5:8]) + camera_id = binary_image_properties[8] + binary_image_name = b"" + current_char = read_next_bytes(fid, 1, "c")[0] + while current_char != b"\x00": # look for the ASCII 0 entry + binary_image_name += current_char + current_char = read_next_bytes(fid, 1, "c")[0] + image_name = binary_image_name.decode("utf-8") + num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + x_y_id_s = read_next_bytes( + fid, + num_bytes=24 * num_points2D, + format_char_sequence="ddq" * num_points2D, + ) + xys = np.column_stack( + [ + tuple(map(float, x_y_id_s[0::3])), + tuple(map(float, x_y_id_s[1::3])), + ] + ) + point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def write_images_text(images, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesText(const std::string& path) + void Reconstruction::WriteImagesText(const std::string& path) + """ + if len(images) == 0: + mean_observations = 0 + else: + mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images) + HEADER = ( + "# Image list with two lines of data per image:\n" + + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" + + "# Number of images: {}, mean observations per image: {}\n".format( + len(images), mean_observations + ) + ) + + with open(path, "w") as fid: + fid.write(HEADER) + for _, img in images.items(): + image_header = [ + img.id, + *img.qvec, + *img.tvec, + img.camera_id, + img.name, + ] + first_line = " ".join(map(str, image_header)) + fid.write(first_line + "\n") + + points_strings = [] + for xy, point3D_id in zip(img.xys, img.point3D_ids): + points_strings.append(" ".join(map(str, [*xy, point3D_id]))) + fid.write(" ".join(points_strings) + "\n") + + +def write_images_binary(images, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(images), "Q") + for _, img in images.items(): + write_next_bytes(fid, img.id, "i") + write_next_bytes(fid, img.qvec.tolist(), "dddd") + write_next_bytes(fid, img.tvec.tolist(), "ddd") + write_next_bytes(fid, img.camera_id, "i") + for char in img.name: + write_next_bytes(fid, char.encode("utf-8"), "c") + write_next_bytes(fid, b"\x00", "c") + write_next_bytes(fid, len(img.point3D_ids), "Q") + for xy, p3d_id in zip(img.xys, img.point3D_ids): + write_next_bytes(fid, [*xy, p3d_id], "ddq") + + +def read_points3D_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + points3D = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + point3D_id = int(elems[0]) + xyz = np.array(tuple(map(float, elems[1:4]))) + rgb = np.array(tuple(map(int, elems[4:7]))) + error = float(elems[7]) + image_ids = np.array(tuple(map(int, elems[8::2]))) + point2D_idxs = np.array(tuple(map(int, elems[9::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) + return points3D + + +def read_points3D_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + points3D = {} + with open(path_to_model_file, "rb") as fid: + num_points = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_points): + binary_point_line_properties = read_next_bytes( + fid, num_bytes=43, format_char_sequence="QdddBBBd" + ) + point3D_id = binary_point_line_properties[0] + xyz = np.array(binary_point_line_properties[1:4]) + rgb = np.array(binary_point_line_properties[4:7]) + error = np.array(binary_point_line_properties[7]) + track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + track_elems = read_next_bytes( + fid, + num_bytes=8 * track_length, + format_char_sequence="ii" * track_length, + ) + image_ids = np.array(tuple(map(int, track_elems[0::2]))) + point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) + return points3D + + +def write_points3D_text(points3D, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + if len(points3D) == 0: + mean_track_length = 0 + else: + mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D) + HEADER = ( + "# 3D point list with one line of data per point:\n" + + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" + + "# Number of points: {}, mean track length: {}\n".format( + len(points3D), mean_track_length + ) + ) + + with open(path, "w") as fid: + fid.write(HEADER) + for _, pt in points3D.items(): + point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error] + fid.write(" ".join(map(str, point_header)) + " ") + track_strings = [] + for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs): + track_strings.append(" ".join(map(str, [image_id, point2D]))) + fid.write(" ".join(track_strings) + "\n") + + +def write_points3D_binary(points3D, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(points3D), "Q") + for _, pt in points3D.items(): + write_next_bytes(fid, pt.id, "Q") + write_next_bytes(fid, pt.xyz.tolist(), "ddd") + write_next_bytes(fid, pt.rgb.tolist(), "BBB") + write_next_bytes(fid, pt.error, "d") + track_length = pt.image_ids.shape[0] + write_next_bytes(fid, track_length, "Q") + for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs): + write_next_bytes(fid, [image_id, point2D_id], "ii") + + +def detect_model_format(path, ext): + if ( + os.path.isfile(os.path.join(path, "cameras" + ext)) + and os.path.isfile(os.path.join(path, "images" + ext)) + and os.path.isfile(os.path.join(path, "points3D" + ext)) + ): + print("Detected model format: '" + ext + "'") + return True + + return False + + +def read_model(path, ext=""): + # try to detect the extension automatically + if ext == "": + if detect_model_format(path, ".bin"): + ext = ".bin" + elif detect_model_format(path, ".txt"): + ext = ".txt" + else: + print("Provide model format: '.bin' or '.txt'") + return + + if ext == ".txt": + cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) + images = read_images_text(os.path.join(path, "images" + ext)) + points3D = read_points3D_text(os.path.join(path, "points3D") + ext) + else: + cameras = read_cameras_binary(os.path.join(path, "cameras" + ext)) + images = read_images_binary(os.path.join(path, "images" + ext)) + points3D = read_points3D_binary(os.path.join(path, "points3D") + ext) + return cameras, images, points3D + + +def write_model(cameras, images, points3D, path, ext=".bin"): + if ext == ".txt": + write_cameras_text(cameras, os.path.join(path, "cameras" + ext)) + write_images_text(images, os.path.join(path, "images" + ext)) + write_points3D_text(points3D, os.path.join(path, "points3D") + ext) + else: + write_cameras_binary(cameras, os.path.join(path, "cameras" + ext)) + write_images_binary(images, os.path.join(path, "images" + ext)) + write_points3D_binary(points3D, os.path.join(path, "points3D") + ext) + return cameras, images, points3D + + +def qvec2rotmat(qvec): + return np.array( + [ + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + ], + [ + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + ], + [ + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ], + ] + ) + + +def rotmat2qvec(R): + Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat + K = ( + np.array( + [ + [Rxx - Ryy - Rzz, 0, 0, 0], + [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], + [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], + [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], + ] + ) + / 3.0 + ) + eigvals, eigvecs = np.linalg.eigh(K) + qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] + if qvec[0] < 0: + qvec *= -1 + return qvec + + +def main(): + parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models") + parser.add_argument("--input_model", help="path to input model folder") + parser.add_argument( + "--input_format", + choices=[".bin", ".txt"], + help="input model format", + default="", + ) + parser.add_argument("--output_model", help="path to output model folder") + parser.add_argument( + "--output_format", + choices=[".bin", ".txt"], + help="output model format", + default=".txt", + ) + args = parser.parse_args() + + cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format) + + print("num_cameras:", len(cameras)) + print("num_images:", len(images)) + print("num_points3D:", len(points3D)) + + if args.output_model is not None: + write_model( + cameras, + images, + points3D, + path=args.output_model, + ext=args.output_format, + ) + + +if __name__ == "__main__": + main() diff --git a/depth_anything_3/utils/registry.py b/depth_anything_3/utils/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..7db16d525e0417110d87aa5e621b792d8bd95596 --- /dev/null +++ b/depth_anything_3/utils/registry.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any +from addict import Dict + + +class Registry(Dict[str, Any]): + def __init__(self): + super().__init__() + self._map = Dict({}) + + def register(self, name=None): + def decorator(cls): + key = name or cls.__name__ + self._map[key] = cls + return cls + + return decorator + + def get(self, name): + return self._map[name] + + def all(self): + return self._map diff --git a/depth_anything_3/utils/sh_helpers.py b/depth_anything_3/utils/sh_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..1a1a4ca8204eb8d858351afb253e41e77dfa2f74 --- /dev/null +++ b/depth_anything_3/utils/sh_helpers.py @@ -0,0 +1,88 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from math import isqrt +import torch +from einops import einsum + +try: + from e3nn.o3 import matrix_to_angles, wigner_D +except ImportError: + from depth_anything_3.utils.logger import logger + + logger.warn("Dependency 'e3nn' not found. Required for rotating the camera space SH coeff") + + +def project_to_so3_strict(M: torch.Tensor) -> torch.Tensor: + if M.shape[-2:] != (3, 3): + raise ValueError("Input must be a batch of 3x3 matrices (i.e., shape [..., 3, 3]).") + + # 1. Compute SVD + U, S, Vh = torch.linalg.svd(M) + V = Vh.mH + + # 2. Handle reflection case (det = -1) + det_U = torch.det(U) + det_V = torch.det(V) + is_reflection = (det_U * det_V) < 0 + correction_sign = torch.where( + is_reflection[..., None], + torch.tensor([1, 1, -1.0], device=M.device, dtype=M.dtype), + torch.tensor([1, 1, 1.0], device=M.device, dtype=M.dtype), + ) + correction_matrix = torch.diag_embed(correction_sign) + U_corrected = U @ correction_matrix + R_so3_initial = U_corrected @ V.transpose(-2, -1) + + # 3. Explicitly ensure determinant is 1 (or extremely close) + current_det = torch.det(R_so3_initial) + det_correction_factor = torch.pow(current_det, -1 / 3)[..., None, None] + R_so3_final = R_so3_initial * det_correction_factor + + return R_so3_final + + +def rotate_sh( + sh_coefficients: torch.Tensor, # "*#batch n" + rotations: torch.Tensor, # "*#batch 3 3" +) -> torch.Tensor: # "*batch n" + # https://github.com/graphdeco-inria/gaussian-splatting/issues/176#issuecomment-2452412653 + device = sh_coefficients.device + dtype = sh_coefficients.dtype + + *_, n = sh_coefficients.shape + + with torch.autocast(device_type=rotations.device.type, enabled=False): + rotations_float32 = rotations.to(torch.float32) + + # switch axes: yzx -> xyz + P = torch.tensor([[0, 0, 1], [1, 0, 0], [0, 1, 0]]).unsqueeze(0).to(rotations_float32) + permuted_rotations = torch.linalg.inv(P) @ rotations_float32 @ P + + # ensure rotation has det == 1 in float32 type + permuted_rotations_so3 = project_to_so3_strict(permuted_rotations) + + alpha, beta, gamma = matrix_to_angles(permuted_rotations_so3) + result = [] + for degree in range(isqrt(n)): + with torch.device(device): + sh_rotations = wigner_D(degree, alpha, -beta, gamma).type(dtype) + sh_rotated = einsum( + sh_rotations, + sh_coefficients[..., degree**2 : (degree + 1) ** 2], + "... i j, ... j -> ... i", + ) + result.append(sh_rotated) + + return torch.cat(result, dim=-1) diff --git a/depth_anything_3/utils/visualize.py b/depth_anything_3/utils/visualize.py new file mode 100644 index 0000000000000000000000000000000000000000..8fd32bddf00e5461f674525e73653409270c0227 --- /dev/null +++ b/depth_anything_3/utils/visualize.py @@ -0,0 +1,120 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import matplotlib +import numpy as np +import torch +from einops import rearrange + +from depth_anything_3.utils.logger import logger + + +def visualize_depth( + depth: np.ndarray, + depth_min=None, + depth_max=None, + percentile=2, + ret_minmax=False, + ret_type=np.uint8, + cmap="Spectral", +): + """ + Visualize a depth map using a colormap. + + Args: + depth: Input depth map array + depth_min: Minimum depth value for normalization. If None, uses percentile + depth_max: Maximum depth value for normalization. If None, uses percentile + percentile: Percentile for min/max computation if not provided + ret_minmax: Whether to return min/max depth values + ret_type: Return array type (uint8 or float) + cmap: Matplotlib colormap name to use + + Returns: + Colored depth visualization as numpy array + If ret_minmax=True, also returns depth_min and depth_max + """ + depth = depth.copy() + depth.copy() + valid_mask = depth > 0 + depth[valid_mask] = 1 / depth[valid_mask] + if depth_min is None: + if valid_mask.sum() <= 10: + depth_min = 0 + else: + depth_min = np.percentile(depth[valid_mask], percentile) + if depth_max is None: + if valid_mask.sum() <= 10: + depth_max = 0 + else: + depth_max = np.percentile(depth[valid_mask], 100 - percentile) + if depth_min == depth_max: + depth_min = depth_min - 1e-6 + depth_max = depth_max + 1e-6 + cm = matplotlib.colormaps[cmap] + depth = ((depth - depth_min) / (depth_max - depth_min)).clip(0, 1) + depth = 1 - depth + img_colored_np = cm(depth[None], bytes=False)[:, :, :, 0:3] # value from 0 to 1 + if ret_type == np.uint8: + img_colored_np = (img_colored_np[0] * 255.0).astype(np.uint8) + elif ret_type == np.float32 or ret_type == np.float64: + img_colored_np = img_colored_np[0] + else: + raise ValueError(f"Invalid return type: {ret_type}") + if ret_minmax: + return img_colored_np, depth_min, depth_max + else: + return img_colored_np + + +# GS video rendering visulization function, since it operates in Tensor space... + + +def vis_depth_map_tensor( + result: torch.Tensor, # "*batch height width" + color_map: str = "Spectral", +) -> torch.Tensor: # "*batch 3 height with" + """ + Color-map the depth map. + """ + far = result.reshape(-1)[:16_000_000].float().quantile(0.99).log().to(result) + try: + near = result[result > 0][:16_000_000].float().quantile(0.01).log().to(result) + except (RuntimeError, ValueError) as e: + logger.error(f"No valid depth values found. Reason: {e}") + near = torch.zeros_like(far) + result = result.log() + result = (result - near) / (far - near) + return apply_color_map_to_image(result, color_map) + + +def apply_color_map( + x: torch.Tensor, # " *batch" + color_map: str = "inferno", +) -> torch.Tensor: # "*batch 3" + cmap = matplotlib.cm.get_cmap(color_map) + + # Convert to NumPy so that Matplotlib color maps can be used. + mapped = cmap(x.float().detach().clip(min=0, max=1).cpu().numpy())[..., :3] + + # Convert back to the original format. + return torch.tensor(mapped, device=x.device, dtype=torch.float32) + + +def apply_color_map_to_image( + image: torch.Tensor, # "*batch height width" + color_map: str = "inferno", +) -> torch.Tensor: # "*batch 3 height with" + image = apply_color_map(image, color_map) + return rearrange(image, "... h w c -> ... c h w") diff --git a/diffsynth/__init__.py b/diffsynth/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bb67a43fa4e5791ab58e7e40260bc3df8b6bc7cc --- /dev/null +++ b/diffsynth/__init__.py @@ -0,0 +1 @@ +from .core import * diff --git a/diffsynth/configs/__init__.py b/diffsynth/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..144a822978b12a8297e341760a74791fb12802c0 --- /dev/null +++ b/diffsynth/configs/__init__.py @@ -0,0 +1,2 @@ +from .model_configs import MODEL_CONFIGS +from .vram_management_module_maps import VRAM_MANAGEMENT_MODULE_MAPS diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py new file mode 100644 index 0000000000000000000000000000000000000000..93a338a2a48399c976c0b6f5bcf03890f37df515 --- /dev/null +++ b/diffsynth/configs/model_configs.py @@ -0,0 +1,506 @@ +qwen_image_series = [ + { + # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors") + "model_hash": "0319a1cb19835fb510907dd3367c95ff", + "model_name": "metaview_dit", + # "model_class": "diffsynth.models.qwen_image_dit.QwenImageDiT", + "model_class": "src.MetaView_dit.MetaViewDiT", + # "extra_kwargs": {'add_attn_type': 'prope_attn'}, + # "extra_kwargs": {'add_attn_type': 'cross_attn', 'add_in_dim': 3072}, + # "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'add_in_dim': 3072}, + # "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'add_in_dim': 4096}, + # "extra_kwargs": {'add_attn_type': 'self_attn_lora_3D', 'add_in_dim': 3072, 'lora_rank': 256}, + # "extra_kwargs": {'add_attn_type': 'self_attn_lora_3D', 'merge_3D': True, '_3d_dim': 6144, 'add_in_dim': 3072, 'lora_rank': 256}, + # "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'merge_3D': True, '_3d_dim': 3072, 'add_in_dim': 3072}, + # "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'merge_3D': True, '_3d_dim': 4096, 'add_in_dim': 3072}, + + "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'merge_3D': True, '_3d_dim': 6144, 'add_in_dim': 3072}, + # "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'merge_3D': True, 'decode_3D': True, '_3d_dim': 4096, 'add_in_dim': 3072}, + }, + { + # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors") + "model_hash": "8004730443f55db63092006dd9f7110e", + "model_name": "qwen_image_text_encoder", + "model_class": "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.qwen_image_text_encoder.QwenImageTextEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors") + "model_hash": "ed4ea5824d55ec3107b09815e318123a", + "model_name": "qwen_image_vae", + "model_class": "diffsynth.models.qwen_image_vae.QwenImageVAE", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth", origin_file_pattern="model.safetensors") + "model_hash": "073bce9cf969e317e5662cd570c3e79c", + "model_name": "qwen_image_blockwise_controlnet", + "model_class": "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint", origin_file_pattern="model.safetensors") + "model_hash": "a9e54e480a628f0b956a688a81c33bab", + "model_name": "qwen_image_blockwise_controlnet", + "model_class": "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet", + "extra_kwargs": {"additional_in_dim": 4}, + }, +] + +wan_series = [ + { + # Example: ModelConfig(model_id="krea/krea-realtime-video", origin_file_pattern="krea-realtime-video-14b.safetensors") + "model_hash": "5ec04e02b42d2580483ad69f4e76346a", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth") + "model_hash": "9c8818c2cbea55eca56c7b447df170da", + "model_name": "wan_video_text_encoder", + "model_class": "diffsynth.models.wan_video_text_encoder.WanTextEncoder", + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="Wan2.1_VAE.pth") + "model_hash": "ccc42284ea13e1ad04693284c7a09be6", + "model_name": "wan_video_vae", + "model_class": "diffsynth.models.wan_video_vae.WanVideoVAE", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vae.WanVideoVAEStateDictConverter", + }, + { + # Example: ModelConfig(model_id="meituan-longcat/LongCat-Video", origin_file_pattern="dit/diffusion_pytorch_model*.safetensors") + "model_hash": "8b27900f680d7251ce44e2dc8ae1ffef", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.longcat_video_dit.LongCatVideoTransformer3DModel", + }, + { + # Example: ModelConfig(model_id="ByteDance/Video-As-Prompt-Wan2.1-14B", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors") + "model_hash": "5f90e66a0672219f12d9a626c8c21f61", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTFromDiffusers" + }, + { + # Example: ModelConfig(model_id="ByteDance/Video-As-Prompt-Wan2.1-14B", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors") + "model_hash": "5f90e66a0672219f12d9a626c8c21f61", + "model_name": "wan_video_vap", + "model_class": "diffsynth.models.wan_video_mot.MotWanModel", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_mot.WanVideoMotStateDictConverter" + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-I2V-14B-480P", origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth") + "model_hash": "5941c53e207d62f20f9025686193c40b", + "model_name": "wan_video_image_encoder", + "model_class": "diffsynth.models.wan_video_image_encoder.WanImageEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_image_encoder.WanImageEncoderStateDictConverter" + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1", origin_file_pattern="model.safetensors") + "model_hash": "dbd5ec76bbf977983f972c151d545389", + "model_name": "wan_video_motion_controller", + "model_class": "diffsynth.models.wan_video_motion_controller.WanMotionControllerModel", + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "9269f8db9040a9d860eaca435be61814", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-FLF2V-14B-720P", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "3ef3b1f8e1dab83d5b71fd7b617f859f", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_image_pos_emb': True} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-1.3B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "349723183fc063b2bfc10bb2835cf677", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-1.3B-InP", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "6d6ccde6845b95ad9114ab993d917893", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-14B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "efa44cddf936c70abd0ea28b6cbe946c", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-14B-InP", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "6bfcfb3b342cb286ce886889d519a77e", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "ac6a5aa74f4a0aab6f64eb9a72f19901", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 32, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-1.3B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "70ddad9d3a133785da5ea371aae09504", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06, 'has_ref_conv': True} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-14B-Control-Camera", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "b61c605c2adbd23124d152ed28e049ae", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 32, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-14B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "26bde73488a92e64cc20b0a7485b9e5b", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': True} + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "aafcfd9672c3a2456dc46e1cb6e52c70", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06} + }, + { + # Example: ModelConfig(model_id="iic/VACE-Wan2.1-1.3B-Preview", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "a61453409b67cd3246cf0c3bebad47ba", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="iic/VACE-Wan2.1-1.3B-Preview", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "a61453409b67cd3246cf0c3bebad47ba", + "model_name": "wan_video_vace", + "model_class": "diffsynth.models.wan_video_vace.VaceWanModel", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vace.VaceWanModelDictConverter" + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-VACE-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "7a513e1f257a861512b1afd387a8ecd9", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.1-VACE-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "7a513e1f257a861512b1afd387a8ecd9", + "model_name": "wan_video_vace", + "model_class": "diffsynth.models.wan_video_vace.VaceWanModel", + "extra_kwargs": {'vace_layers': (0, 5, 10, 15, 20, 25, 30, 35), 'vace_in_dim': 96, 'patch_size': (1, 2, 2), 'has_image_input': False, 'dim': 5120, 'num_heads': 40, 'ffn_dim': 13824, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vace.VaceWanModelDictConverter" + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-Animate-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "31fa352acb8a1b1d33cd8764273d80a2", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter" + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-Animate-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "31fa352acb8a1b1d33cd8764273d80a2", + "model_name": "wan_video_animate_adapter", + "model_class": "diffsynth.models.wan_video_animate_adapter.WanAnimateAdapter", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_animate_adapter.WanAnimateAdapterStateDictConverter" + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.2-Fun-A14B-Control-Camera", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors") + "model_hash": "47dbeab5e560db3180adf51dc0232fb1", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24, 'require_clip_embedding': False} + }, + { + # Example: ModelConfig(model_id="PAI/Wan2.2-Fun-A14B-Control", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors") + "model_hash": "2267d489f0ceb9f21836532952852ee5", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 52, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': True, 'require_clip_embedding': False}, + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-I2V-A14B", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors") + "model_hash": "5b013604280dd715f8457c6ed6d6a626", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'require_clip_embedding': False} + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-S2V-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "966cffdcc52f9c46c391768b27637614", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit_s2v.WanS2VModel", + "extra_kwargs": {'dim': 5120, 'in_dim': 16, 'ffn_dim': 13824, 'out_dim': 16, 'text_dim': 4096, 'freq_dim': 256, 'eps': 1e-06, 'patch_size': (1, 2, 2), 'num_heads': 40, 'num_layers': 40, 'cond_dim': 16, 'audio_dim': 1024, 'num_audio_token': 4} + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="diffusion_pytorch_model*.safetensors") + "model_hash": "1f5ab7703c6fc803fdded85ff040c316", + "model_name": "wan_video_dit", + "model_class": "diffsynth.models.wan_video_dit.WanModel", + "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 3072, 'ffn_dim': 14336, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 48, 'num_heads': 24, 'num_layers': 30, 'eps': 1e-06, 'seperated_timestep': True, 'require_clip_embedding': False, 'require_vae_embedding': False, 'fuse_vae_embedding_in_latents': True} + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="Wan2.2_VAE.pth") + "model_hash": "e1de6c02cdac79f8b739f4d3698cd216", + "model_name": "wan_video_vae", + "model_class": "diffsynth.models.wan_video_vae.WanVideoVAE38", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vae.WanVideoVAEStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Wan-AI/Wan2.2-S2V-14B", origin_file_pattern="wav2vec2-large-xlsr-53-english/model.safetensors") + "model_hash": "06be60f3a4526586d8431cd038a71486", + "model_name": "wans2v_audio_encoder", + "model_class": "diffsynth.models.wav2vec.WanS2VAudioEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.wans2v_audio_encoder.WanS2VAudioEncoderStateDictConverter", + }, +] + +flux_series = [ + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors") + "model_hash": "a29710fea6dddb0314663ee823598e50", + "model_name": "flux_dit", + "model_class": "src.flux_kontext_dit.FluxDiTKontext", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + "extra_kwargs": {'add_attn_type': 'self_attn_3D', 'merge_3D': True, '_3d_dim': 6144, 'add_in_dim': 3072}, + }, + # { + # # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors") + # "model_hash": "a29710fea6dddb0314663ee823598e50", + # "model_name": "flux_dit", + # "model_class": "diffsynth.models.flux_dit.FluxDiT", + # "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + # }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors") + "model_hash": "94eefa3dac9cec93cb1ebaf1747d7b78", + "model_name": "flux_text_encoder_clip", + "model_class": "diffsynth.models.flux_text_encoder_clip.FluxTextEncoderClip", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_text_encoder_clip.FluxTextEncoderClipStateDictConverter", + }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors") + "model_hash": "22540b49eaedbc2f2784b2091a234c7c", + "model_name": "flux_text_encoder_t5", + "model_class": "diffsynth.models.flux_text_encoder_t5.FluxTextEncoderT5", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_text_encoder_t5.FluxTextEncoderT5StateDictConverter", + }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors") + "model_hash": "21ea55f476dfc4fd135587abb59dfe5d", + "model_name": "flux_vae_encoder", + "model_class": "diffsynth.models.flux_vae.FluxVAEEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors") + "model_hash": "21ea55f476dfc4fd135587abb59dfe5d", + "model_name": "flux_vae_decoder", + "model_class": "diffsynth.models.flux_vae.FluxVAEDecoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEDecoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="ostris/Flex.2-preview", origin_file_pattern="Flex.2-preview.safetensors") + "model_hash": "d02f41c13549fa5093d3521f62a5570a", + "model_name": "flux_dit", + "model_class": "diffsynth.models.flux_dit.FluxDiT", + "extra_kwargs": {'input_dim': 196, 'num_blocks': 8}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors") + "model_hash": "0629116fce1472503a66992f96f3eb1a", + "model_name": "flux_value_controller", + "model_class": "diffsynth.models.flux_value_control.SingleValueEncoder", + }, + { + # Example: ModelConfig(model_id="alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", origin_file_pattern="diffusion_pytorch_model.safetensors") + "model_hash": "52357cb26250681367488a8954c271e8", + "model_name": "flux_controlnet", + "model_class": "diffsynth.models.flux_controlnet.FluxControlNet", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter", + "extra_kwargs": {"num_joint_blocks": 6, "num_single_blocks": 0, "additional_input_dim": 4}, + }, + { + # Example: ModelConfig(model_id="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors") + "model_hash": "78d18b9101345ff695f312e7e62538c0", + "model_name": "flux_controlnet", + "model_class": "diffsynth.models.flux_controlnet.FluxControlNet", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter", + "extra_kwargs": {"num_mode": 10, "mode_dict": {"canny": 0, "tile": 1, "depth": 2, "blur": 3, "pose": 4, "gray": 5, "lq": 6}}, + }, + { + # Example: ModelConfig(model_id="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors") + "model_hash": "b001c89139b5f053c715fe772362dd2a", + "model_name": "flux_controlnet", + "model_class": "diffsynth.models.flux_controlnet.FluxControlNet", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter", + "extra_kwargs": {"num_single_blocks": 0}, + }, + { + # Example: ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/image_proj_model.bin") + "model_hash": "c07c0f04f5ff55e86b4e937c7a40d481", + "model_name": "infiniteyou_image_projector", + "model_class": "diffsynth.models.flux_infiniteyou.InfiniteYouImageProjector", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_infiniteyou.FluxInfiniteYouImageProjectorStateDictConverter", + }, + { + # Example: ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors") + "model_hash": "7f9583eb8ba86642abb9a21a4b2c9e16", + "model_name": "flux_controlnet", + "model_class": "diffsynth.models.flux_controlnet.FluxControlNet", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter", + "extra_kwargs": {"num_joint_blocks": 4, "num_single_blocks": 10}, + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev", origin_file_pattern="model.safetensors") + "model_hash": "77c2e4dd2440269eb33bfaa0d004f6ab", + "model_name": "flux_lora_encoder", + "model_class": "diffsynth.models.flux_lora_encoder.FluxLoRAEncoder", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev", origin_file_pattern="model.safetensors") + "model_hash": "30143afb2dea73d1ac580e0787628f8c", + "model_name": "flux_lora_patcher", + "model_class": "diffsynth.models.flux_lora_patcher.FluxLoraPatcher", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors") + "model_hash": "2bd19e845116e4f875a0a048e27fc219", + "model_name": "nexus_gen_llm", + "model_class": "diffsynth.models.nexus_gen.NexusGenAutoregressiveModel", + "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen.NexusGenAutoregressiveModelStateDictConverter", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin") + "model_hash": "63c969fd37cce769a90aa781fbff5f81", + "model_name": "nexus_gen_editing_adapter", + "model_class": "diffsynth.models.nexus_gen_projector.NexusGenImageEmbeddingMerger", + "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen_projector.NexusGenMergerStateDictConverter", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin") + "model_hash": "63c969fd37cce769a90aa781fbff5f81", + "model_name": "flux_dit", + "model_class": "diffsynth.models.flux_dit.FluxDiT", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin") + "model_hash": "3e6c61b0f9471135fc9c6d6a98e98b6d", + "model_name": "nexus_gen_generation_adapter", + "model_class": "diffsynth.models.nexus_gen_projector.NexusGenAdapter", + "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen_projector.NexusGenAdapterStateDictConverter", + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin") + "model_hash": "3e6c61b0f9471135fc9c6d6a98e98b6d", + "model_name": "flux_dit", + "model_class": "diffsynth.models.flux_dit.FluxDiT", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + }, + { + # Example: ModelConfig(model_id="InstantX/FLUX.1-dev-IP-Adapter", origin_file_pattern="ip-adapter.bin") + "model_hash": "4daaa66cc656a8fe369908693dad0a35", + "model_name": "flux_ipadapter", + "model_class": "diffsynth.models.flux_ipadapter.FluxIpAdapter", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.FluxIpAdapterStateDictConverter", + }, + { + # Example: ModelConfig(model_id="google/siglip-so400m-patch14-384", origin_file_pattern="model.safetensors") + "model_hash": "04d8c1e20a1f1b25f7434f111992a33f", + "model_name": "siglip_vision_model", + "model_class": "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.SiglipStateDictConverter", + }, + { + # Example: ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"), + "model_hash": "d30fb9e02b1dbf4e509142f05cf7dd50", + "model_name": "step1x_connector", + "model_class": "diffsynth.models.step1x_connector.Qwen2Connector", + "state_dict_converter": "diffsynth.utils.state_dict_converters.step1x_connector.Qwen2ConnectorStateDictConverter", + }, + { + # Example: ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"), + "model_hash": "d30fb9e02b1dbf4e509142f05cf7dd50", + "model_name": "flux_dit", + "model_class": "diffsynth.models.flux_dit.FluxDiT", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter", + "extra_kwargs": {"disable_guidance_embedder": True}, + }, +] + +flux2_series = [ + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="text_encoder/*.safetensors") + "model_hash": "28fca3d8e5bf2a2d1271748a773f6757", + "model_name": "flux2_text_encoder", + "model_class": "diffsynth.models.flux2_text_encoder.Flux2TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux2_text_encoder.Flux2TextEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="transformer/*.safetensors") + "model_hash": "d38e1d5c5aec3b0a11e79327ac6e3b0f", + "model_name": "flux2_dit", + "model_class": "diffsynth.models.flux2_dit.Flux2DiT", + }, + { + # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="vae/diffusion_pytorch_model.safetensors") + "model_hash": "c54288e3ee12ca215898840682337b95", + "model_name": "flux2_vae", + "model_class": "diffsynth.models.flux2_vae.Flux2VAE", + }, +] + +z_image_series = [ + { + # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="transformer/*.safetensors") + "model_hash": "fc3a8a1247fe185ce116ccbe0e426c28", + "model_name": "z_image_dit", + "model_class": "diffsynth.models.z_image_dit.ZImageDiT", + }, + { + # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors") + "model_hash": "0f050f62a88876fea6eae0a18dac5a2e", + "model_name": "z_image_text_encoder", + "model_class": "diffsynth.models.z_image_text_encoder.ZImageTextEncoder", + }, + { + # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/vae/diffusion_pytorch_model.safetensors") + "model_hash": "1aafa3cc91716fb6b300cc1cd51b85a3", + "model_name": "flux_vae_encoder", + "model_class": "diffsynth.models.flux_vae.FluxVAEEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEEncoderStateDictConverterDiffusers", + "extra_kwargs": {"use_conv_attention": False}, + }, + { + # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/vae/diffusion_pytorch_model.safetensors") + "model_hash": "1aafa3cc91716fb6b300cc1cd51b85a3", + "model_name": "flux_vae_decoder", + "model_class": "diffsynth.models.flux_vae.FluxVAEDecoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEDecoderStateDictConverterDiffusers", + "extra_kwargs": {"use_conv_attention": False}, + }, +] + +MODEL_CONFIGS = qwen_image_series + wan_series + flux_series + flux2_series + z_image_series diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py new file mode 100644 index 0000000000000000000000000000000000000000..219983d0e68e3f448656f86b9c6ba6f643721819 --- /dev/null +++ b/diffsynth/configs/vram_management_module_maps.py @@ -0,0 +1,178 @@ +flux_general_vram_config = { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.general_modules.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.flux_lora_encoder.LoRALayerBlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.flux_lora_patcher.LoraMerger": "diffsynth.core.vram.layers.AutoWrappedModule", +} + +VRAM_MANAGEMENT_MODULE_MAPS = { + "diffsynth.models.qwen_image_dit.QwenImageDiT": { + "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + }, + "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VisionPatchEmbed": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VisionRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.qwen_image_vae.QwenImageVAE": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.qwen_image_vae.QwenImageRMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.qwen_image_controlnet.BlockWiseControlBlock": { + "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + }, + "diffsynth.models.wan_video_animate_adapter.WanAnimateAdapter": { + "diffsynth.models.wan_video_animate_adapter.FaceEncoder": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_animate_adapter.EqualLinear": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_animate_adapter.ConvLayer": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_animate_adapter.FusedLeakyReLU": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_animate_adapter.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_dit_s2v.WanS2VModel": { + "diffsynth.models.wan_video_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit_s2v.WanS2VDiTBlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit_s2v.CausalAudioEncoder": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_dit.WanModel": { + "diffsynth.models.wan_video_dit.MLP": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit.DiTBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule", + "diffsynth.models.wan_video_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_image_encoder.WanImageEncoder": { + "diffsynth.models.wan_video_image_encoder.VisionTransformer": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_mot.MotWanModel": { + "diffsynth.models.wan_video_mot.MotWanAttentionBlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_motion_controller.WanMotionControllerModel": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + }, + "diffsynth.models.wan_video_text_encoder.WanTextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_text_encoder.T5RelativeEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_text_encoder.T5LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_vace.VaceWanModel": { + "diffsynth.models.wan_video_dit.DiTBlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_vae.WanVideoVAE": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.RMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.CausalConv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.Upsample": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wan_video_vae.WanVideoVAE38": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.RMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.CausalConv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.wan_video_vae.Upsample": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.wav2vec.WanS2VAudioEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.longcat_video_dit.LongCatVideoTransformer3DModel": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.longcat_video_dit.RMSNorm_FP32": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.longcat_video_dit.LayerNorm_FP32": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux_dit.FluxDiT": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "diffsynth.models.flux_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux_text_encoder_clip.FluxTextEncoderClip": flux_general_vram_config, + "diffsynth.models.flux_vae.FluxVAEEncoder": flux_general_vram_config, + "diffsynth.models.flux_vae.FluxVAEDecoder": flux_general_vram_config, + "diffsynth.models.flux_controlnet.FluxControlNet": flux_general_vram_config, + "diffsynth.models.flux_infiniteyou.InfiniteYouImageProjector": flux_general_vram_config, + "diffsynth.models.flux_ipadapter.FluxIpAdapter": flux_general_vram_config, + "diffsynth.models.flux_lora_patcher.FluxLoraPatcher": flux_general_vram_config, + "diffsynth.models.step1x_connector.Qwen2Connector": flux_general_vram_config, + "diffsynth.models.flux_lora_encoder.FluxLoRAEncoder": flux_general_vram_config, + "diffsynth.models.flux_text_encoder_t5.FluxTextEncoderT5": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.t5.modeling_t5.T5LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.t5.modeling_t5.T5DenseActDense": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.t5.modeling_t5.T5DenseGatedActDense": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M": { + "transformers.models.siglip.modeling_siglip.SiglipVisionEmbeddings": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.siglip.modeling_siglip.SiglipEncoder": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.siglip.modeling_siglip.SiglipMultiheadAttentionPoolingHead": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.MultiheadAttention": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux2_dit.Flux2DiT": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux2_text_encoder.Flux2TextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.mistral.modeling_mistral.MistralRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.flux2_vae.Flux2VAE": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.z_image_text_encoder.ZImageTextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.z_image_dit.ZImageDiT": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "diffsynth.models.z_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, +} diff --git a/diffsynth/core/__init__.py b/diffsynth/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72e501fc500321f0710c3f507930d301ab847033 --- /dev/null +++ b/diffsynth/core/__init__.py @@ -0,0 +1,5 @@ +from .attention import * +from .data import * +from .gradient import * +from .loader import * +from .vram import * diff --git a/diffsynth/core/attention/__init__.py b/diffsynth/core/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..45cf8a4382397aa4ff6558f37191c726d821b95a --- /dev/null +++ b/diffsynth/core/attention/__init__.py @@ -0,0 +1 @@ +from .attention import attention_forward diff --git a/diffsynth/core/attention/attention.py b/diffsynth/core/attention/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..15b55a43aa0ee248245fd6652f731f966091b6f7 --- /dev/null +++ b/diffsynth/core/attention/attention.py @@ -0,0 +1,121 @@ +import torch, os +from einops import rearrange + + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + +try: + import xformers.ops as xops + XFORMERS_AVAILABLE = True +except ModuleNotFoundError: + XFORMERS_AVAILABLE = False + + +def initialize_attention_priority(): + if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None: + return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').lower() + elif FLASH_ATTN_3_AVAILABLE: + return "flash_attention_3" + elif FLASH_ATTN_2_AVAILABLE: + return "flash_attention_2" + elif SAGE_ATTN_AVAILABLE: + return "sage_attention" + elif XFORMERS_AVAILABLE: + return "xformers" + else: + return "torch" + + +ATTENTION_IMPLEMENTATION = initialize_attention_priority() + + +def rearrange_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", required_in_pattern="b n s d", dims=None): + dims = {} if dims is None else dims + if q_pattern != required_in_pattern: + q = rearrange(q, f"{q_pattern} -> {required_in_pattern}", **dims) + if k_pattern != required_in_pattern: + k = rearrange(k, f"{k_pattern} -> {required_in_pattern}", **dims) + if v_pattern != required_in_pattern: + v = rearrange(v, f"{q_pattern} -> {required_in_pattern}", **dims) + return q, k, v + + +def rearrange_out(out: torch.Tensor, out_pattern="b n s d", required_out_pattern="b n s d", dims=None): + dims = {} if dims is None else dims + if out_pattern != required_out_pattern: + out = rearrange(out, f"{required_out_pattern} -> {out_pattern}", **dims) + return out + + +def torch_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None): + required_in_pattern, required_out_pattern= "b n s d", "b n s d" + q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims) + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale) + out = rearrange_out(out, out_pattern, required_out_pattern, dims) + return out + + +def flash_attention_3(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None): + required_in_pattern, required_out_pattern= "b s n d", "b s n d" + q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims) + out = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=scale) + if isinstance(out, tuple): + out = out[0] + out = rearrange_out(out, out_pattern, required_out_pattern, dims) + return out + + +def flash_attention_2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None): + required_in_pattern, required_out_pattern= "b s n d", "b s n d" + q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims) + out = flash_attn.flash_attn_func(q, k, v, softmax_scale=scale) + out = rearrange_out(out, out_pattern, required_out_pattern, dims) + return out + + +def sage_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None): + required_in_pattern, required_out_pattern= "b n s d", "b n s d" + q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims) + out = sageattn(q, k, v, sm_scale=scale) + out = rearrange_out(out, out_pattern, required_out_pattern, dims) + return out + + +def xformers_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None): + required_in_pattern, required_out_pattern= "b s n d", "b s n d" + q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims) + out = xops.memory_efficient_attention(q, k, v, scale=scale) + out = rearrange_out(out, out_pattern, required_out_pattern, dims) + return out + + +def attention_forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, compatibility_mode=False): + if compatibility_mode or (attn_mask is not None): + return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale) + else: + if ATTENTION_IMPLEMENTATION == "flash_attention_3": + return flash_attention_3(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) + elif ATTENTION_IMPLEMENTATION == "flash_attention_2": + return flash_attention_2(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) + elif ATTENTION_IMPLEMENTATION == "sage_attention": + return sage_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) + elif ATTENTION_IMPLEMENTATION == "xformers": + return xformers_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) + else: + return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) diff --git a/diffsynth/core/data/__init__.py b/diffsynth/core/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d494a277d81eeb2a9575155eb983d8bc3879590a --- /dev/null +++ b/diffsynth/core/data/__init__.py @@ -0,0 +1 @@ +from .unified_dataset import UnifiedDataset diff --git a/diffsynth/core/data/operators.py b/diffsynth/core/data/operators.py new file mode 100644 index 0000000000000000000000000000000000000000..c14944ea3de15dafaaff34b5badcb87c8839f6da --- /dev/null +++ b/diffsynth/core/data/operators.py @@ -0,0 +1,218 @@ +import torch, torchvision, imageio, os +import imageio.v3 as iio +from PIL import Image + + +class DataProcessingPipeline: + def __init__(self, operators=None): + self.operators: list[DataProcessingOperator] = [] if operators is None else operators + + def __call__(self, data): + for operator in self.operators: + data = operator(data) + return data + + def __rshift__(self, pipe): + if isinstance(pipe, DataProcessingOperator): + pipe = DataProcessingPipeline([pipe]) + return DataProcessingPipeline(self.operators + pipe.operators) + + +class DataProcessingOperator: + def __call__(self, data): + raise NotImplementedError("DataProcessingOperator cannot be called directly.") + + def __rshift__(self, pipe): + if isinstance(pipe, DataProcessingOperator): + pipe = DataProcessingPipeline([pipe]) + return DataProcessingPipeline([self]).__rshift__(pipe) + + +class DataProcessingOperatorRaw(DataProcessingOperator): + def __call__(self, data): + return data + + +class ToInt(DataProcessingOperator): + def __call__(self, data): + return int(data) + + +class ToFloat(DataProcessingOperator): + def __call__(self, data): + return float(data) + + +class ToStr(DataProcessingOperator): + def __init__(self, none_value=""): + self.none_value = none_value + + def __call__(self, data): + if data is None: data = self.none_value + return str(data) + + +class LoadImage(DataProcessingOperator): + def __init__(self, convert_RGB=True): + self.convert_RGB = convert_RGB + + def __call__(self, data: str): + image = Image.open(data) + if self.convert_RGB: image = image.convert("RGB") + return image + + +class ImageCropAndResize(DataProcessingOperator): + def __init__(self, height=None, width=None, max_pixels=None, height_division_factor=1, width_division_factor=1): + self.height = height + self.width = width + self.max_pixels = max_pixels + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + + def crop_and_resize(self, image, target_height, target_width): + width, height = image.size + scale = max(target_width / width, target_height / height) + image = torchvision.transforms.functional.resize( + image, + (round(height*scale), round(width*scale)), + interpolation=torchvision.transforms.InterpolationMode.BILINEAR + ) + image = torchvision.transforms.functional.center_crop(image, (target_height, target_width)) + return image + + def get_height_width(self, image): + if self.height is None or self.width is None: + width, height = image.size + if width * height > self.max_pixels: + scale = (width * height / self.max_pixels) ** 0.5 + height, width = int(height / scale), int(width / scale) + height = height // self.height_division_factor * self.height_division_factor + width = width // self.width_division_factor * self.width_division_factor + else: + height, width = self.height, self.width + return height, width + + def __call__(self, data: Image.Image): + image = self.crop_and_resize(data, *self.get_height_width(data)) + return image + + +class ToList(DataProcessingOperator): + def __call__(self, data): + return [data] + + +class LoadVideo(DataProcessingOperator): + def __init__(self, num_frames=81, time_division_factor=4, time_division_remainder=1, frame_processor=lambda x: x): + self.num_frames = num_frames + self.time_division_factor = time_division_factor + self.time_division_remainder = time_division_remainder + # frame_processor is build in the video loader for high efficiency. + self.frame_processor = frame_processor + + def get_num_frames(self, reader): + num_frames = self.num_frames + if int(reader.count_frames()) < num_frames: + num_frames = int(reader.count_frames()) + while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder: + num_frames -= 1 + return num_frames + + def __call__(self, data: str): + reader = imageio.get_reader(data) + num_frames = self.get_num_frames(reader) + frames = [] + for frame_id in range(num_frames): + frame = reader.get_data(frame_id) + frame = Image.fromarray(frame) + frame = self.frame_processor(frame) + frames.append(frame) + reader.close() + return frames + + +class SequencialProcess(DataProcessingOperator): + def __init__(self, operator=lambda x: x): + self.operator = operator + + def __call__(self, data): + return [self.operator(i) for i in data] + + +class LoadGIF(DataProcessingOperator): + def __init__(self, num_frames=81, time_division_factor=4, time_division_remainder=1, frame_processor=lambda x: x): + self.num_frames = num_frames + self.time_division_factor = time_division_factor + self.time_division_remainder = time_division_remainder + # frame_processor is build in the video loader for high efficiency. + self.frame_processor = frame_processor + + def get_num_frames(self, path): + num_frames = self.num_frames + images = iio.imread(path, mode="RGB") + if len(images) < num_frames: + num_frames = len(images) + while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder: + num_frames -= 1 + return num_frames + + def __call__(self, data: str): + num_frames = self.get_num_frames(data) + frames = [] + images = iio.imread(data, mode="RGB") + for img in images: + frame = Image.fromarray(img) + frame = self.frame_processor(frame) + frames.append(frame) + if len(frames) >= num_frames: + break + return frames + + +class RouteByExtensionName(DataProcessingOperator): + def __init__(self, operator_map): + self.operator_map = operator_map + + def __call__(self, data: str): + file_ext_name = data.split(".")[-1].lower() + for ext_names, operator in self.operator_map: + if ext_names is None or file_ext_name in ext_names: + return operator(data) + raise ValueError(f"Unsupported file: {data}") + + +class RouteByType(DataProcessingOperator): + def __init__(self, operator_map): + self.operator_map = operator_map + + def __call__(self, data): + for dtype, operator in self.operator_map: + if dtype is None or isinstance(data, dtype): + return operator(data) + raise ValueError(f"Unsupported data: {data}") + + +class LoadTorchPickle(DataProcessingOperator): + def __init__(self, map_location="cpu"): + self.map_location = map_location + + def __call__(self, data): + return torch.load(data, map_location=self.map_location, weights_only=False) + + +class ToAbsolutePath(DataProcessingOperator): + def __init__(self, base_path=""): + self.base_path = base_path + + def __call__(self, data): + return os.path.join(self.base_path, data) + + +class LoadAudio(DataProcessingOperator): + def __init__(self, sr=16000): + self.sr = sr + def __call__(self, data: str): + import librosa + input_audio, sample_rate = librosa.load(data, sr=self.sr) + return input_audio diff --git a/diffsynth/core/data/unified_dataset.py b/diffsynth/core/data/unified_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..074208cf81a5c77bbd54ea08f4699ce63891916e --- /dev/null +++ b/diffsynth/core/data/unified_dataset.py @@ -0,0 +1,112 @@ +from .operators import * +import torch, json, pandas + + +class UnifiedDataset(torch.utils.data.Dataset): + def __init__( + self, + base_path=None, metadata_path=None, + repeat=1, + data_file_keys=tuple(), + main_data_operator=lambda x: x, + special_operator_map=None, + ): + self.base_path = base_path + self.metadata_path = metadata_path + self.repeat = repeat + self.data_file_keys = data_file_keys + self.main_data_operator = main_data_operator + self.cached_data_operator = LoadTorchPickle() + self.special_operator_map = {} if special_operator_map is None else special_operator_map + self.data = [] + self.cached_data = [] + self.load_from_cache = metadata_path is None + self.load_metadata(metadata_path) + + @staticmethod + def default_image_operator( + base_path="", + max_pixels=1920*1080, height=None, width=None, + height_division_factor=16, width_division_factor=16, + ): + return RouteByType(operator_map=[ + (str, ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor)), + (list, SequencialProcess(ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor))), + ]) + + @staticmethod + def default_video_operator( + base_path="", + max_pixels=1920*1080, height=None, width=None, + height_division_factor=16, width_division_factor=16, + num_frames=81, time_division_factor=4, time_division_remainder=1, + ): + return RouteByType(operator_map=[ + (str, ToAbsolutePath(base_path) >> RouteByExtensionName(operator_map=[ + (("jpg", "jpeg", "png", "webp"), LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor) >> ToList()), + (("gif",), LoadGIF( + num_frames, time_division_factor, time_division_remainder, + frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), + )), + (("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), LoadVideo( + num_frames, time_division_factor, time_division_remainder, + frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), + )), + ])), + ]) + + def search_for_cached_data_files(self, path): + for file_name in os.listdir(path): + subpath = os.path.join(path, file_name) + if os.path.isdir(subpath): + self.search_for_cached_data_files(subpath) + elif subpath.endswith(".pth"): + self.cached_data.append(subpath) + + def load_metadata(self, metadata_path): + if metadata_path is None: + print("No metadata_path. Searching for cached data files.") + self.search_for_cached_data_files(self.base_path) + print(f"{len(self.cached_data)} cached data files found.") + elif metadata_path.endswith(".json"): + with open(metadata_path, "r") as f: + metadata = json.load(f) + self.data = metadata + elif metadata_path.endswith(".jsonl"): + metadata = [] + with open(metadata_path, 'r') as f: + for line in f: + metadata.append(json.loads(line.strip())) + self.data = metadata + else: + metadata = pandas.read_csv(metadata_path) + self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] + + def __getitem__(self, data_id): + if self.load_from_cache: + data = self.cached_data[data_id % len(self.cached_data)] + data = self.cached_data_operator(data) + else: + data = self.data[data_id % len(self.data)].copy() + for key in self.data_file_keys: + if key in data: + if key in self.special_operator_map: + data[key] = self.special_operator_map[key](data[key]) + elif key in self.data_file_keys: + data[key] = self.main_data_operator(data[key]) + return data + + def __len__(self): + if self.load_from_cache: + return len(self.cached_data) * self.repeat + else: + return len(self.data) * self.repeat + + def check_data_equal(self, data1, data2): + # Debug only + if len(data1) != len(data2): + return False + for k in data1: + if data1[k] != data2[k]: + return False + return True diff --git a/diffsynth/core/gradient/__init__.py b/diffsynth/core/gradient/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..57914792a78ec32f69c3c99ae37535598efc8d52 --- /dev/null +++ b/diffsynth/core/gradient/__init__.py @@ -0,0 +1 @@ +from .gradient_checkpoint import gradient_checkpoint_forward diff --git a/diffsynth/core/gradient/gradient_checkpoint.py b/diffsynth/core/gradient/gradient_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..b356415a004f3d74afdd45840f1fc4caf6659e16 --- /dev/null +++ b/diffsynth/core/gradient/gradient_checkpoint.py @@ -0,0 +1,34 @@ +import torch + + +def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + +def gradient_checkpoint_forward( + model, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + *args, + **kwargs, +): + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + else: + model_output = model(*args, **kwargs) + return model_output diff --git a/diffsynth/core/loader/__init__.py b/diffsynth/core/loader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1f56d814bae40436f66bca583e33d180d6e11247 --- /dev/null +++ b/diffsynth/core/loader/__init__.py @@ -0,0 +1,3 @@ +from .file import load_state_dict, hash_state_dict_keys, hash_model_file +from .model import load_model, load_model_with_disk_offload +from .config import ModelConfig diff --git a/diffsynth/core/loader/config.py b/diffsynth/core/loader/config.py new file mode 100644 index 0000000000000000000000000000000000000000..562675f3055cb4f4644f7bc4c524f3cbf6a54f76 --- /dev/null +++ b/diffsynth/core/loader/config.py @@ -0,0 +1,117 @@ +import torch, glob, os +from typing import Optional, Union +from dataclasses import dataclass +from modelscope import snapshot_download +from huggingface_hub import snapshot_download as hf_snapshot_download +from typing import Optional + + +@dataclass +class ModelConfig: + path: Union[str, list[str]] = None + model_id: str = None + origin_file_pattern: Union[str, list[str]] = None + download_source: str = None + local_model_path: str = None + skip_download: bool = None + offload_device: Optional[Union[str, torch.device]] = None + offload_dtype: Optional[torch.dtype] = None + onload_device: Optional[Union[str, torch.device]] = None + onload_dtype: Optional[torch.dtype] = None + preparing_device: Optional[Union[str, torch.device]] = None + preparing_dtype: Optional[torch.dtype] = None + computation_device: Optional[Union[str, torch.device]] = None + computation_dtype: Optional[torch.dtype] = None + clear_parameters: bool = False + + def check_input(self): + if self.path is None and self.model_id is None: + raise ValueError(f"""No valid model files. Please use `ModelConfig(path="xxx")` or `ModelConfig(model_id="xxx/yyy", origin_file_pattern="zzz")`. `skip_download=True` only supports the first one.""") + + def parse_original_file_pattern(self): + if self.origin_file_pattern is None or self.origin_file_pattern == "": + return "*" + elif self.origin_file_pattern.endswith("/"): + return self.origin_file_pattern + "*" + else: + return self.origin_file_pattern + + def parse_download_source(self): + if self.download_source is None: + if os.environ.get('DIFFSYNTH_DOWNLOAD_SOURCE') is not None: + return os.environ.get('DIFFSYNTH_DOWNLOAD_SOURCE') + else: + return "modelscope" + else: + return self.download_source + + def parse_skip_download(self): + if self.skip_download is None: + if os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD') is not None: + if os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD').lower() == "true": + return True + elif os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD').lower() == "false": + return False + else: + return False + else: + return self.skip_download + + def download(self): + origin_file_pattern = self.parse_original_file_pattern() + downloaded_files = glob.glob(origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id)) + download_source = self.parse_download_source() + if download_source.lower() == "modelscope": + snapshot_download( + self.model_id, + local_dir=os.path.join(self.local_model_path, self.model_id), + allow_file_pattern=origin_file_pattern, + ignore_file_pattern=downloaded_files, + local_files_only=False + ) + elif download_source.lower() == "huggingface": + hf_snapshot_download( + self.model_id, + local_dir=os.path.join(self.local_model_path, self.model_id), + allow_patterns=origin_file_pattern, + ignore_patterns=downloaded_files, + local_files_only=False + ) + else: + raise ValueError("`download_source` should be `modelscope` or `huggingface`.") + + def require_downloading(self): + if self.path is not None: + return False + skip_download = self.parse_skip_download() + return not skip_download + + def reset_local_model_path(self): + if os.environ.get('DIFFSYNTH_MODEL_BASE_PATH') is not None: + self.local_model_path = os.environ.get('DIFFSYNTH_MODEL_BASE_PATH') + elif self.local_model_path is None: + self.local_model_path = "./models" + + def download_if_necessary(self): + self.check_input() + self.reset_local_model_path() + if self.require_downloading(): + self.download() + if self.origin_file_pattern is None or self.origin_file_pattern == "": + self.path = os.path.join(self.local_model_path, self.model_id) + else: + self.path = glob.glob(os.path.join(self.local_model_path, self.model_id, self.origin_file_pattern)) + if isinstance(self.path, list) and len(self.path) == 1: + self.path = self.path[0] + + def vram_config(self): + return { + "offload_device": self.offload_device, + "offload_dtype": self.offload_dtype, + "onload_device": self.onload_device, + "onload_dtype": self.onload_dtype, + "preparing_device": self.preparing_device, + "preparing_dtype": self.preparing_dtype, + "computation_device": self.computation_device, + "computation_dtype": self.computation_dtype, + } diff --git a/diffsynth/core/loader/file.py b/diffsynth/core/loader/file.py new file mode 100644 index 0000000000000000000000000000000000000000..8f66961f25d4fc547a2ec638f9d6a93be851afb9 --- /dev/null +++ b/diffsynth/core/loader/file.py @@ -0,0 +1,121 @@ +from safetensors import safe_open +import torch, hashlib + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if isinstance(file_path, list): + state_dict = {} + for file_path_ in file_path: + state_dict.update(load_state_dict(file_path_, torch_dtype, device)) + return state_dict + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=str(device)) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if len(state_dict) == 1: + if "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + elif "module" in state_dict: + state_dict = state_dict["module"] + elif "model_state" in state_dict: + state_dict = state_dict["model_state"] + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() + + +def load_keys_dict(file_path): + if isinstance(file_path, list): + state_dict = {} + for file_path_ in file_path: + state_dict.update(load_keys_dict(file_path_)) + return state_dict + if file_path.endswith(".safetensors"): + return load_keys_dict_from_safetensors(file_path) + else: + return load_keys_dict_from_bin(file_path) + + +def load_keys_dict_from_safetensors(file_path): + keys_dict = {} + with safe_open(file_path, framework="pt", device="cpu") as f: + for k in f.keys(): + keys_dict[k] = f.get_slice(k).get_shape() + return keys_dict + + +def convert_state_dict_to_keys_dict(state_dict): + keys_dict = {} + for k, v in state_dict.items(): + if isinstance(v, torch.Tensor): + keys_dict[k] = list(v.shape) + else: + keys_dict[k] = convert_state_dict_to_keys_dict(v) + return keys_dict + + +def load_keys_dict_from_bin(file_path): + state_dict = load_state_dict_from_bin(file_path) + keys_dict = convert_state_dict_to_keys_dict(state_dict) + return keys_dict + + +def convert_keys_dict_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, dict): + keys.append(key + "|" + convert_keys_dict_to_single_str(value, with_shape=with_shape)) + else: + if with_shape: + shape = "_".join(map(str, list(value))) + keys.append(key + ":" + shape) + keys.append(key) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def hash_model_file(path, with_shape=True): + keys_dict = load_keys_dict(path) + keys_str = convert_keys_dict_to_single_str(keys_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() diff --git a/diffsynth/core/loader/model.py b/diffsynth/core/loader/model.py new file mode 100644 index 0000000000000000000000000000000000000000..84dd9ea76202516136e989b1aadecb812f8e70a2 --- /dev/null +++ b/diffsynth/core/loader/model.py @@ -0,0 +1,82 @@ +from ..vram.initialization import skip_model_initialization +from ..vram.disk_map import DiskMap +from ..vram.layers import enable_vram_management +from .file import load_state_dict +import torch + + +def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, device="cpu", state_dict_converter=None, use_disk_map=False, module_map=None, vram_config=None, vram_limit=None): + config = {} if config is None else config + # Why do we use `skip_model_initialization`? + # It skips the random initialization of model parameters, + # thereby speeding up model loading and avoiding excessive memory usage. + # with skip_model_initialization(): + # model = model_class(**config) + model = model_class(**config) + + # What is `module_map`? + # This is a module mapping table for VRAM management. + if module_map is not None: + devices = [vram_config["offload_device"], vram_config["onload_device"], vram_config["preparing_device"], vram_config["computation_device"]] + device = [d for d in devices if d != "disk"][0] + dtypes = [vram_config["offload_dtype"], vram_config["onload_dtype"], vram_config["preparing_dtype"], vram_config["computation_dtype"]] + dtype = [d for d in dtypes if d != "disk"][0] + if vram_config["offload_device"] != "disk": + state_dict = DiskMap(path, device, torch_dtype=dtype) + if state_dict_converter is not None: + state_dict = state_dict_converter(state_dict) + else: + state_dict = {i: state_dict[i] for i in state_dict} + model.load_state_dict(state_dict, assign=True) + model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=None, vram_limit=vram_limit) + else: + disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter) + model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=vram_limit) + else: + # Why do we use `DiskMap`? + # Sometimes a model file contains multiple models, + # and DiskMap can load only the parameters of a single model, + # avoiding the need to load all parameters in the file. + if use_disk_map: + state_dict = DiskMap(path, device, torch_dtype=torch_dtype) + else: + state_dict = load_state_dict(path, torch_dtype, device) + # Why do we use `state_dict_converter`? + # Some models are saved in complex formats, + # and we need to convert the state dict into the appropriate format. + if state_dict_converter is not None: + state_dict = state_dict_converter(state_dict) + else: + state_dict = {i: state_dict[i] for i in state_dict} + model.load_state_dict(state_dict, assign=True, strict=False) + # Why do we call `to()`? + # Because some models override the behavior of `to()`, + # especially those from libraries like Transformers. + model = model.to(dtype=torch_dtype, device=device) + # model = model.to_empty(dtype=torch_dtype, device=device) + if hasattr(model, "eval"): + model = model.eval() + return model + + +def load_model_with_disk_offload(model_class, path, config=None, torch_dtype=torch.bfloat16, device="cpu", state_dict_converter=None, module_map=None): + if isinstance(path, str): + path = [path] + config = {} if config is None else config + with skip_model_initialization(): + model = model_class(**config) + if hasattr(model, "eval"): + model = model.eval() + disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter) + vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": device, + "computation_dtype": torch_dtype, + "computation_device": device, + } + enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=80) + return model diff --git a/diffsynth/core/vram/__init__.py b/diffsynth/core/vram/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..32763bb9b4abfa4d5b2617827661c520d7e9fcae --- /dev/null +++ b/diffsynth/core/vram/__init__.py @@ -0,0 +1,2 @@ +from .initialization import skip_model_initialization +from .layers import * diff --git a/diffsynth/core/vram/disk_map.py b/diffsynth/core/vram/disk_map.py new file mode 100644 index 0000000000000000000000000000000000000000..a666590fa99a9cc4de05dc3f5fa84c212e43de38 --- /dev/null +++ b/diffsynth/core/vram/disk_map.py @@ -0,0 +1,93 @@ +from safetensors import safe_open +import torch, os + + +class SafetensorsCompatibleTensor: + def __init__(self, tensor): + self.tensor = tensor + + def get_shape(self): + return list(self.tensor.shape) + + +class SafetensorsCompatibleBinaryLoader: + def __init__(self, path, device): + print("Detected non-safetensors files, which may cause slower loading. It's recommended to convert it to a safetensors file.") + self.state_dict = torch.load(path, weights_only=True, map_location=device) + + def keys(self): + return self.state_dict.keys() + + def get_tensor(self, name): + return self.state_dict[name] + + def get_slice(self, name): + return SafetensorsCompatibleTensor(self.state_dict[name]) + + +class DiskMap: + + def __init__(self, path, device, torch_dtype=None, state_dict_converter=None, buffer_size=10**9): + self.path = path if isinstance(path, list) else [path] + self.device = device + self.torch_dtype = torch_dtype + if os.environ.get('DIFFSYNTH_DISK_MAP_BUFFER_SIZE') is not None: + self.buffer_size = int(os.environ.get('DIFFSYNTH_DISK_MAP_BUFFER_SIZE')) + else: + self.buffer_size = buffer_size + self.files = [] + self.flush_files() + self.name_map = {} + for file_id, file in enumerate(self.files): + for name in file.keys(): + self.name_map[name] = file_id + self.rename_dict = self.fetch_rename_dict(state_dict_converter) + + def flush_files(self): + if len(self.files) == 0: + for path in self.path: + if path.endswith(".safetensors"): + self.files.append(safe_open(path, framework="pt", device=str(self.device))) + else: + self.files.append(SafetensorsCompatibleBinaryLoader(path, device=self.device)) + else: + for i, path in enumerate(self.path): + if path.endswith(".safetensors"): + self.files[i] = safe_open(path, framework="pt", device=str(self.device)) + self.num_params = 0 + + def __getitem__(self, name): + if self.rename_dict is not None: name = self.rename_dict[name] + file_id = self.name_map[name] + param = self.files[file_id].get_tensor(name) + if self.torch_dtype is not None and isinstance(param, torch.Tensor): + param = param.to(self.torch_dtype) + if isinstance(param, torch.Tensor) and param.device == "cpu": + param = param.clone() + if isinstance(param, torch.Tensor): + self.num_params += param.numel() + if self.num_params > self.buffer_size: + self.flush_files() + return param + + def fetch_rename_dict(self, state_dict_converter): + if state_dict_converter is None: + return None + state_dict = {} + for file in self.files: + for name in file.keys(): + state_dict[name] = name + state_dict = state_dict_converter(state_dict) + return state_dict + + def __iter__(self): + if self.rename_dict is not None: + return self.rename_dict.__iter__() + else: + return self.name_map.__iter__() + + def __contains__(self, x): + if self.rename_dict is not None: + return x in self.rename_dict + else: + return x in self.name_map diff --git a/diffsynth/core/vram/initialization.py b/diffsynth/core/vram/initialization.py new file mode 100644 index 0000000000000000000000000000000000000000..bff2498b526638bfdd1c114c78aa0b98c251a47d --- /dev/null +++ b/diffsynth/core/vram/initialization.py @@ -0,0 +1,21 @@ +import torch +from contextlib import contextmanager + + +@contextmanager +def skip_model_initialization(device=torch.device("meta")): + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + old_register_parameter = torch.nn.Module.register_parameter + torch.nn.Module.register_parameter = register_empty_parameter + try: + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter diff --git a/diffsynth/core/vram/layers.py b/diffsynth/core/vram/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..852a8f33855f4906b26d4bbf177e9ea7b4725164 --- /dev/null +++ b/diffsynth/core/vram/layers.py @@ -0,0 +1,475 @@ +import torch, copy +from typing import Union +from .initialization import skip_model_initialization +from .disk_map import DiskMap + + +class AutoTorchModule(torch.nn.Module): + + def __init__( + self, + offload_dtype: torch.dtype = None, + offload_device: Union[str, torch.device] = None, + onload_dtype: torch.dtype = None, + onload_device: Union[str, torch.device] = None, + preparing_dtype: torch.dtype = None, + preparing_device: Union[str, torch.device] = None, + computation_dtype: torch.dtype = None, + computation_device: Union[str, torch.device] = None, + vram_limit: float = None, + ): + super().__init__() + self.set_dtype_and_device( + offload_dtype, + offload_device, + onload_dtype, + onload_device, + preparing_dtype, + preparing_device, + computation_dtype, + computation_device, + vram_limit, + ) + self.state = 0 + self.name = "" + + def set_dtype_and_device( + self, + offload_dtype: torch.dtype = None, + offload_device: Union[str, torch.device] = None, + onload_dtype: torch.dtype = None, + onload_device: Union[str, torch.device] = None, + preparing_dtype: torch.dtype = None, + preparing_device: Union[str, torch.device] = None, + computation_dtype: torch.dtype = None, + computation_device: Union[str, torch.device] = None, + vram_limit: float = None, + ): + self.offload_dtype = offload_dtype or computation_dtype + self.offload_device = offload_device or computation_dtype + self.onload_dtype = onload_dtype or computation_dtype + self.onload_device = onload_device or computation_dtype + self.preparing_dtype = preparing_dtype or computation_dtype + self.preparing_device = preparing_device or computation_dtype + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + + def cast_to(self, weight, dtype, device): + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight) + return r + + def check_free_vram(self): + gpu_mem_state = torch.cuda.mem_get_info(self.computation_device) + used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024**3) + return used_memory < self.vram_limit + + def offload(self): + if self.state != 0: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + if self.state != 1: + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def param_name(self, name): + if self.name == "": + return name + else: + return self.name + "." + name + + +class AutoWrappedModule(AutoTorchModule): + + def __init__( + self, + module: torch.nn.Module, + offload_dtype: torch.dtype = None, + offload_device: Union[str, torch.device] = None, + onload_dtype: torch.dtype = None, + onload_device: Union[str, torch.device] = None, + preparing_dtype: torch.dtype = None, + preparing_device: Union[str, torch.device] = None, + computation_dtype: torch.dtype = None, + computation_device: Union[str, torch.device] = None, + vram_limit: float = None, + name: str = "", + disk_map: DiskMap = None, + **kwargs + ): + super().__init__( + offload_dtype, + offload_device, + onload_dtype, + onload_device, + preparing_dtype, + preparing_device, + computation_dtype, + computation_device, + vram_limit, + ) + self.module = module + if offload_dtype == "disk": + self.name = name + self.disk_map = disk_map + self.required_params = [name for name, _ in self.module.named_parameters()] + self.disk_offload = True + else: + self.disk_offload = False + + def load_from_disk(self, torch_dtype, device, copy_module=False): + if copy_module: + module = copy.deepcopy(self.module) + else: + module = self.module + state_dict = {} + for name in self.required_params: + param = self.disk_map[self.param_name(name)] + param = param.to(dtype=torch_dtype, device=device) + state_dict[name] = param + module.load_state_dict(state_dict, assign=True) + module.to(dtype=torch_dtype, device=device) + return module + + def offload_to_disk(self, model: torch.nn.Module): + for buf in model.buffers(): + # If there are some parameters are registed in buffers (not in state dict), + # We cannot offload the model. + for children in model.children(): + self.offload_to_disk(children) + break + else: + model.to("meta") + + def offload(self): + # offload / onload / preparing -> offload + if self.state != 0: + if self.disk_offload: + self.offload_to_disk(self.module) + else: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + # offload / onload / preparing -> onload + if self.state < 1: + if self.disk_offload and self.onload_device != "disk" and self.offload_device == "disk": + self.load_from_disk(self.onload_dtype, self.onload_device) + elif self.onload_device != "disk": + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def preparing(self): + # onload / preparing -> preparing + if self.state != 2: + if self.disk_offload and self.preparing_device != "disk" and self.onload_device == "disk": + self.load_from_disk(self.preparing_dtype, self.preparing_device) + elif self.preparing_device != "disk": + self.to(dtype=self.preparing_dtype, device=self.preparing_device) + self.state = 2 + + def cast_to(self, module, dtype, device): + return copy.deepcopy(module).to(dtype=dtype, device=device) + + def computation(self): + # onload / preparing -> computation (temporary) + if self.state == 2: + torch_dtype, device = self.preparing_dtype, self.preparing_device + else: + torch_dtype, device = self.onload_dtype, self.onload_device + if torch_dtype == self.computation_dtype and device == self.computation_device: + module = self.module + elif self.disk_offload and device == "disk": + module = self.load_from_disk(self.computation_dtype, self.computation_device, copy_module=True) + else: + module = self.cast_to(self.module, dtype=self.computation_dtype, device=self.computation_device) + return module + + def forward(self, *args, **kwargs): + if self.state == 1 and (self.vram_limit is None or self.check_free_vram()): + self.preparing() + module = self.computation() + return module(*args, **kwargs) + + def __getattr__(self, name): + if name in self.__dict__ or name == "module": + return super().__getattr__(name) + else: + return getattr(self.module, name) + + +class AutoWrappedNonRecurseModule(AutoWrappedModule): + + def __init__( + self, + module: torch.nn.Module, + offload_dtype: torch.dtype = None, + offload_device: Union[str, torch.device] = None, + onload_dtype: torch.dtype = None, + onload_device: Union[str, torch.device] = None, + preparing_dtype: torch.dtype = None, + preparing_device: Union[str, torch.device] = None, + computation_dtype: torch.dtype = None, + computation_device: Union[str, torch.device] = None, + vram_limit: float = None, + name: str = "", + disk_map: DiskMap = None, + **kwargs + ): + super().__init__( + module, + offload_dtype, + offload_device, + onload_dtype, + onload_device, + preparing_dtype, + preparing_device, + computation_dtype, + computation_device, + vram_limit, + name, + disk_map, + **kwargs + ) + if self.disk_offload: + self.required_params = [name for name, _ in self.module.named_parameters(recurse=False)] + + def load_from_disk(self, torch_dtype, device, copy_module=False): + if copy_module: + module = copy.deepcopy(self.module) + else: + module = self.module + state_dict = {} + for name in self.required_params: + param = self.disk_map[self.param_name(name)] + param = param.to(dtype=torch_dtype, device=device) + state_dict[name] = param + module.load_state_dict(state_dict, assign=True, strict=False) + return module + + def offload_to_disk(self, model: torch.nn.Module): + for name in self.required_params: + getattr(self, name).to("meta") + + def cast_to(self, module, dtype, device): + # Parameter casting is implemented in the model architecture. + return module + + def __getattr__(self, name): + if name in self.__dict__ or name == "module": + return super().__getattr__(name) + else: + return getattr(self.module, name) + + +class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule): + def __init__( + self, + module: torch.nn.Linear, + offload_dtype: torch.dtype = None, + offload_device: Union[str, torch.device] = None, + onload_dtype: torch.dtype = None, + onload_device: Union[str, torch.device] = None, + preparing_dtype: torch.dtype = None, + preparing_device: Union[str, torch.device] = None, + computation_dtype: torch.dtype = None, + computation_device: Union[str, torch.device] = None, + vram_limit: float = None, + name: str = "", + disk_map: DiskMap = None, + **kwargs + ): + with skip_model_initialization(): + super().__init__( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + ) + self.set_dtype_and_device( + offload_dtype, + offload_device, + onload_dtype, + onload_device, + preparing_dtype, + preparing_device, + computation_dtype, + computation_device, + vram_limit, + ) + self.weight = module.weight + self.bias = module.bias + self.state = 0 + self.name = name + self.lora_A_weights = [] + self.lora_B_weights = [] + self.lora_merger = None + self.enable_fp8 = computation_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz] + + if offload_dtype == "disk": + self.disk_map = disk_map + self.disk_offload = True + else: + self.disk_offload = False + + def fp8_linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ) -> torch.Tensor: + device = input.device + origin_dtype = input.dtype + origin_shape = input.shape + input = input.reshape(-1, origin_shape[-1]) + + x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values + fp8_max = 448.0 + # For float8_e4m3fnuz, the maximum representable value is half of that of e4m3fn. + # To avoid overflow and ensure numerical compatibility during FP8 computation, + # we scale down the input by 2.0 in advance. + # This scaling will be compensated later during the final result scaling. + if self.computation_dtype == torch.float8_e4m3fnuz: + fp8_max = fp8_max / 2.0 + scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device) + scale_b = torch.ones((weight.shape[0], 1)).to(device=device) + input = input / (scale_a + 1e-8) + input = input.to(self.computation_dtype) + weight = weight.to(self.computation_dtype) + bias = bias.to(torch.bfloat16) + + result = torch._scaled_mm( + input, + weight.T, + scale_a=scale_a, + scale_b=scale_b.T, + bias=bias, + out_dtype=origin_dtype, + ) + new_shape = origin_shape[:-1] + result.shape[-1:] + result = result.reshape(new_shape) + return result + + def load_from_disk(self, torch_dtype, device, assign=True): + weight = self.disk_map[self.name + ".weight"].to(dtype=torch_dtype, device=device) + bias = None if self.bias is None else self.disk_map[self.name + ".bias"].to(dtype=torch_dtype, device=device) + if assign: + state_dict = {"weight": weight} + if bias is not None: state_dict["bias"] = bias + self.load_state_dict(state_dict, assign=True) + return weight, bias + + def offload(self): + # offload / onload / preparing -> offload + if self.state != 0: + if self.disk_offload: + self.to("meta") + else: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + # offload / onload / preparing -> onload + if self.state < 1: + if self.disk_offload and self.onload_device != "disk" and self.offload_device == "disk": + self.load_from_disk(self.onload_dtype, self.onload_device) + elif self.onload_device != "disk": + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def preparing(self): + # onload / preparing -> preparing + if self.state != 2: + if self.disk_offload and self.preparing_device != "disk" and self.onload_device == "disk": + self.load_from_disk(self.preparing_dtype, self.preparing_device) + elif self.preparing_device != "disk": + self.to(dtype=self.preparing_dtype, device=self.preparing_device) + self.state = 2 + + def computation(self): + # onload / preparing -> computation (temporary) + if self.state == 2: + torch_dtype, device = self.preparing_dtype, self.preparing_device + else: + torch_dtype, device = self.onload_dtype, self.onload_device + if torch_dtype == self.computation_dtype and device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.disk_offload and device == "disk": + weight, bias = self.load_from_disk(self.computation_dtype, self.computation_device, assign=False) + else: + weight = self.cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else self.cast_to(self.bias, self.computation_dtype, self.computation_device) + return weight, bias + + def linear_forward(self, x, weight, bias): + if self.enable_fp8: + out = self.fp8_linear(x, weight, bias) + else: + out = torch.nn.functional.linear(x, weight, bias) + return out + + def lora_forward(self, x, out): + if self.lora_merger is None: + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + out = out + x @ lora_A.T @ lora_B.T + else: + lora_output = [] + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + lora_output.append(x @ lora_A.T @ lora_B.T) + lora_output = torch.stack(lora_output) + out = self.lora_merger(out, lora_output) + return out + + def forward(self, x, *args, **kwargs): + if self.state == 1 and (self.vram_limit is None or self.check_free_vram()): + self.preparing() + weight, bias = self.computation() + out = self.linear_forward(x, weight, bias) + if len(self.lora_A_weights) > 0: + out = self.lora_forward(x, out) + return out + + +def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, vram_config: dict, vram_limit=None, name_prefix="", disk_map=None, **kwargs): + if isinstance(model, AutoWrappedNonRecurseModule): + model = model.module + for name, module in model.named_children(): + layer_name = name if name_prefix == "" else name_prefix + "." + name + for source_module, target_module in module_map.items(): + if isinstance(module, source_module): + module_ = target_module(module, **vram_config, vram_limit=vram_limit, name=layer_name, disk_map=disk_map, **kwargs) + if isinstance(module_, AutoWrappedNonRecurseModule): + enable_vram_management_recursively(module_, module_map, vram_config, vram_limit=vram_limit, name_prefix=layer_name, disk_map=disk_map, **kwargs) + setattr(model, name, module_) + break + else: + enable_vram_management_recursively(module, module_map, vram_config, vram_limit=vram_limit, name_prefix=layer_name, disk_map=disk_map, **kwargs) + + +def fill_vram_config(model, vram_config): + vram_config_ = vram_config.copy() + vram_config_["onload_dtype"] = vram_config["computation_dtype"] + vram_config_["onload_device"] = vram_config["computation_device"] + vram_config_["preparing_dtype"] = vram_config["computation_dtype"] + vram_config_["preparing_device"] = vram_config["computation_device"] + for k in vram_config: + if vram_config[k] != vram_config_[k]: + print(f"No fine-grained VRAM configuration is provided for {model.__class__.__name__}. [`onload`, `preparing`, `computation`] will be the same state. `vram_config` is set to {vram_config_}") + break + return vram_config_ + + +def enable_vram_management(model: torch.nn.Module, module_map: dict, vram_config: dict, vram_limit=None, disk_map=None, **kwargs): + for source_module, target_module in module_map.items(): + # If no fine-grained VRAM configuration is provided, the entire model will be managed uniformly. + if isinstance(model, source_module): + vram_config = fill_vram_config(model, vram_config) + model = target_module(model, **vram_config, vram_limit=vram_limit, disk_map=disk_map, **kwargs) + break + else: + enable_vram_management_recursively(model, module_map, vram_config, vram_limit=vram_limit, disk_map=disk_map, **kwargs) + # `vram_management_enabled` is a flag that allows the pipeline to determine whether VRAM management is enabled. + model.vram_management_enabled = True + return model diff --git a/diffsynth/diffusion/__init__.py b/diffsynth/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0a4a0873a7b3d09e95aa00cfe340d653c58a834b --- /dev/null +++ b/diffsynth/diffusion/__init__.py @@ -0,0 +1,6 @@ +from .flow_match import FlowMatchScheduler +from .training_module import DiffusionTrainingModule +from .logger import ModelLogger +from .runner import launch_training_task, launch_data_process_task +from .parsers import * +from .loss import * diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..2c52327f65b5ec7a1ca3aaad75cf6989edfe487b --- /dev/null +++ b/diffsynth/diffusion/base_pipeline.py @@ -0,0 +1,440 @@ +from PIL import Image +import torch +import numpy as np +from einops import repeat, reduce +from typing import Union +from ..core import AutoTorchModule, AutoWrappedLinear, load_state_dict, ModelConfig +from ..utils.lora import GeneralLoRALoader +from ..models.model_loader import ModelPool +from ..utils.controlnet import ControlNetInput + + +class PipelineUnit: + def __init__( + self, + seperate_cfg: bool = False, + take_over: bool = False, + input_params: tuple[str] = None, + output_params: tuple[str] = None, + input_params_posi: dict[str, str] = None, + input_params_nega: dict[str, str] = None, + onload_model_names: tuple[str] = None + ): + self.seperate_cfg = seperate_cfg + self.take_over = take_over + self.input_params = input_params + self.output_params = output_params + self.input_params_posi = input_params_posi + self.input_params_nega = input_params_nega + self.onload_model_names = onload_model_names + + def fetch_input_params(self): + params = [] + if self.input_params is not None: + for param in self.input_params: + params.append(param) + if self.input_params_posi is not None: + for _, param in self.input_params_posi.items(): + params.append(param) + if self.input_params_nega is not None: + for _, param in self.input_params_nega.items(): + params.append(param) + params = sorted(list(set(params))) + return params + + def fetch_output_params(self): + params = [] + if self.output_params is not None: + for param in self.output_params: + params.append(param) + return params + + def process(self, pipe, **kwargs) -> dict: + return {} + + def post_process(self, pipe, **kwargs) -> dict: + return {} + + +class BasePipeline(torch.nn.Module): + + def __init__( + self, + device="cuda", torch_dtype=torch.float16, + height_division_factor=64, width_division_factor=64, + time_division_factor=None, time_division_remainder=None, + ): + super().__init__() + # The device and torch_dtype is used for the storage of intermediate variables, not models. + self.device = device + self.torch_dtype = torch_dtype + # The following parameters are used for shape check. + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + self.time_division_factor = time_division_factor + self.time_division_remainder = time_division_remainder + # VRAM management + self.vram_management_enabled = False + # Pipeline Unit Runner + self.unit_runner = PipelineUnitRunner() + # LoRA Loader + self.lora_loader = GeneralLoRALoader + + + def to(self, *args, **kwargs): + device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) + if device is not None: + self.device = device + if dtype is not None: + self.torch_dtype = dtype + super().to(*args, **kwargs) + return self + + + def check_resize_height_width(self, height, width, num_frames=None): + # Shape check + if height % self.height_division_factor != 0: + height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor + print(f"height % {self.height_division_factor} != 0. We round it up to {height}.") + if width % self.width_division_factor != 0: + width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor + print(f"width % {self.width_division_factor} != 0. We round it up to {width}.") + if num_frames is None: + return height, width + else: + if num_frames % self.time_division_factor != self.time_division_remainder: + num_frames = (num_frames + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder + print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.") + return height, width, num_frames + + + def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1): + # Transform a PIL.Image to torch.Tensor + image = torch.Tensor(np.array(image, dtype=np.float32)) + # print(image.shape) + image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) + image = image * ((max_value - min_value) / 255) + min_value + image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {})) + return image + + + def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1): + # Transform a list of PIL.Image to torch.Tensor + video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video] + video = torch.stack(video, dim=pattern.index("T") // 2) + return video + + + def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1): + # Transform a torch.Tensor to PIL.Image + if pattern != "H W C": + vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean") + image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255) + image = image.to(device="cpu", dtype=torch.uint8) + image = Image.fromarray(image.numpy()) + return image + + + def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1): + # Transform a torch.Tensor to list of PIL.Image + if pattern != "T H W C": + vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean") + video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output] + return video + + + def load_models_to_device(self, model_names): + if self.vram_management_enabled: + # offload models + for name, model in self.named_children(): + if name not in model_names: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + if hasattr(model, "offload"): + model.offload() + else: + for module in model.modules(): + if hasattr(module, "offload"): + module.offload() + torch.cuda.empty_cache() + # onload models + for name, model in self.named_children(): + if name in model_names: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + if hasattr(model, "onload"): + model.onload() + else: + for module in model.modules(): + if hasattr(module, "onload"): + module.onload() + + + def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None): + # Initialize Gaussian noise + generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) + noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype) + noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) + return noise + + + def get_vram(self): + return torch.cuda.mem_get_info(self.device)[1] / (1024 ** 3) + + def get_module(self, model, name): + if "." in name: + name, suffix = name[:name.index(".")], name[name.index(".") + 1:] + if name.isdigit(): + return self.get_module(model[int(name)], suffix) + else: + return self.get_module(getattr(model, name), suffix) + else: + return getattr(model, name) + + def freeze_except(self, model_names): + self.eval() + self.requires_grad_(False) + for name in model_names: + module = self.get_module(self, name) + if module is None: + print(f"No {name} models in the pipeline. We cannot enable training on the model. If this occurs during the data processing stage, it is normal.") + continue + module.train() + module.requires_grad_(True) + + + def blend_with_mask(self, base, addition, mask): + return base * (1 - mask) + addition * mask + + + def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs): + timestep = scheduler.timesteps[progress_id] + if inpaint_mask is not None: + noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents) + noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask) + latents_next = scheduler.step(noise_pred, timestep, latents) + return latents_next + + + def split_pipeline_units(self, model_names: list[str]): + return PipelineUnitGraph().split_pipeline_units(self.units, model_names) + + + def flush_vram_management_device(self, device): + for module in self.modules(): + if isinstance(module, AutoTorchModule): + module.offload_device = device + module.onload_device = device + module.preparing_device = device + module.computation_device = device + + + def load_lora( + self, + module: torch.nn.Module, + lora_config: Union[ModelConfig, str] = None, + alpha=1, + hotload=None, + state_dict=None, + ): + if state_dict is None: + if isinstance(lora_config, str): + lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device) + else: + lora_config.download_if_necessary() + lora = load_state_dict(lora_config.path, torch_dtype=self.torch_dtype, device=self.device) + else: + lora = state_dict + lora_loader = self.lora_loader(torch_dtype=self.torch_dtype, device=self.device) + lora = lora_loader.convert_state_dict(lora) + if hotload is None: + hotload = hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled") + if hotload: + if not (hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")): + raise ValueError("VRAM Management is not enabled. LoRA hotloading is not supported.") + updated_num = 0 + for _, module in module.named_modules(): + if isinstance(module, AutoWrappedLinear): + name = module.name + lora_a_name = f'{name}.lora_A.weight' + lora_b_name = f'{name}.lora_B.weight' + if lora_a_name in lora and lora_b_name in lora: + updated_num += 1 + module.lora_A_weights.append(lora[lora_a_name] * alpha) + module.lora_B_weights.append(lora[lora_b_name]) + print(f"{updated_num} tensors are patched by LoRA. You can use `pipe.clear_lora()` to clear all LoRA layers.") + else: + lora_loader.fuse_lora_to_base_model(module, lora, alpha=alpha) + + + def clear_lora(self): + cleared_num = 0 + for name, module in self.named_modules(): + if isinstance(module, AutoWrappedLinear): + if hasattr(module, "lora_A_weights"): + if len(module.lora_A_weights) > 0: + cleared_num += 1 + module.lora_A_weights.clear() + if hasattr(module, "lora_B_weights"): + module.lora_B_weights.clear() + print(f"{cleared_num} LoRA layers are cleared.") + + + def download_and_load_models(self, model_configs: list[ModelConfig] = [], vram_limit: float = None): + model_pool = ModelPool() + for model_config in model_configs: + model_config.download_if_necessary() + vram_config = model_config.vram_config() + vram_config["computation_dtype"] = vram_config["computation_dtype"] or self.torch_dtype + vram_config["computation_device"] = vram_config["computation_device"] or self.device + model_pool.auto_load_model( + model_config.path, + vram_config=vram_config, + vram_limit=vram_limit, + clear_parameters=model_config.clear_parameters, + ) + return model_pool + + + def check_vram_management_state(self): + vram_management_enabled = False + for module in self.children(): + if hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled"): + vram_management_enabled = True + return vram_management_enabled + + + def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **inputs_others): + noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others) + if cfg_scale != 1.0: + noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others) + noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega) + else: + noise_pred = noise_pred_posi + return noise_pred + + +class PipelineUnitGraph: + def __init__(self): + pass + + def build_edges(self, units: list[PipelineUnit]): + # Establish dependencies between units + # to search for subsequent related computation units. + last_compute_unit_id = {} + edges = [] + for unit_id, unit in enumerate(units): + for input_param in unit.fetch_input_params(): + if input_param in last_compute_unit_id: + edges.append((last_compute_unit_id[input_param], unit_id)) + for output_param in unit.fetch_output_params(): + last_compute_unit_id[output_param] = unit_id + return edges + + def build_chains(self, units: list[PipelineUnit]): + # Establish updating chains for each variable + # to track their computation process. + params = sum([unit.fetch_input_params() + unit.fetch_output_params() for unit in units], []) + params = sorted(list(set(params))) + chains = {param: [] for param in params} + for unit_id, unit in enumerate(units): + for param in unit.fetch_output_params(): + chains[param].append(unit_id) + return chains + + def search_direct_unit_ids(self, units: list[PipelineUnit], model_names: list[str]): + # Search for units that directly participate in the model's computation. + related_unit_ids = [] + for unit_id, unit in enumerate(units): + for model_name in model_names: + if unit.onload_model_names is not None and model_name in unit.onload_model_names: + related_unit_ids.append(unit_id) + break + return related_unit_ids + + def search_related_unit_ids(self, edges, start_unit_ids, direction="target"): + # Search for subsequent related computation units. + related_unit_ids = [unit_id for unit_id in start_unit_ids] + while True: + neighbors = [] + for source, target in edges: + if direction == "target" and source in related_unit_ids and target not in related_unit_ids: + neighbors.append(target) + elif direction == "source" and source not in related_unit_ids and target in related_unit_ids: + neighbors.append(source) + neighbors = sorted(list(set(neighbors))) + if len(neighbors) == 0: + break + else: + related_unit_ids.extend(neighbors) + related_unit_ids = sorted(list(set(related_unit_ids))) + return related_unit_ids + + def search_updating_unit_ids(self, units: list[PipelineUnit], chains, related_unit_ids): + # If the input parameters of this subgraph are updated outside the subgraph, + # search for the units where these updates occur. + first_compute_unit_id = {} + for unit_id in related_unit_ids: + for param in units[unit_id].fetch_input_params(): + if param not in first_compute_unit_id: + first_compute_unit_id[param] = unit_id + updating_unit_ids = [] + for param in first_compute_unit_id: + unit_id = first_compute_unit_id[param] + chain = chains[param] + if unit_id in chain and chain.index(unit_id) != len(chain) - 1: + for unit_id_ in chain[chain.index(unit_id) + 1:]: + if unit_id_ not in related_unit_ids: + updating_unit_ids.append(unit_id_) + related_unit_ids.extend(updating_unit_ids) + related_unit_ids = sorted(list(set(related_unit_ids))) + return related_unit_ids + + def split_pipeline_units(self, units: list[PipelineUnit], model_names: list[str]): + # Split the computation graph, + # separating all model-related computations. + related_unit_ids = self.search_direct_unit_ids(units, model_names) + edges = self.build_edges(units) + chains = self.build_chains(units) + while True: + num_related_unit_ids = len(related_unit_ids) + related_unit_ids = self.search_related_unit_ids(edges, related_unit_ids, "target") + related_unit_ids = self.search_updating_unit_ids(units, chains, related_unit_ids) + if len(related_unit_ids) == num_related_unit_ids: + break + else: + num_related_unit_ids = len(related_unit_ids) + related_units = [units[i] for i in related_unit_ids] + unrelated_units = [units[i] for i in range(len(units)) if i not in related_unit_ids] + return related_units, unrelated_units + + +class PipelineUnitRunner: + def __init__(self): + pass + + def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]: + if unit.take_over: + # Let the pipeline unit take over this function. + inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega) + elif unit.seperate_cfg: + # Positive side + processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()} + if unit.input_params is not None: + for name in unit.input_params: + processor_inputs[name] = inputs_shared.get(name) + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_posi.update(processor_outputs) + # Negative side + if inputs_shared["cfg_scale"] != 1: + processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()} + if unit.input_params is not None: + for name in unit.input_params: + processor_inputs[name] = inputs_shared.get(name) + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_nega.update(processor_outputs) + else: + inputs_nega.update(processor_outputs) + else: + processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params} + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_shared.update(processor_outputs) + return inputs_shared, inputs_posi, inputs_nega diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py new file mode 100644 index 0000000000000000000000000000000000000000..bb5fbc52855d730354daa7f92d1a0ed3715f670b --- /dev/null +++ b/diffsynth/diffusion/flow_match.py @@ -0,0 +1,179 @@ +import torch, math +from typing_extensions import Literal + + +class FlowMatchScheduler(): + + def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z-Image"] = "FLUX.1"): + self.set_timesteps_fn = { + "FLUX.1": FlowMatchScheduler.set_timesteps_flux, + "Wan": FlowMatchScheduler.set_timesteps_wan, + "Qwen-Image": FlowMatchScheduler.set_timesteps_qwen_image, + "FLUX.2": FlowMatchScheduler.set_timesteps_flux2, + "Z-Image": FlowMatchScheduler.set_timesteps_z_image, + }.get(template, FlowMatchScheduler.set_timesteps_flux) + self.num_train_timesteps = 1000 + + @staticmethod + def set_timesteps_flux(num_inference_steps=100, denoising_strength=1.0, shift=None): + sigma_min = 0.003/1.002 + sigma_max = 1.0 + shift = 3 if shift is None else shift + num_train_timesteps = 1000 + sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength + sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps) + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + timesteps = sigmas * num_train_timesteps + return sigmas, timesteps + + @staticmethod + def set_timesteps_wan(num_inference_steps=100, denoising_strength=1.0, shift=None): + sigma_min = 0.0 + sigma_max = 1.0 + shift = 5 if shift is None else shift + num_train_timesteps = 1000 + sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength + sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1] + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + timesteps = sigmas * num_train_timesteps + return sigmas, timesteps + + @staticmethod + def _calculate_shift_qwen_image(image_seq_len, base_seq_len=256, max_seq_len=8192, base_shift=0.5, max_shift=0.9): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + @staticmethod + def set_timesteps_qwen_image(num_inference_steps=100, denoising_strength=1.0, exponential_shift_mu=None, dynamic_shift_len=None): + sigma_min = 0.0 + sigma_max = 1.0 + num_train_timesteps = 1000 + shift_terminal = 0.02 + # Sigmas + sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength + sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1] + # Mu + if exponential_shift_mu is not None: + mu = exponential_shift_mu + elif dynamic_shift_len is not None: + mu = FlowMatchScheduler._calculate_shift_qwen_image(dynamic_shift_len) + else: + mu = 0.8 + sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1)) + # Shift terminal + one_minus_z = 1 - sigmas + scale_factor = one_minus_z[-1] / (1 - shift_terminal) + sigmas = 1 - (one_minus_z / scale_factor) + # Timesteps + timesteps = sigmas * num_train_timesteps + return sigmas, timesteps + + @staticmethod + def compute_empirical_mu(image_seq_len, num_steps): + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + + if image_seq_len > 4300: + mu = a2 * image_seq_len + b2 + return float(mu) + + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + mu = a * num_steps + b + + return float(mu) + + @staticmethod + def set_timesteps_flux2(num_inference_steps=100, denoising_strength=1.0, dynamic_shift_len=1024//16*1024//16): + sigma_min = 1 / num_inference_steps + sigma_max = 1.0 + num_train_timesteps = 1000 + sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength + sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps) + mu = FlowMatchScheduler.compute_empirical_mu(dynamic_shift_len, num_inference_steps) + sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1)) + timesteps = sigmas * num_train_timesteps + return sigmas, timesteps + + @staticmethod + def set_timesteps_z_image(num_inference_steps=100, denoising_strength=1.0, shift=None, target_timesteps=None): + sigma_min = 0.0 + sigma_max = 1.0 + shift = 3 if shift is None else shift + num_train_timesteps = 1000 + sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength + sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1] + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + timesteps = sigmas * num_train_timesteps + if target_timesteps is not None: + target_timesteps = target_timesteps.to(dtype=timesteps.dtype, device=timesteps.device) + for timestep in target_timesteps: + timestep_id = torch.argmin((timesteps - timestep).abs()) + timesteps[timestep_id] = timestep + return sigmas, timesteps + + def set_training_weight(self): + steps = 1000 + x = self.timesteps + y = torch.exp(-2 * ((x - steps / 2) / steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (steps / y_shifted.sum()) + if len(self.timesteps) != 1000: + # This is an empirical formula. + bsmntw_weighing = bsmntw_weighing * (len(self.timesteps) / steps) + bsmntw_weighing = bsmntw_weighing + bsmntw_weighing[1] + self.linear_timesteps_weights = bsmntw_weighing + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, **kwargs): + self.sigmas, self.timesteps = self.set_timesteps_fn( + num_inference_steps=num_inference_steps, + denoising_strength=denoising_strength, + **kwargs, + ) + if training: + self.set_training_weight() + self.training = True + else: + self.training = False + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + weights = self.linear_timesteps_weights[timestep_id] + return weights diff --git a/diffsynth/diffusion/logger.py b/diffsynth/diffusion/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..217d36dea1d57bbdd6036b38b6c7cf0de80d51d8 --- /dev/null +++ b/diffsynth/diffusion/logger.py @@ -0,0 +1,43 @@ +import os, torch +from accelerate import Accelerator + + +class ModelLogger: + def __init__(self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x:x): + self.output_path = output_path + self.remove_prefix_in_ckpt = remove_prefix_in_ckpt + self.state_dict_converter = state_dict_converter + self.num_steps = 0 + + + def on_step_end(self, accelerator: Accelerator, model: torch.nn.Module, save_steps=None): + self.num_steps += 1 + if (save_steps is not None and self.num_steps % save_steps == 0) or self.num_steps == 1: + self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors") + + + def on_epoch_end(self, accelerator: Accelerator, model: torch.nn.Module, epoch_id): + accelerator.wait_for_everyone() + if accelerator.is_main_process: + state_dict = accelerator.get_state_dict(model) + state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt) + state_dict = self.state_dict_converter(state_dict) + os.makedirs(self.output_path, exist_ok=True) + path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors") + accelerator.save(state_dict, path, safe_serialization=True) + + + def on_training_end(self, accelerator: Accelerator, model: torch.nn.Module, save_steps=None): + if save_steps is not None and self.num_steps % save_steps != 0: + self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors") + + + def save_model(self, accelerator: Accelerator, model: torch.nn.Module, file_name): + accelerator.wait_for_everyone() + if accelerator.is_main_process: + state_dict = accelerator.get_state_dict(model) + state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt) + state_dict = self.state_dict_converter(state_dict) + os.makedirs(self.output_path, exist_ok=True) + path = os.path.join(self.output_path, file_name) + accelerator.save(state_dict, path, safe_serialization=True) diff --git a/diffsynth/diffusion/loss.py b/diffsynth/diffusion/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..b0e5a4fc5bbb1e770cad858109b01056723f6f81 --- /dev/null +++ b/diffsynth/diffusion/loss.py @@ -0,0 +1,120 @@ +from .base_pipeline import BasePipeline +import torch + + +def FlowMatchSFTLoss(pipe: BasePipeline, **inputs): + max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps)) + min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps)) + + timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,)) + timestep = pipe.scheduler.timesteps[timestep_id].to(dtype=pipe.torch_dtype, device=pipe.device) + + noise = torch.randn_like(inputs["input_latents"]) + inputs["latents"] = pipe.scheduler.add_noise(inputs["input_latents"], noise, timestep) + training_target = pipe.scheduler.training_target(inputs["input_latents"], noise, timestep) + + models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} + + noise_pred = pipe.model_fn(**models, **inputs, timestep=timestep) + + loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float()) + loss = loss * pipe.scheduler.training_weight(timestep) + return loss + + +def DirectDistillLoss(pipe: BasePipeline, **inputs): + pipe.scheduler.set_timesteps(inputs["num_inference_steps"]) + pipe.scheduler.training = True + models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} + for progress_id, timestep in enumerate(pipe.scheduler.timesteps): + timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) + noise_pred = pipe.model_fn(**models, **inputs, timestep=timestep, progress_id=progress_id) + inputs["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs) + loss = torch.nn.functional.mse_loss(inputs["latents"].float(), inputs["input_latents"].float()) + return loss + + +class TrajectoryImitationLoss(torch.nn.Module): + def __init__(self): + super().__init__() + self.initialized = False + + def initialize(self, device): + import lpips # TODO: remove it + self.loss_fn = lpips.LPIPS(net='alex').to(device) + self.initialized = True + + def fetch_trajectory(self, pipe: BasePipeline, timesteps_student, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): + trajectory = [inputs_shared["latents"].clone()] + + pipe.scheduler.set_timesteps(num_inference_steps, target_timesteps=timesteps_student) + models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} + for progress_id, timestep in enumerate(pipe.scheduler.timesteps): + timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) + noise_pred = pipe.cfg_guided_model_fn( + pipe.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared) + + trajectory.append(inputs_shared["latents"].clone()) + return pipe.scheduler.timesteps, trajectory + + def align_trajectory(self, pipe: BasePipeline, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): + loss = 0 + pipe.scheduler.set_timesteps(num_inference_steps, training=True) + models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} + for progress_id, timestep in enumerate(pipe.scheduler.timesteps): + timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) + + progress_id_teacher = torch.argmin((timesteps_teacher - timestep).abs()) + inputs_shared["latents"] = trajectory_teacher[progress_id_teacher] + + noise_pred = pipe.cfg_guided_model_fn( + pipe.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + + sigma = pipe.scheduler.sigmas[progress_id] + sigma_ = 0 if progress_id + 1 >= len(pipe.scheduler.timesteps) else pipe.scheduler.sigmas[progress_id + 1] + if progress_id + 1 >= len(pipe.scheduler.timesteps): + latents_ = trajectory_teacher[-1] + else: + progress_id_teacher = torch.argmin((timesteps_teacher - pipe.scheduler.timesteps[progress_id + 1]).abs()) + latents_ = trajectory_teacher[progress_id_teacher] + + target = (latents_ - inputs_shared["latents"]) / (sigma_ - sigma) + loss = loss + torch.nn.functional.mse_loss(noise_pred.float(), target.float()) * pipe.scheduler.training_weight(timestep) + return loss + + def compute_regularization(self, pipe: BasePipeline, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): + inputs_shared["latents"] = trajectory_teacher[0] + pipe.scheduler.set_timesteps(num_inference_steps) + models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} + for progress_id, timestep in enumerate(pipe.scheduler.timesteps): + timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) + noise_pred = pipe.cfg_guided_model_fn( + pipe.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared) + + image_pred = pipe.vae_decoder(inputs_shared["latents"]) + image_real = pipe.vae_decoder(trajectory_teacher[-1]) + loss = self.loss_fn(image_pred.float(), image_real.float()) + return loss + + def forward(self, pipe: BasePipeline, inputs_shared, inputs_posi, inputs_nega): + if not self.initialized: + self.initialize(pipe.device) + with torch.no_grad(): + pipe.scheduler.set_timesteps(8) + timesteps_teacher, trajectory_teacher = self.fetch_trajectory(inputs_shared["teacher"], pipe.scheduler.timesteps, inputs_shared, inputs_posi, inputs_nega, 50, 2) + timesteps_teacher = timesteps_teacher.to(dtype=pipe.torch_dtype, device=pipe.device) + loss_1 = self.align_trajectory(pipe, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1) + loss_2 = self.compute_regularization(pipe, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1) + loss = loss_1 + loss_2 + return loss diff --git a/diffsynth/diffusion/parsers.py b/diffsynth/diffusion/parsers.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c6c6afdd3868aab45b1b51a3a47fe2e37d77f4 --- /dev/null +++ b/diffsynth/diffusion/parsers.py @@ -0,0 +1,70 @@ +import argparse + + +def add_dataset_base_config(parser: argparse.ArgumentParser): + parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.") + parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.") + parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.") + parser.add_argument("--dataset_num_workers", type=int, default=0, help="Number of workers for data loading.") + parser.add_argument("--data_file_keys", type=str, default="image,video", help="Data file keys in the metadata. Comma-separated.") + return parser + +def add_image_size_config(parser: argparse.ArgumentParser): + parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.") + parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.") + parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution.") + return parser + +def add_video_size_config(parser: argparse.ArgumentParser): + parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.") + parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.") + parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution.") + parser.add_argument("--num_frames", type=int, default=81, help="Number of frames per video. Frames are sampled from the video prefix.") + return parser + +def add_model_config(parser: argparse.ArgumentParser): + parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.") + parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.") + parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.") + parser.add_argument("--fp8_models", default=None, help="Models with FP8 precision, comma-separated.") + parser.add_argument("--offload_models", default=None, help="Models with offload, comma-separated. Only used in splited training.") + return parser + +def add_training_config(parser: argparse.ArgumentParser): + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.") + parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.") + parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.") + parser.add_argument("--find_unused_parameters", default=False, action="store_true", help="Whether to find unused parameters in DDP.") + parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay.") + parser.add_argument("--task", type=str, default="sft", required=False, help="Task type.") + return parser + +def add_output_config(parser: argparse.ArgumentParser): + parser.add_argument("--output_path", type=str, default="./models", help="Output save path.") + parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.") + parser.add_argument("--save_steps", type=int, default=None, help="Number of checkpoint saving invervals. If None, checkpoints will be saved every epoch.") + return parser + +def add_lora_config(parser: argparse.ArgumentParser): + parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.") + parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.") + parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.") + parser.add_argument("--lora_checkpoint", type=str, default=None, help="Path to the LoRA checkpoint. If provided, LoRA will be loaded from this checkpoint.") + parser.add_argument("--preset_lora_path", type=str, default=None, help="Path to the preset LoRA checkpoint. If provided, this LoRA will be fused to the base model.") + parser.add_argument("--preset_lora_model", type=str, default=None, help="Which model the preset LoRA is fused to.") + return parser + +def add_gradient_config(parser: argparse.ArgumentParser): + parser.add_argument("--use_gradient_checkpointing", default=False, action="store_true", help="Whether to use gradient checkpointing.") + parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.") + return parser + +def add_general_config(parser: argparse.ArgumentParser): + parser = add_dataset_base_config(parser) + parser = add_model_config(parser) + parser = add_training_config(parser) + parser = add_output_config(parser) + parser = add_lora_config(parser) + parser = add_gradient_config(parser) + return parser diff --git a/diffsynth/diffusion/runner.py b/diffsynth/diffusion/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..09b2a1495846974c3ea2fd8ec3420d5f80719977 --- /dev/null +++ b/diffsynth/diffusion/runner.py @@ -0,0 +1,131 @@ +import os, torch +from tqdm import tqdm +from accelerate import Accelerator +from .training_module import DiffusionTrainingModule +from src.training_module import MetaViewTrainingModule +from .logger import ModelLogger + +from PIL import Image + + +def collate_fn(batch): + if len(batch) == 0: + return {} + + collated = {} + keys = batch[0].keys() + for key in keys: + values = [sample[key] for sample in batch] + first_val = values[0] + if isinstance(first_val, torch.Tensor): + # 对于 Tensor，使用 torch.stack 沿第0维堆叠（要求所有张量形状相同） + collated[key] = torch.stack(values, dim=0) + elif isinstance(first_val, Image.Image): + collated[key] = values + elif isinstance(first_val, str): + collated[key] = values + else: + collated[key] = values + + return collated + +def launch_training_task( + accelerator: Accelerator, + dataset: torch.utils.data.Dataset, + model: MetaViewTrainingModule, + model_logger: ModelLogger, + learning_rate: float = 1e-5, + weight_decay: float = 1e-2, + num_workers: int = 1, + save_steps: int = None, + num_epochs: int = 1, + batch_size: int = 1, + args = None, +): + if args is not None: + learning_rate = args.learning_rate + weight_decay = args.weight_decay + num_workers = args.dataset_num_workers + save_steps = args.save_steps + num_epochs = args.num_epochs + + optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay) + scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer) + # dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, num_workers=num_workers) + + model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler) + + def grad_hook(name): + def hook(grad): + # # print(f"{name} gradient norm:{grad.norm().item():.6f}") + # if grad.norm().item() > 1: + # print(f" gradient over 1: {name} {grad.norm().item():.6f}") + if torch.isnan(grad).any(): + print(f"!!! NaN gradient: {name}") + if torch.isinf(grad).any(): + print(f"!!! Inf gradient: {name}") + return grad + return hook + + for name, param in model.named_parameters(): + if param.requires_grad: + param.register_hook(grad_hook(name)) + + NaN_step = 0 + for epoch_id in range(num_epochs): + for data in tqdm(dataloader): + with accelerator.accumulate(model): + # print(type(data)) + # print(data["prompt"]) + optimizer.zero_grad() + if dataset.load_from_cache: + loss = model({}, inputs=data) + else: + loss = model(data) + + if torch.isnan(loss).any(): + print(f"!!! Loss is NaN at step {model_logger.num_steps}! Skipping...") + NaN_step += 1 + print(data["name"]) + exit(0) + + accelerator.backward(loss) + + max_norm = 5.0 + if accelerator.sync_gradients: + grad_norm = accelerator.clip_grad_norm_(model.parameters(), max_norm=max_norm) + if accelerator.is_main_process: + if grad_norm > 5.0: + print(f"gradient over 5: {grad_norm:.4f}") + + optimizer.step() + model_logger.on_step_end(accelerator, model, save_steps) + scheduler.step() + if save_steps is None: + model_logger.on_epoch_end(accelerator, model, epoch_id) + model_logger.on_training_end(accelerator, model, save_steps) + + +def launch_data_process_task( + accelerator: Accelerator, + dataset: torch.utils.data.Dataset, + model: DiffusionTrainingModule, + model_logger: ModelLogger, + num_workers: int = 8, + args = None, +): + if args is not None: + num_workers = args.dataset_num_workers + + dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers) + model, dataloader = accelerator.prepare(model, dataloader) + + for data_id, data in enumerate(tqdm(dataloader)): + with accelerator.accumulate(model): + with torch.no_grad(): + folder = os.path.join(model_logger.output_path, str(accelerator.process_index)) + os.makedirs(folder, exist_ok=True) + save_path = os.path.join(model_logger.output_path, str(accelerator.process_index), f"{data_id}.pth") + data = model(data) + torch.save(data, save_path) diff --git a/diffsynth/diffusion/training_module.py b/diffsynth/diffusion/training_module.py new file mode 100644 index 0000000000000000000000000000000000000000..e9d3ccd01f803f6354252c4c0af31c6509ec98db --- /dev/null +++ b/diffsynth/diffusion/training_module.py @@ -0,0 +1,214 @@ +import torch, json +from ..core import ModelConfig, load_state_dict +from ..utils.controlnet import ControlNetInput +from peft import LoraConfig, inject_adapter_in_model + + +class DiffusionTrainingModule(torch.nn.Module): + def __init__(self): + super().__init__() + + + def to(self, *args, **kwargs): + for name, model in self.named_children(): + model.to(*args, **kwargs) + return self + + + def trainable_modules(self): + trainable_modules = filter(lambda p: p.requires_grad, self.parameters()) + return trainable_modules + + + def trainable_param_names(self): + trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters())) + trainable_param_names = set([named_param[0] for named_param in trainable_param_names]) + return trainable_param_names + + + def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None): + if lora_alpha is None: + lora_alpha = lora_rank + if isinstance(target_modules, list) and len(target_modules) == 1: + target_modules = target_modules[0] + lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules) + model = inject_adapter_in_model(lora_config, model) + if upcast_dtype is not None: + for param in model.parameters(): + if param.requires_grad: + param.data = param.to(upcast_dtype) + return model + + + def mapping_lora_state_dict(self, state_dict): + new_state_dict = {} + for key, value in state_dict.items(): + if "lora_A.weight" in key or "lora_B.weight" in key: + new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight") + new_state_dict[new_key] = value + elif "lora_A.default.weight" in key or "lora_B.default.weight" in key: + new_state_dict[key] = value + return new_state_dict + + + def export_trainable_state_dict(self, state_dict, remove_prefix=None): + trainable_param_names = self.trainable_param_names() + state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names} + if remove_prefix is not None: + state_dict_ = {} + for name, param in state_dict.items(): + if name.startswith(remove_prefix): + name = name[len(remove_prefix):] + state_dict_[name] = param + state_dict = state_dict_ + return state_dict + + + def transfer_data_to_device(self, data, device, torch_float_dtype=None): + if data is None: + return data + elif isinstance(data, torch.Tensor): + data = data.to(device) + if torch_float_dtype is not None and data.dtype in [torch.float, torch.float16, torch.bfloat16]: + data = data.to(torch_float_dtype) + return data + elif isinstance(data, tuple): + data = tuple(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data) + return data + elif isinstance(data, list): + data = list(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data) + return data + elif isinstance(data, dict): + data = {i: self.transfer_data_to_device(data[i], device, torch_float_dtype) for i in data} + return data + else: + return data + + def parse_vram_config(self, fp8=False, offload=False, device="cpu"): + if fp8: + return { + "offload_dtype": torch.float8_e4m3fn, + "offload_device": device, + "onload_dtype": torch.float8_e4m3fn, + "onload_device": device, + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": device, + "computation_dtype": torch.bfloat16, + "computation_device": device, + } + elif offload: + return { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": device, + "computation_dtype": torch.bfloat16, + "computation_device": device, + "clear_parameters": True, + } + else: + return {} + + def parse_model_configs(self, model_paths, model_id_with_origin_paths, fp8_models=None, offload_models=None, device="cpu"): + fp8_models = [] if fp8_models is None else fp8_models.split(",") + offload_models = [] if offload_models is None else offload_models.split(",") + model_configs = [] + if model_paths is not None: + model_paths = json.loads(model_paths) + for path in model_paths: + vram_config = self.parse_vram_config( + fp8=path in fp8_models, + offload=path in offload_models, + device=device + ) + model_configs.append(ModelConfig(path=path, **vram_config)) + if model_id_with_origin_paths is not None: + model_id_with_origin_paths = model_id_with_origin_paths.split(",") + for model_id_with_origin_path in model_id_with_origin_paths: + model_id, origin_file_pattern = model_id_with_origin_path.split(":") + vram_config = self.parse_vram_config( + fp8=model_id_with_origin_path in fp8_models, + offload=model_id_with_origin_path in offload_models, + device=device + ) + model_configs.append(ModelConfig(model_id=model_id, origin_file_pattern=origin_file_pattern, **vram_config)) + return model_configs + + + def switch_pipe_to_training_mode( + self, + pipe, + trainable_models=None, + lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None, + preset_lora_path=None, preset_lora_model=None, + task="sft", + ): + # Scheduler + pipe.scheduler.set_timesteps(1000, training=True) + + # Freeze untrainable models + pipe.freeze_except([] if trainable_models is None else trainable_models.split(",")) + + # Preset LoRA + if preset_lora_path is not None: + pipe.load_lora(getattr(pipe, preset_lora_model), preset_lora_path) + + # FP8 + # FP8 relies on a model-specific memory management scheme. + # It is delegated to the subclass. + + # Add LoRA to the base models + if lora_base_model is not None and not task.endswith(":data_process"): + if (not hasattr(pipe, lora_base_model)) or getattr(pipe, lora_base_model) is None: + print(f"No {lora_base_model} models in the pipeline. We cannot patch LoRA on the model. If this occurs during the data processing stage, it is normal.") + return + model = self.add_lora_to_model( + getattr(pipe, lora_base_model), + target_modules=lora_target_modules.split(","), + lora_rank=lora_rank, + upcast_dtype=pipe.torch_dtype, + ) + if lora_checkpoint is not None: + state_dict = load_state_dict(lora_checkpoint) + state_dict = self.mapping_lora_state_dict(state_dict) + load_result = model.load_state_dict(state_dict, strict=False) + print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys") + if len(load_result[1]) > 0: + print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}") + setattr(pipe, lora_base_model, model) + + + def split_pipeline_units(self, task, pipe, trainable_models=None, lora_base_model=None): + models_require_backward = [] + if trainable_models is not None: + models_require_backward += trainable_models.split(",") + if lora_base_model is not None: + models_require_backward += [lora_base_model] + if task.endswith(":data_process"): + print("!!!! task.endswith(:data_process)", ) + _, pipe.units = pipe.split_pipeline_units(models_require_backward) + elif task.endswith(":train"): + print("!!!! task.endswith(:train)", ) + pipe.units, _ = pipe.split_pipeline_units(models_require_backward) + return pipe + + def parse_extra_inputs(self, data, extra_inputs, inputs_shared): + controlnet_keys_map = ( + ("blockwise_controlnet_", "blockwise_controlnet_inputs",), + ("controlnet_", "controlnet_inputs"), + ) + controlnet_inputs = {} + for extra_input in extra_inputs: + for prefix, name in controlnet_keys_map: + if extra_input.startswith(prefix): + if name not in controlnet_inputs: + controlnet_inputs[name] = {} + controlnet_inputs[name][extra_input.replace(prefix, "")] = data[extra_input] + break + else: + inputs_shared[extra_input] = data[extra_input] + for name, params in controlnet_inputs.items(): + inputs_shared[name] = [ControlNetInput(**params)] + return inputs_shared diff --git a/diffsynth/models/flux2_dit.py b/diffsynth/models/flux2_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..a08c579a77c8aa1c2b35c0c168517335a636097a --- /dev/null +++ b/diffsynth/models/flux2_dit.py @@ -0,0 +1,1057 @@ +import inspect +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch, math +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..core.attention import attention_forward +from ..core.gradient import gradient_checkpoint_forward + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + + Args + timesteps (torch.Tensor): + a 1-D Tensor of N indices, one per batch element. These may be fractional. + embedding_dim (int): + the dimension of the output. + flip_sin_to_cos (bool): + Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) + downscale_freq_shift (float): + Controls the delta between frequencies between dimensions + scale (float): + Scaling factor applied to the embeddings. + max_period (int): + Controls the maximum frequency of the embeddings + Returns + torch.Tensor: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + sample_proj_bias=True, + ): + super().__init__() + + self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = torch.nn.SiLU() + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Timesteps(nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + return t_emb + + +class AdaLayerNormContinuous(nn.Module): + r""" + Adaptive normalization layer with a norm layer (layer_norm or rms_norm). + + Args: + embedding_dim (`int`): Embedding dimension to use during projection. + conditioning_embedding_dim (`int`): Dimension of the input condition. + elementwise_affine (`bool`, defaults to `True`): + Boolean flag to denote if affine transformation should be applied. + eps (`float`, defaults to 1e-5): Epsilon factor. + bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use. + norm_type (`str`, defaults to `"layer_norm"`): + Normalization layer to use. Values supported: "layer_norm", "rms_norm". + """ + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters + # because the output is immediately scaled and shifted by the projected conditioning embeddings. + # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. + # However, this is how it was implemented in the original code, and it's rather likely you should + # set `elementwise_affine` to False. + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) + emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[np.ndarray, int], + theta: float = 10000.0, + use_real=False, + linear_factor=1.0, + ntk_factor=1.0, + repeat_interleave_real=True, + freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux) +): + """ + Precompute the frequency tensor for complex exponentials (cis) with given dimensions. + + This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end + index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64 + data type. + + Args: + dim (`int`): Dimension of the frequency tensor. + pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar + theta (`float`, *optional*, defaults to 10000.0): + Scaling factor for frequency computation. Defaults to 10000.0. + use_real (`bool`, *optional*): + If True, return real part and imaginary part separately. Otherwise, return complex numbers. + linear_factor (`float`, *optional*, defaults to 1.0): + Scaling factor for the context extrapolation. Defaults to 1.0. + ntk_factor (`float`, *optional*, defaults to 1.0): + Scaling factor for the NTK-Aware RoPE. Defaults to 1.0. + repeat_interleave_real (`bool`, *optional*, defaults to `True`): + If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`. + Otherwise, they are concateanted with themselves. + freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`): + the dtype of the frequency tensor. + Returns: + `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2] + """ + assert dim % 2 == 0 + + if isinstance(pos, int): + pos = torch.arange(pos) + if isinstance(pos, np.ndarray): + pos = torch.from_numpy(pos) # type: ignore # [S] + + theta = theta * ntk_factor + freqs = ( + 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) / dim)) / linear_factor + ) # [D/2] + freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2] + is_npu = freqs.device.type == "npu" + if is_npu: + freqs = freqs.float() + if use_real and repeat_interleave_real: + # flux, hunyuan-dit, cogvideox + freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] + return freqs_cos, freqs_sin + elif use_real: + # stable audio, allegro + freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D] + freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D] + return freqs_cos, freqs_sin + else: + # lumina + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2] + return freqs_cis + + +def apply_rotary_emb( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + use_real: bool = True, + use_real_unbind_dim: int = -1, + sequence_dim: int = 2, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings + to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are + reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting + tensors contain rotary embeddings and are returned as real tensors. + + Args: + x (`torch.Tensor`): + Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply + freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + if sequence_dim == 2: + cos = cos[None, None, :, :] + sin = sin[None, None, :, :] + elif sequence_dim == 1: + cos = cos[None, :, None, :] + sin = sin[None, :, None, :] + else: + raise ValueError(f"`sequence_dim={sequence_dim}` but should be 1 or 2.") + + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, S, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, H, S, D//2] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + # used for lumina + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + +def _get_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None): + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + encoder_query = encoder_key = encoder_value = None + if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_fused_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None): + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + + encoder_query = encoder_key = encoder_value = (None,) + if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"): + encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_qkv_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None): + return _get_projections(attn, hidden_states, encoder_hidden_states) + + +class Flux2SwiGLU(nn.Module): + """ + Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection + layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters. + """ + + def __init__(self): + super().__init__() + self.gate_fn = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + x = self.gate_fn(x1) * x2 + return x + + +class Flux2FeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: float = 3.0, + inner_dim: Optional[int] = None, + bias: bool = False, + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out or dim + + # Flux2SwiGLU will reduce the dimension by half + self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias) + self.act_fn = Flux2SwiGLU() + self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear_in(x) + x = self.act_fn(x) + x = self.linear_out(x) + return x + + +class Flux2AttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") + + def __call__( + self, + attn: "Flux2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( + attn, hidden_states, encoder_hidden_states + ) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if attn.added_kv_proj_dim is not None: + encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + hidden_states = attention_forward( + query, + key, + value, + q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d", + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if encoder_hidden_states is not None: + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 + ) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + if encoder_hidden_states is not None: + return hidden_states, encoder_hidden_states + else: + return hidden_states + + +class Flux2Attention(torch.nn.Module): + _default_processor_cls = Flux2AttnProcessor + _available_processors = [Flux2AttnProcessor] + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: Optional[int] = None, + added_proj_bias: Optional[bool] = True, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + elementwise_affine: bool = True, + processor=None, + ): + super().__init__() + + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.heads = out_dim // dim_head if out_dim is not None else heads + + self.use_bias = bias + self.dropout = dropout + + self.added_kv_proj_dim = added_kv_proj_dim + self.added_proj_bias = added_proj_bias + + self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + + # QK Norm + self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + + self.to_out = torch.nn.ModuleList([]) + self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(torch.nn.Dropout(dropout)) + + if added_kv_proj_dim is not None: + self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps) + self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps) + self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias) + + if processor is None: + processor = self._default_processor_cls() + self.processor = processor + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs) + + +class Flux2ParallelSelfAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") + + def __call__( + self, + attn: "Flux2ParallelSelfAttention", + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # Parallel in (QKV + MLP in) projection + hidden_states = attn.to_qkv_mlp_proj(hidden_states) + qkv, mlp_hidden_states = torch.split( + hidden_states, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1 + ) + + # Handle the attention logic + query, key, value = qkv.chunk(3, dim=-1) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + hidden_states = attention_forward( + query, + key, + value, + q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d", + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # Handle the feedforward (FF) logic + mlp_hidden_states = attn.mlp_act_fn(mlp_hidden_states) + + # Concatenate and parallel output projection + hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) + hidden_states = attn.to_out(hidden_states) + + return hidden_states + + +class Flux2ParallelSelfAttention(torch.nn.Module): + """ + Flux 2 parallel self-attention for the Flux 2 single-stream transformer blocks. + + This implements a parallel transformer block, where the attention QKV projections are fused to the feedforward (FF) + input projections, and the attention output projections are fused to the FF output projections. See the [ViT-22B + paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block. + """ + + _default_processor_cls = Flux2ParallelSelfAttnProcessor + _available_processors = [Flux2ParallelSelfAttnProcessor] + # Does not support QKV fusion as the QKV projections are always fused + _supports_qkv_fusion = False + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + elementwise_affine: bool = True, + mlp_ratio: float = 4.0, + mlp_mult_factor: int = 2, + processor=None, + ): + super().__init__() + + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.heads = out_dim // dim_head if out_dim is not None else heads + + self.use_bias = bias + self.dropout = dropout + + self.mlp_ratio = mlp_ratio + self.mlp_hidden_dim = int(query_dim * self.mlp_ratio) + self.mlp_mult_factor = mlp_mult_factor + + # Fused QKV projections + MLP input projection + self.to_qkv_mlp_proj = torch.nn.Linear( + self.query_dim, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias + ) + self.mlp_act_fn = Flux2SwiGLU() + + # QK Norm + self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + + # Fused attention output projection + MLP output projection + self.to_out = torch.nn.Linear(self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias) + + if processor is None: + processor = self._default_processor_cls() + self.processor = processor + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + return self.processor(self, hidden_states, attention_mask, image_rotary_emb, **kwargs) + + +class Flux2SingleTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + ): + super().__init__() + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + # Note that the MLP in/out linear layers are fused with the attention QKV/out projections, respectively; this + # is often called a "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442) + # for a visual depiction of this type of transformer block. + self.attn = Flux2ParallelSelfAttention( + query_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=bias, + out_bias=bias, + eps=eps, + mlp_ratio=mlp_ratio, + mlp_mult_factor=2, + processor=Flux2ParallelSelfAttnProcessor(), + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor], + temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + split_hidden_states: bool = False, + text_seq_len: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already + # concatenated + if encoder_hidden_states is not None: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + mod_shift, mod_scale, mod_gate = temb_mod_params + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift + + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = hidden_states + mod_gate * attn_output + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + if split_hidden_states: + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states + else: + return hidden_states + + +class Flux2TransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + ): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + self.attn = Flux2Attention( + query_dim=dim, + added_kv_proj_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=bias, + added_proj_bias=bias, + out_bias=bias, + eps=eps, + processor=Flux2AttnProcessor(), + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb_mod_params_img: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], + temb_mod_params_txt: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + joint_attention_kwargs = joint_attention_kwargs or {} + + # Modulation parameters shape: [1, 1, self.dim] + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img + (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt + + # Img stream + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + + # Conditioning txt stream + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) + norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa + + # Attention on concatenated img + txt stream + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + attn_output, context_attn_output = attention_outputs + + # Process attention outputs for the image stream (`hidden_states`). + attn_output = gate_msa * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + ff_output = self.ff(norm_hidden_states) + hidden_states = hidden_states + gate_mlp * ff_output + + # Process attention outputs for the text stream (`encoder_hidden_states`). + context_attn_output = c_gate_msa * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class Flux2PosEmbed(nn.Module): + # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 + def __init__(self, theta: int, axes_dim: List[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + # Expected ids shape: [S, len(self.axes_dim)] + cos_out = [] + sin_out = [] + pos = ids.float() + is_mps = ids.device.type == "mps" + is_npu = ids.device.type == "npu" + freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1] + for i in range(len(self.axes_dim)): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[..., i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class Flux2TimestepGuidanceEmbeddings(nn.Module): + def __init__(self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False): + super().__init__() + + self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding( + in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias + ) + + self.guidance_embedder = TimestepEmbedding( + in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias + ) + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) # (N, D) + + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) # (N, D) + + time_guidance_emb = timesteps_emb + guidance_emb + + return time_guidance_emb + + +class Flux2Modulation(nn.Module): + def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): + super().__init__() + self.mod_param_sets = mod_param_sets + + self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias) + self.act_fn = nn.SiLU() + + def forward(self, temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + mod = self.act_fn(temb) + mod = self.linear(mod) + + if mod.ndim == 2: + mod = mod.unsqueeze(1) + mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) + # Return tuple of 3-tuples of modulation params shift/scale/gate + return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets)) + + +class Flux2DiT(torch.nn.Module): + def __init__( + self, + patch_size: int = 1, + in_channels: int = 128, + out_channels: Optional[int] = None, + num_layers: int = 8, + num_single_layers: int = 48, + attention_head_dim: int = 128, + num_attention_heads: int = 48, + joint_attention_dim: int = 15360, + timestep_guidance_channels: int = 256, + mlp_ratio: float = 3.0, + axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32), + rope_theta: int = 2000, + eps: float = 1e-6, + ): + super().__init__() + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + # 1. Sinusoidal positional embedding for RoPE on image and text tokens + self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope) + + # 2. Combined timestep + guidance embedding + self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( + in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, bias=False + ) + + # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.) + # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks + self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) + self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) + # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream + self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False) + + # 4. Input projections + self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False) + self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False) + + # 5. Double Stream Transformer Blocks + self.transformer_blocks = nn.ModuleList( + [ + Flux2TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + ) + for _ in range(num_layers) + ] + ) + + # 6. Single Stream Transformer Blocks + self.single_transformer_blocks = nn.ModuleList( + [ + Flux2SingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + ) + for _ in range(num_single_layers) + ] + ) + + # 7. Output layers + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False + ) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + ) -> Union[torch.Tensor]: + """ + The [`FluxTransformer2DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Input `hidden_states`. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + timestep ( `torch.LongTensor`): + Used to indicate denoising step. + block_controlnet_hidden_states: (`list` of `torch.Tensor`): + A list of tensors that if specified are added to the residuals of transformer blocks. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # 0. Handle input arguments + if joint_attention_kwargs is not None: + joint_attention_kwargs = joint_attention_kwargs.copy() + lora_scale = joint_attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + num_txt_tokens = encoder_hidden_states.shape[1] + + # 1. Calculate timestep embedding and modulation parameters + timestep = timestep.to(hidden_states.dtype) * 1000 + guidance = guidance.to(hidden_states.dtype) * 1000 + + temb = self.time_guidance_embed(timestep, guidance) + + double_stream_mod_img = self.double_stream_modulation_img(temb) + double_stream_mod_txt = self.double_stream_modulation_txt(temb) + single_stream_mod = self.single_stream_modulation(temb)[0] + + # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states) + hidden_states = self.x_embedder(hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + # 3. Calculate RoPE embeddings from image and text tokens + # NOTE: the below logic means that we can't support batched inference with images of different resolutions or + # text prompts of differents lengths. Is this a use case we want to support? + if img_ids.ndim == 3: + img_ids = img_ids[0] + if txt_ids.ndim == 3: + txt_ids = txt_ids[0] + + image_rotary_emb = self.pos_embed(img_ids) + text_rotary_emb = self.pos_embed(txt_ids) + concat_rotary_emb = ( + torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), + torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), + ) + + # 4. Double Stream Transformer Blocks + for index_block, block in enumerate(self.transformer_blocks): + encoder_hidden_states, hidden_states = gradient_checkpoint_forward( + block, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb_mod_params_img=double_stream_mod_img, + temb_mod_params_txt=double_stream_mod_txt, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + # Concatenate text and image streams for single-block inference + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # 5. Single Stream Transformer Blocks + for index_block, block in enumerate(self.single_transformer_blocks): + hidden_states = gradient_checkpoint_forward( + block, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + hidden_states=hidden_states, + encoder_hidden_states=None, + temb_mod_params=single_stream_mod, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + # Remove text tokens from concatenated stream + hidden_states = hidden_states[:, num_txt_tokens:, ...] + + # 6. Output layers + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + return output diff --git a/diffsynth/models/flux2_text_encoder.py b/diffsynth/models/flux2_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..f3c68411f3160655ebefd49dfb6424b19373a301 --- /dev/null +++ b/diffsynth/models/flux2_text_encoder.py @@ -0,0 +1,58 @@ +from transformers import Mistral3ForConditionalGeneration, Mistral3Config + + +class Flux2TextEncoder(Mistral3ForConditionalGeneration): + def __init__(self): + config = Mistral3Config(**{ + "architectures": [ + "Mistral3ForConditionalGeneration" + ], + "dtype": "bfloat16", + "image_token_index": 10, + "model_type": "mistral3", + "multimodal_projector_bias": False, + "projector_hidden_act": "gelu", + "spatial_merge_size": 2, + "text_config": { + "attention_dropout": 0.0, + "dtype": "bfloat16", + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 5120, + "initializer_range": 0.02, + "intermediate_size": 32768, + "max_position_embeddings": 131072, + "model_type": "mistral", + "num_attention_heads": 32, + "num_hidden_layers": 40, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-05, + "rope_theta": 1000000000.0, + "sliding_window": None, + "use_cache": True, + "vocab_size": 131072 + }, + "transformers_version": "4.57.1", + "vision_config": { + "attention_dropout": 0.0, + "dtype": "bfloat16", + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 1540, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "pixtral", + "num_attention_heads": 16, + "num_channels": 3, + "num_hidden_layers": 24, + "patch_size": 14, + "rope_theta": 10000.0 + }, + "vision_feature_layer": -1 + }) + super().__init__(config) + + def forward(self, input_ids = None, pixel_values = None, attention_mask = None, position_ids = None, past_key_values = None, inputs_embeds = None, labels = None, use_cache = None, output_attentions = None, output_hidden_states = None, return_dict = None, cache_position = None, logits_to_keep = 0, image_sizes = None, **kwargs): + return super().forward(input_ids, pixel_values, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, logits_to_keep, image_sizes, **kwargs) + diff --git a/diffsynth/models/flux2_vae.py b/diffsynth/models/flux2_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..c7904b17618cc3c0811e42fde0f80ecd8f15f7ee --- /dev/null +++ b/diffsynth/models/flux2_vae.py @@ -0,0 +1,2322 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from typing import Dict, Optional, Tuple, Union, Callable + +import torch +import torch.nn as nn +from einops import rearrange +import torch.nn.functional as F +import inspect + +ACT2CLS = { + "swish": nn.SiLU, + "silu": nn.SiLU, + "mish": nn.Mish, + "gelu": nn.GELU, + "relu": nn.ReLU, +} + +def get_activation(act_fn: str) -> nn.Module: + """Helper function to get activation function from string. + + Args: + act_fn (str): Name of activation function. + + Returns: + nn.Module: Activation function. + """ + + act_fn = act_fn.lower() + if act_fn in ACT2CLS: + return ACT2CLS[act_fn]() + else: + raise ValueError(f"activation function {act_fn} not found in ACT2FN mapping {list(ACT2CLS.keys())}") + +class ResnetBlock2D(nn.Module): + r""" + A Resnet block. + + Parameters: + in_channels (`int`): The number of channels in the input. + out_channels (`int`, *optional*, default to be `None`): + The number of output channels for the first conv2d layer. If None, same as `in_channels`. + dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. + temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. + groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. + groups_out (`int`, *optional*, default to None): + The number of groups to use for the second normalization layer. if set to None, same as `groups`. + eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. + non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. + time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config. + By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" for a + stronger conditioning with scale and shift. + kernel (`torch.Tensor`, optional, default to None): FIR filter, see + [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. + output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. + use_in_shortcut (`bool`, *optional*, default to `True`): + If `True`, add a 1x1 nn.conv2d layer for skip-connection. + up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. + down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. + conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the + `conv_shortcut` output. + conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. + If None, same as `out_channels`. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + time_embedding_norm: str = "default", # default, scale_shift, + kernel: Optional[torch.Tensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, + conv_shortcut_bias: bool = True, + conv_2d_out_channels: Optional[int] = None, + ): + super().__init__() + if time_embedding_norm == "ada_group": + raise ValueError( + "This class cannot be used with `time_embedding_norm==ada_group`, please use `ResnetBlockCondNorm2D` instead", + ) + if time_embedding_norm == "spatial": + raise ValueError( + "This class cannot be used with `time_embedding_norm==spatial`, please use `ResnetBlockCondNorm2D` instead", + ) + + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.up = up + self.down = down + self.output_scale_factor = output_scale_factor + self.time_embedding_norm = time_embedding_norm + self.skip_time_act = skip_time_act + + if groups_out is None: + groups_out = groups + + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + self.time_emb_proj = nn.Linear(temb_channels, out_channels) + elif self.time_embedding_norm == "scale_shift": + self.time_emb_proj = nn.Linear(temb_channels, 2 * out_channels) + else: + raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ") + else: + self.time_emb_proj = None + + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) + + self.dropout = torch.nn.Dropout(dropout) + conv_2d_out_channels = conv_2d_out_channels or out_channels + self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1) + + self.nonlinearity = get_activation(non_linearity) + + self.upsample = self.downsample = None + if self.up: + if kernel == "fir": + fir_kernel = (1, 3, 3, 1) + self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel) + elif kernel == "sde_vp": + self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest") + else: + self.upsample = Upsample2D(in_channels, use_conv=False) + elif self.down: + if kernel == "fir": + fir_kernel = (1, 3, 3, 1) + self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel) + elif kernel == "sde_vp": + self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2) + else: + self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op") + + self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = nn.Conv2d( + in_channels, + conv_2d_out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=conv_shortcut_bias, + ) + + def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." + deprecate("scale", "1.0.0", deprecation_message) + + hidden_states = input_tensor + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor) + hidden_states = self.upsample(hidden_states) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor) + hidden_states = self.downsample(hidden_states) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb)[:, :, None, None] + + if self.time_embedding_norm == "default": + if temb is not None: + hidden_states = hidden_states + temb + hidden_states = self.norm2(hidden_states) + elif self.time_embedding_norm == "scale_shift": + if temb is None: + raise ValueError( + f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}" + ) + time_scale, time_shift = torch.chunk(temb, 2, dim=1) + hidden_states = self.norm2(hidden_states) + hidden_states = hidden_states * (1 + time_scale) + time_shift + else: + hidden_states = self.norm2(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor.contiguous()) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + +class Downsample2D(nn.Module): + """A 2D downsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + padding (`int`, default `1`): + padding for the convolution. + name (`str`, default `conv`): + name of the downsampling 2D layer. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + kernel_size=3, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + self.name = name + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + if use_conv: + conv = nn.Conv2d( + self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias + ) + else: + assert self.channels == self.out_channels + conv = nn.AvgPool2d(kernel_size=stride, stride=stride) + + # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." + deprecate("scale", "1.0.0", deprecation_message) + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + if self.use_conv and self.padding == 0: + pad = (0, 1, 0, 1) + hidden_states = F.pad(hidden_states, pad, mode="constant", value=0) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states) + + return hidden_states + +class Upsample2D(nn.Module): + """A 2D upsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + use_conv_transpose (`bool`, default `False`): + option to use a convolution transpose. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + name (`str`, default `conv`): + name of the upsampling 2D layer. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + kernel_size: Optional[int] = None, + padding=1, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + interpolate=True, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + self.interpolate = interpolate + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + conv = None + if use_conv_transpose: + if kernel_size is None: + kernel_size = 4 + conv = nn.ConvTranspose2d( + channels, self.out_channels, kernel_size=kernel_size, stride=2, padding=padding, bias=bias + ) + elif use_conv: + if kernel_size is None: + kernel_size = 3 + conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=kernel_size, padding=padding, bias=bias) + + # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." + deprecate("scale", "1.0.0", deprecation_message) + + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + if self.use_conv_transpose: + return self.conv(hidden_states) + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 until PyTorch 2.1 + # https://github.com/pytorch/pytorch/issues/86679#issuecomment-1783978767 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if self.interpolate: + # upsample_nearest_nhwc also fails when the number of output elements is large + # https://github.com/pytorch/pytorch/issues/141831 + scale_factor = ( + 2 if output_size is None else max([f / s for f, s in zip(output_size, hidden_states.shape[-2:])]) + ) + if hidden_states.numel() * scale_factor > pow(2, 31): + hidden_states = hidden_states.contiguous() + + if output_size is None: + hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest") + else: + hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest") + + # Cast back to original dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + + +class Attention(nn.Module): + r""" + A cross attention layer. + + Parameters: + query_dim (`int`): + The number of channels in the query. + cross_attention_dim (`int`, *optional*): + The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. + heads (`int`, *optional*, defaults to 8): + The number of heads to use for multi-head attention. + kv_heads (`int`, *optional*, defaults to `None`): + The number of key and value heads to use for multi-head attention. Defaults to `heads`. If + `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi + Query Attention (MQA) otherwise GQA is used. + dim_head (`int`, *optional*, defaults to 64): + The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): + The dropout probability to use. + bias (`bool`, *optional*, defaults to False): + Set to `True` for the query, key, and value linear layers to contain a bias parameter. + upcast_attention (`bool`, *optional*, defaults to False): + Set to `True` to upcast the attention computation to `float32`. + upcast_softmax (`bool`, *optional*, defaults to False): + Set to `True` to upcast the softmax computation to `float32`. + cross_attention_norm (`str`, *optional*, defaults to `None`): + The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`. + cross_attention_norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups to use for the group norm in the cross attention. + added_kv_proj_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the added key and value projections. If `None`, no projection is used. + norm_num_groups (`int`, *optional*, defaults to `None`): + The number of groups to use for the group norm in the attention. + spatial_norm_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the spatial normalization. + out_bias (`bool`, *optional*, defaults to `True`): + Set to `True` to use a bias in the output linear layer. + scale_qk (`bool`, *optional*, defaults to `True`): + Set to `True` to scale the query and key by `1 / sqrt(dim_head)`. + only_cross_attention (`bool`, *optional*, defaults to `False`): + Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if + `added_kv_proj_dim` is not `None`. + eps (`float`, *optional*, defaults to 1e-5): + An additional value added to the denominator in group normalization that is used for numerical stability. + rescale_output_factor (`float`, *optional*, defaults to 1.0): + A factor to rescale the output by dividing it with this value. + residual_connection (`bool`, *optional*, defaults to `False`): + Set to `True` to add the residual connection to the output. + _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`): + Set to `True` if the attention block is loaded from a deprecated state dict. + processor (`AttnProcessor`, *optional*, defaults to `None`): + The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and + `AttnProcessor` otherwise. + """ + + def __init__( + self, + query_dim: int, + cross_attention_dim: Optional[int] = None, + heads: int = 8, + kv_heads: Optional[int] = None, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + upcast_attention: bool = False, + upcast_softmax: bool = False, + cross_attention_norm: Optional[str] = None, + cross_attention_norm_num_groups: int = 32, + qk_norm: Optional[str] = None, + added_kv_proj_dim: Optional[int] = None, + added_proj_bias: Optional[bool] = True, + norm_num_groups: Optional[int] = None, + spatial_norm_dim: Optional[int] = None, + out_bias: bool = True, + scale_qk: bool = True, + only_cross_attention: bool = False, + eps: float = 1e-5, + rescale_output_factor: float = 1.0, + residual_connection: bool = False, + _from_deprecated_attn_block: bool = False, + processor: Optional["AttnProcessor"] = None, + out_dim: int = None, + out_context_dim: int = None, + context_pre_only=None, + pre_only=False, + elementwise_affine: bool = True, + is_causal: bool = False, + ): + super().__init__() + + # To prevent circular import. + # from .normalization import FP32LayerNorm, LpNorm, RMSNorm + + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads + self.query_dim = query_dim + self.use_bias = bias + self.is_cross_attention = cross_attention_dim is not None + self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.upcast_attention = upcast_attention + self.upcast_softmax = upcast_softmax + self.rescale_output_factor = rescale_output_factor + self.residual_connection = residual_connection + self.dropout = dropout + self.fused_projections = False + self.out_dim = out_dim if out_dim is not None else query_dim + self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim + self.context_pre_only = context_pre_only + self.pre_only = pre_only + self.is_causal = is_causal + + # we make use of this private variable to know whether this class is loaded + # with an deprecated state dict so that we can convert it on the fly + self._from_deprecated_attn_block = _from_deprecated_attn_block + + self.scale_qk = scale_qk + self.scale = dim_head**-0.5 if self.scale_qk else 1.0 + + self.heads = out_dim // dim_head if out_dim is not None else heads + # for slice_size > 0 the attention score computation + # is split across the batch axis to save memory + # You can set slice_size with `set_attention_slice` + self.sliceable_head_dim = heads + + self.added_kv_proj_dim = added_kv_proj_dim + self.only_cross_attention = only_cross_attention + + if self.added_kv_proj_dim is None and self.only_cross_attention: + raise ValueError( + "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`." + ) + + if norm_num_groups is not None: + self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) + else: + self.group_norm = None + + if spatial_norm_dim is not None: + self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim) + else: + self.spatial_norm = None + + if qk_norm is None: + self.norm_q = None + self.norm_k = None + elif qk_norm == "layer_norm": + self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + elif qk_norm == "fp32_layer_norm": + self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) + self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) + elif qk_norm == "layer_norm_across_heads": + # Lumina applies qk norm across all heads + self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps) + self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps) + elif qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim_head * heads, eps=eps) + self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps) + elif qk_norm == "l2": + self.norm_q = LpNorm(p=2, dim=-1, eps=eps) + self.norm_k = LpNorm(p=2, dim=-1, eps=eps) + else: + raise ValueError( + f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'." + ) + + if cross_attention_norm is None: + self.norm_cross = None + elif cross_attention_norm == "layer_norm": + self.norm_cross = nn.LayerNorm(self.cross_attention_dim) + elif cross_attention_norm == "group_norm": + if self.added_kv_proj_dim is not None: + # The given `encoder_hidden_states` are initially of shape + # (batch_size, seq_len, added_kv_proj_dim) before being projected + # to (batch_size, seq_len, cross_attention_dim). The norm is applied + # before the projection, so we need to use `added_kv_proj_dim` as + # the number of channels for the group norm. + norm_cross_num_channels = added_kv_proj_dim + else: + norm_cross_num_channels = self.cross_attention_dim + + self.norm_cross = nn.GroupNorm( + num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True + ) + else: + raise ValueError( + f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" + ) + + self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) + + if not self.only_cross_attention: + # only relevant for the `AddedKVProcessor` classes + self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + else: + self.to_k = None + self.to_v = None + + self.added_proj_bias = added_proj_bias + if self.added_kv_proj_dim is not None: + self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias) + self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias) + if self.context_pre_only is not None: + self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + else: + self.add_q_proj = None + self.add_k_proj = None + self.add_v_proj = None + + if not self.pre_only: + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(dropout)) + else: + self.to_out = None + + if self.context_pre_only is not None and not self.context_pre_only: + self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) + else: + self.to_add_out = None + + if qk_norm is not None and added_kv_proj_dim is not None: + if qk_norm == "layer_norm": + self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + elif qk_norm == "fp32_layer_norm": + self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) + self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) + elif qk_norm == "rms_norm": + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # Wan applies qk norm across all heads + # Wan also doesn't apply a q norm + self.norm_added_q = None + self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps) + else: + raise ValueError( + f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`" + ) + else: + self.norm_added_q = None + self.norm_added_k = None + + # set attention processor + # We use the AttnProcessor2_0 by default when torch 2.x is used which uses + # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention + # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 + if processor is None: + processor = ( + AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() + ) + self.set_processor(processor) + + def set_use_xla_flash_attention( + self, + use_xla_flash_attention: bool, + partition_spec: Optional[Tuple[Optional[str], ...]] = None, + is_flux=False, + ) -> None: + r""" + Set whether to use xla flash attention from `torch_xla` or not. + + Args: + use_xla_flash_attention (`bool`): + Whether to use pallas flash attention kernel from `torch_xla` or not. + partition_spec (`Tuple[]`, *optional*): + Specify the partition specification if using SPMD. Otherwise None. + """ + if use_xla_flash_attention: + if not is_torch_xla_available: + raise "torch_xla is not available" + elif is_torch_xla_version("<", "2.3"): + raise "flash attention pallas kernel is supported from torch_xla version 2.3" + elif is_spmd() and is_torch_xla_version("<", "2.4"): + raise "flash attention pallas kernel using SPMD is supported from torch_xla version 2.4" + else: + if is_flux: + processor = XLAFluxFlashAttnProcessor2_0(partition_spec) + else: + processor = XLAFlashAttnProcessor2_0(partition_spec) + else: + processor = ( + AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() + ) + self.set_processor(processor) + + def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: + r""" + Set whether to use npu flash attention from `torch_npu` or not. + + """ + if use_npu_flash_attention: + processor = AttnProcessorNPU() + else: + # set attention processor + # We use the AttnProcessor2_0 by default when torch 2.x is used which uses + # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention + # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 + processor = ( + AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() + ) + self.set_processor(processor) + + def set_use_memory_efficient_attention_xformers( + self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None + ) -> None: + r""" + Set whether to use memory efficient attention from `xformers` or not. + + Args: + use_memory_efficient_attention_xformers (`bool`): + Whether to use memory efficient attention from `xformers` or not. + attention_op (`Callable`, *optional*): + The attention operation to use. Defaults to `None` which uses the default attention operation from + `xformers`. + """ + is_custom_diffusion = hasattr(self, "processor") and isinstance( + self.processor, + (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0), + ) + is_added_kv_processor = hasattr(self, "processor") and isinstance( + self.processor, + ( + AttnAddedKVProcessor, + AttnAddedKVProcessor2_0, + SlicedAttnAddedKVProcessor, + XFormersAttnAddedKVProcessor, + ), + ) + is_ip_adapter = hasattr(self, "processor") and isinstance( + self.processor, + (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor), + ) + is_joint_processor = hasattr(self, "processor") and isinstance( + self.processor, + ( + JointAttnProcessor2_0, + XFormersJointAttnProcessor, + ), + ) + + if use_memory_efficient_attention_xformers: + if is_added_kv_processor and is_custom_diffusion: + raise NotImplementedError( + f"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}" + ) + if not is_xformers_available(): + raise ModuleNotFoundError( + ( + "Refer to https://github.com/facebookresearch/xformers for more information on how to install" + " xformers" + ), + name="xformers", + ) + elif not torch.cuda.is_available(): + raise ValueError( + "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is" + " only available for GPU " + ) + else: + try: + # Make sure we can run the memory efficient attention + dtype = None + if attention_op is not None: + op_fw, op_bw = attention_op + dtype, *_ = op_fw.SUPPORTED_DTYPES + q = torch.randn((1, 2, 40), device="cuda", dtype=dtype) + _ = xformers.ops.memory_efficient_attention(q, q, q) + except Exception as e: + raise e + + if is_custom_diffusion: + processor = CustomDiffusionXFormersAttnProcessor( + train_kv=self.processor.train_kv, + train_q_out=self.processor.train_q_out, + hidden_size=self.processor.hidden_size, + cross_attention_dim=self.processor.cross_attention_dim, + attention_op=attention_op, + ) + processor.load_state_dict(self.processor.state_dict()) + if hasattr(self.processor, "to_k_custom_diffusion"): + processor.to(self.processor.to_k_custom_diffusion.weight.device) + elif is_added_kv_processor: + # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP + # which uses this type of cross attention ONLY because the attention mask of format + # [0, ..., -10.000, ..., 0, ...,] is not supported + # throw warning + logger.info( + "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation." + ) + processor = XFormersAttnAddedKVProcessor(attention_op=attention_op) + elif is_ip_adapter: + processor = IPAdapterXFormersAttnProcessor( + hidden_size=self.processor.hidden_size, + cross_attention_dim=self.processor.cross_attention_dim, + num_tokens=self.processor.num_tokens, + scale=self.processor.scale, + attention_op=attention_op, + ) + processor.load_state_dict(self.processor.state_dict()) + if hasattr(self.processor, "to_k_ip"): + processor.to( + device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype + ) + elif is_joint_processor: + processor = XFormersJointAttnProcessor(attention_op=attention_op) + else: + processor = XFormersAttnProcessor(attention_op=attention_op) + else: + if is_custom_diffusion: + attn_processor_class = ( + CustomDiffusionAttnProcessor2_0 + if hasattr(F, "scaled_dot_product_attention") + else CustomDiffusionAttnProcessor + ) + processor = attn_processor_class( + train_kv=self.processor.train_kv, + train_q_out=self.processor.train_q_out, + hidden_size=self.processor.hidden_size, + cross_attention_dim=self.processor.cross_attention_dim, + ) + processor.load_state_dict(self.processor.state_dict()) + if hasattr(self.processor, "to_k_custom_diffusion"): + processor.to(self.processor.to_k_custom_diffusion.weight.device) + elif is_ip_adapter: + processor = IPAdapterAttnProcessor2_0( + hidden_size=self.processor.hidden_size, + cross_attention_dim=self.processor.cross_attention_dim, + num_tokens=self.processor.num_tokens, + scale=self.processor.scale, + ) + processor.load_state_dict(self.processor.state_dict()) + if hasattr(self.processor, "to_k_ip"): + processor.to( + device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype + ) + else: + # set attention processor + # We use the AttnProcessor2_0 by default when torch 2.x is used which uses + # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention + # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 + processor = ( + AttnProcessor2_0() + if hasattr(F, "scaled_dot_product_attention") and self.scale_qk + else AttnProcessor() + ) + + self.set_processor(processor) + + def set_attention_slice(self, slice_size: int) -> None: + r""" + Set the slice size for attention computation. + + Args: + slice_size (`int`): + The slice size for attention computation. + """ + if slice_size is not None and slice_size > self.sliceable_head_dim: + raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.") + + if slice_size is not None and self.added_kv_proj_dim is not None: + processor = SlicedAttnAddedKVProcessor(slice_size) + elif slice_size is not None: + processor = SlicedAttnProcessor(slice_size) + elif self.added_kv_proj_dim is not None: + processor = AttnAddedKVProcessor() + else: + # set attention processor + # We use the AttnProcessor2_0 by default when torch 2.x is used which uses + # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention + # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 + processor = ( + AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() + ) + + self.set_processor(processor) + + def set_processor(self, processor: "AttnProcessor") -> None: + r""" + Set the attention processor to use. + + Args: + processor (`AttnProcessor`): + The attention processor to use. + """ + # if current processor is in `self._modules` and if passed `processor` is not, we need to + # pop `processor` from `self._modules` + if ( + hasattr(self, "processor") + and isinstance(self.processor, torch.nn.Module) + and not isinstance(processor, torch.nn.Module) + ): + logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}") + self._modules.pop("processor") + + self.processor = processor + + def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor": + r""" + Get the attention processor in use. + + Args: + return_deprecated_lora (`bool`, *optional*, defaults to `False`): + Set to `True` to return the deprecated LoRA attention processor. + + Returns: + "AttentionProcessor": The attention processor in use. + """ + if not return_deprecated_lora: + return self.processor + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + r""" + The forward method of the `Attention` class. + + Args: + hidden_states (`torch.Tensor`): + The hidden states of the query. + encoder_hidden_states (`torch.Tensor`, *optional*): + The hidden states of the encoder. + attention_mask (`torch.Tensor`, *optional*): + The attention mask to use. If `None`, no mask is applied. + **cross_attention_kwargs: + Additional keyword arguments to pass along to the cross attention. + + Returns: + `torch.Tensor`: The output of the attention layer. + """ + # The `Attention` class can call different attention processors / attention functions + # here we simply pass along all tensors to the selected processor class + # For standard processors that are defined here, `**cross_attention_kwargs` is empty + + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"} + unused_kwargs = [ + k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters + ] + if len(unused_kwargs) > 0: + logger.warning( + f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + ) + cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters} + + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor: + r""" + Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads` + is the number of heads initialized while constructing the `Attention` class. + + Args: + tensor (`torch.Tensor`): The tensor to reshape. + + Returns: + `torch.Tensor`: The reshaped tensor. + """ + head_size = self.heads + batch_size, seq_len, dim = tensor.shape + tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) + tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size) + return tensor + + def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor: + r""" + Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is + the number of heads initialized while constructing the `Attention` class. + + Args: + tensor (`torch.Tensor`): The tensor to reshape. + out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is + reshaped to `[batch_size * heads, seq_len, dim // heads]`. + + Returns: + `torch.Tensor`: The reshaped tensor. + """ + head_size = self.heads + if tensor.ndim == 3: + batch_size, seq_len, dim = tensor.shape + extra_dim = 1 + else: + batch_size, extra_dim, seq_len, dim = tensor.shape + tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size) + tensor = tensor.permute(0, 2, 1, 3) + + if out_dim == 3: + tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size) + + return tensor + + def get_attention_scores( + self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + r""" + Compute the attention scores. + + Args: + query (`torch.Tensor`): The query tensor. + key (`torch.Tensor`): The key tensor. + attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied. + + Returns: + `torch.Tensor`: The attention probabilities/scores. + """ + dtype = query.dtype + if self.upcast_attention: + query = query.float() + key = key.float() + + if attention_mask is None: + baddbmm_input = torch.empty( + query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device + ) + beta = 0 + else: + baddbmm_input = attention_mask + beta = 1 + + attention_scores = torch.baddbmm( + baddbmm_input, + query, + key.transpose(-1, -2), + beta=beta, + alpha=self.scale, + ) + del baddbmm_input + + if self.upcast_softmax: + attention_scores = attention_scores.float() + + attention_probs = attention_scores.softmax(dim=-1) + del attention_scores + + attention_probs = attention_probs.to(dtype) + + return attention_probs + + def prepare_attention_mask( + self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 + ) -> torch.Tensor: + r""" + Prepare the attention mask for the attention computation. + + Args: + attention_mask (`torch.Tensor`): + The attention mask to prepare. + target_length (`int`): + The target length of the attention mask. This is the length of the attention mask after padding. + batch_size (`int`): + The batch size, which is used to repeat the attention mask. + out_dim (`int`, *optional*, defaults to `3`): + The output dimension of the attention mask. Can be either `3` or `4`. + + Returns: + `torch.Tensor`: The prepared attention mask. + """ + head_size = self.heads + if attention_mask is None: + return attention_mask + + current_length: int = attention_mask.shape[-1] + if current_length != target_length: + if attention_mask.device.type == "mps": + # HACK: MPS: Does not support padding by greater than dimension of input tensor. + # Instead, we can manually construct the padding tensor. + padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length) + padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) + attention_mask = torch.cat([attention_mask, padding], dim=2) + else: + # TODO: for pipelines such as stable-diffusion, padding cross-attn mask: + # we want to instead pad by (0, remaining_length), where remaining_length is: + # remaining_length: int = target_length - current_length + # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding + attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) + + if out_dim == 3: + if attention_mask.shape[0] < batch_size * head_size: + attention_mask = attention_mask.repeat_interleave( + head_size, dim=0, output_size=attention_mask.shape[0] * head_size + ) + elif out_dim == 4: + attention_mask = attention_mask.unsqueeze(1) + attention_mask = attention_mask.repeat_interleave( + head_size, dim=1, output_size=attention_mask.shape[1] * head_size + ) + + return attention_mask + + def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the + `Attention` class. + + Args: + encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder. + + Returns: + `torch.Tensor`: The normalized encoder hidden states. + """ + assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states" + + if isinstance(self.norm_cross, nn.LayerNorm): + encoder_hidden_states = self.norm_cross(encoder_hidden_states) + elif isinstance(self.norm_cross, nn.GroupNorm): + # Group norm norms along the channels dimension and expects + # input to be in the shape of (N, C, *). In this case, we want + # to norm along the hidden dimension, so we need to move + # (batch_size, sequence_length, hidden_size) -> + # (batch_size, hidden_size, sequence_length) + encoder_hidden_states = encoder_hidden_states.transpose(1, 2) + encoder_hidden_states = self.norm_cross(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.transpose(1, 2) + else: + assert False + + return encoder_hidden_states + + @torch.no_grad() + def fuse_projections(self, fuse=True): + device = self.to_q.weight.data.device + dtype = self.to_q.weight.data.dtype + + if not self.is_cross_attention: + # fetch weight matrices. + concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) + in_features = concatenated_weights.shape[1] + out_features = concatenated_weights.shape[0] + + # create a new single projection layer and copy over the weights. + self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype) + self.to_qkv.weight.copy_(concatenated_weights) + if self.use_bias: + concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) + self.to_qkv.bias.copy_(concatenated_bias) + + else: + concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) + in_features = concatenated_weights.shape[1] + out_features = concatenated_weights.shape[0] + + self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype) + self.to_kv.weight.copy_(concatenated_weights) + if self.use_bias: + concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) + self.to_kv.bias.copy_(concatenated_bias) + + # handle added projections for SD3 and others. + if ( + getattr(self, "add_q_proj", None) is not None + and getattr(self, "add_k_proj", None) is not None + and getattr(self, "add_v_proj", None) is not None + ): + concatenated_weights = torch.cat( + [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data] + ) + in_features = concatenated_weights.shape[1] + out_features = concatenated_weights.shape[0] + + self.to_added_qkv = nn.Linear( + in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype + ) + self.to_added_qkv.weight.copy_(concatenated_weights) + if self.added_proj_bias: + concatenated_bias = torch.cat( + [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data] + ) + self.to_added_qkv.bias.copy_(concatenated_bias) + + self.fused_projections = fuse + +class AttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + temb: Optional[torch.Tensor] = None, + *args, + **kwargs, + ) -> torch.Tensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." + deprecate("scale", "1.0.0", deprecation_message) + + residual = hidden_states + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + +class UNetMidBlock2D(nn.Module): + """ + A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks. + + Args: + in_channels (`int`): The number of input channels. + temb_channels (`int`): The number of temporal embedding channels. + dropout (`float`, *optional*, defaults to 0.0): The dropout rate. + num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. + resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. + resnet_time_scale_shift (`str`, *optional*, defaults to `default`): + The type of normalization to apply to the time embeddings. This can help to improve the performance of the + model on tasks with long-range temporal dependencies. + resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. + resnet_groups (`int`, *optional*, defaults to 32): + The number of groups to use in the group normalization layers of the resnet blocks. + attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks. + resnet_pre_norm (`bool`, *optional*, defaults to `True`): + Whether to use pre-normalization for the resnet blocks. + add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. + attention_head_dim (`int`, *optional*, defaults to 1): + Dimension of a single attention head. The number of attention heads is determined based on this value and + the number of input channels. + output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. + + Returns: + `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, in_channels, + height, width)`. + + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + if attn_groups is None: + attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None + + # there is always at least one resnet + if resnet_time_scale_shift == "spatial": + resnets = [ + ResnetBlockCondNorm2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm="spatial", + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + ) + ] + else: + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warning( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + if resnet_time_scale_shift == "spatial": + resnets.append( + ResnetBlockCondNorm2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm="spatial", + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + ) + ) + else: + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if torch.is_grad_enabled() and self.gradient_checkpointing: + if attn is not None: + hidden_states = attn(hidden_states, temb=temb) + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) + else: + if attn is not None: + hidden_states = attn(hidden_states, temb=temb) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + +class DownEncoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + if resnet_time_scale_shift == "spatial": + resnets.append( + ResnetBlockCondNorm2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm="spatial", + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + ) + ) + else: + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." + deprecate("scale", "1.0.0", deprecation_message) + + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class UpDecoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + temb_channels: Optional[int] = None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + if resnet_time_scale_shift == "spatial": + resnets.append( + ResnetBlockCondNorm2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm="spatial", + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + ) + ) + else: + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.resolution_idx = resolution_idx + + def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + +class Encoder(nn.Module): + r""" + The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available + options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + double_z (`bool`, *optional*, defaults to `True`): + Whether to double the number of output channels for the last block. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=3, + stride=1, + padding=1, + ) + + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = DownEncoderBlock2D( + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=not is_final_block, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + # attention_head_dim=output_channel, + # temb_channels=None, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlock2D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + ) + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + r"""The forward method of the `Encoder` class.""" + + sample = self.conv_in(sample) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + # down + for down_block in self.down_blocks: + sample = self._gradient_checkpointing_func(down_block, sample) + # middle + sample = self._gradient_checkpointing_func(self.mid_block, sample) + + else: + # down + for down_block in self.down_blocks: + sample = down_block(sample) + + # middle + sample = self.mid_block(sample) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + +class Decoder(nn.Module): + r""" + The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + norm_type (`str`, *optional*, defaults to `"group"`): + The normalization type to use. Can be either `"group"` or `"spatial"`. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1, + padding=1, + ) + + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlock2D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + + is_final_block = i == len(block_out_channels) - 1 + + up_block = UpDecoderBlock2D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + # prev_output_channel=prev_output_channel, + add_upsample=not is_final_block, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + # attention_head_dim=output_channel, + temb_channels=temb_channels, + resnet_time_scale_shift=norm_type, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) + else: + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward( + self, + sample: torch.Tensor, + latent_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + r"""The forward method of the `Decoder` class.""" + + sample = self.conv_in(sample) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + # middle + sample = self._gradient_checkpointing_func(self.mid_block, sample, latent_embeds) + + # up + for up_block in self.up_blocks: + sample = self._gradient_checkpointing_func(up_block, sample, latent_embeds) + else: + # middle + sample = self.mid_block(sample, latent_embeds) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds) + + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class Flux2VAE(torch.nn.Module): + r""" + A VAE model with KL loss for encoding images into latents and decoding latent representations into images. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + in_channels (int, *optional*, defaults to 3): Number of channels in the input image. + out_channels (int, *optional*, defaults to 3): Number of channels in the output. + down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + Tuple of downsample block types. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + Tuple of upsample block types. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): + Tuple of block output channels. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. + sample_size (`int`, *optional*, defaults to `32`): Sample input size. + force_upcast (`bool`, *optional*, default to `True`): + If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE + can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` + can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix + mid_block_add_attention (`bool`, *optional*, default to `True`): + If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the + mid_block will only have resnet blocks + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"] + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ( + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ), + up_block_types: Tuple[str, ...] = ( + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ), + block_out_channels: Tuple[int, ...] = ( + 128, + 256, + 512, + 512, + ), + layers_per_block: int = 2, + act_fn: str = "silu", + latent_channels: int = 32, + norm_num_groups: int = 32, + sample_size: int = 1024, # YiYi notes: not sure + force_upcast: bool = True, + use_quant_conv: bool = True, + use_post_quant_conv: bool = True, + mid_block_add_attention: bool = True, + batch_norm_eps: float = 1e-4, + batch_norm_momentum: float = 0.1, + patch_size: Tuple[int, int] = (2, 2), + ): + super().__init__() + + # pass init params to Encoder + self.encoder = Encoder( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + mid_block_add_attention=mid_block_add_attention, + ) + + # pass init params to Decoder + self.decoder = Decoder( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + mid_block_add_attention=mid_block_add_attention, + ) + + self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None + self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None + + self.bn = nn.BatchNorm2d( + math.prod(patch_size) * latent_channels, + eps=batch_norm_eps, + momentum=batch_norm_momentum, + affine=False, + track_running_stats=True, + ) + + self.use_slicing = False + self.use_tiling = False + + @property + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self): + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor() + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor(self, processor): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, height, width = x.shape + + if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size): + return self._tiled_encode(x) + + enc = self.encoder(x) + if self.quant_conv is not None: + enc = self.quant_conv(enc) + + return enc + + def encode( + self, x: torch.Tensor, return_dict: bool = True + ): + """ + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + + + h = rearrange(h, "B C (H P) (W Q) -> B (C P Q) H W", P=2, Q=2) + h = h[:, :128] + latents_bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(h.device, h.dtype) + latents_bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + 0.0001).to( + h.device, h.dtype + ) + h = (h - latents_bn_mean) / latents_bn_std + return h + + def _decode(self, z: torch.Tensor, return_dict: bool = True): + if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): + return self.tiled_decode(z, return_dict=return_dict) + + if self.post_quant_conv is not None: + z = self.post_quant_conv(z) + + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return dec + + def decode( + self, z: torch.FloatTensor, return_dict: bool = True, generator=None + ): + latents_bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(z.device, z.dtype) + latents_bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + 0.0001).to( + z.device, z.dtype + ) + z = z * latents_bn_std + latents_bn_mean + z = rearrange(z, "B (C P Q) H W -> B C (H P) (W Q)", P=2, Q=2) + """ + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + + return decoded + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[2], b.shape[2], blend_extent) + for y in range(blend_extent): + b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for x in range(blend_extent): + b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent) + return b + + def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.Tensor`): Input batch of images. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split the image into 512x512 tiles and encode them separately. + rows = [] + for i in range(0, x.shape[2], overlap_size): + row = [] + for j in range(0, x.shape[3], overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + if self.config.use_quant_conv: + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=3)) + + enc = torch.cat(result_rows, dim=2) + return enc + + def tiled_encode(self, x: torch.Tensor, return_dict: bool = True): + r"""Encode a batch of images using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split the image into 512x512 tiles and encode them separately. + rows = [] + for i in range(0, x.shape[2], overlap_size): + row = [] + for j in range(0, x.shape[3], overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + if self.config.use_quant_conv: + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=3)) + + moments = torch.cat(result_rows, dim=2) + return moments + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping 64x64 tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[2], overlap_size): + row = [] + for j in range(0, z.shape[3], overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + if self.config.use_post_quant_conv: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=3)) + + dec = torch.cat(result_rows, dim=2) + if not return_dict: + return (dec,) + + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ): + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z).sample + + if not return_dict: + return (dec,) + + return dec diff --git a/diffsynth/models/flux_controlnet.py b/diffsynth/models/flux_controlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb1138bb74f7b55c6e92ac312098e4168829f0a --- /dev/null +++ b/diffsynth/models/flux_controlnet.py @@ -0,0 +1,384 @@ +import torch +from einops import rearrange, repeat +from .flux_dit import RoPEEmbedding, TimestepEmbeddings, FluxJointTransformerBlock, FluxSingleTransformerBlock, RMSNorm +# from .utils import hash_state_dict_keys, init_weights_on_device +from contextlib import contextmanager + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() + +@contextmanager +def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + +class FluxControlNet(torch.nn.Module): + def __init__(self, disable_guidance_embedder=False, num_joint_blocks=5, num_single_blocks=10, num_mode=0, mode_dict={}, additional_input_dim=0): + super().__init__() + self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56]) + self.time_embedder = TimestepEmbeddings(256, 3072) + self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072) + self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072)) + self.context_embedder = torch.nn.Linear(4096, 3072) + self.x_embedder = torch.nn.Linear(64, 3072) + + self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_joint_blocks)]) + self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(num_single_blocks)]) + + self.controlnet_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_joint_blocks)]) + self.controlnet_single_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_single_blocks)]) + + self.mode_dict = mode_dict + self.controlnet_mode_embedder = torch.nn.Embedding(num_mode, 3072) if len(mode_dict) > 0 else None + self.controlnet_x_embedder = torch.nn.Linear(64 + additional_input_dim, 3072) + + + def prepare_image_ids(self, latents): + batch_size, _, height, width = latents.shape + latent_image_ids = torch.zeros(height // 2, width // 2, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1) + latent_image_ids = latent_image_ids.reshape( + batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + latent_image_ids = latent_image_ids.to(device=latents.device, dtype=latents.dtype) + + return latent_image_ids + + + def patchify(self, hidden_states): + hidden_states = rearrange(hidden_states, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) + return hidden_states + + + def align_res_stack_to_original_blocks(self, res_stack, num_blocks, hidden_states): + if len(res_stack) == 0: + return [torch.zeros_like(hidden_states)] * num_blocks + interval = (num_blocks + len(res_stack) - 1) // len(res_stack) + aligned_res_stack = [res_stack[block_id // interval] for block_id in range(num_blocks)] + return aligned_res_stack + + + def forward( + self, + hidden_states, + controlnet_conditioning, + timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None, + processor_id=None, + tiled=False, tile_size=128, tile_stride=64, + **kwargs + ): + if image_ids is None: + image_ids = self.prepare_image_ids(hidden_states) + + conditioning = self.time_embedder(timestep, hidden_states.dtype) + self.pooled_text_embedder(pooled_prompt_emb) + if self.guidance_embedder is not None: + guidance = guidance * 1000 + conditioning = conditioning + self.guidance_embedder(guidance, hidden_states.dtype) + prompt_emb = self.context_embedder(prompt_emb) + if self.controlnet_mode_embedder is not None: # Different from FluxDiT + processor_id = torch.tensor([self.mode_dict[processor_id]], dtype=torch.int) + processor_id = repeat(processor_id, "D -> B D", B=1).to(text_ids.device) + prompt_emb = torch.concat([self.controlnet_mode_embedder(processor_id), prompt_emb], dim=1) + text_ids = torch.cat([text_ids[:, :1], text_ids], dim=1) + image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1)) + + hidden_states = self.patchify(hidden_states) + hidden_states = self.x_embedder(hidden_states) + controlnet_conditioning = self.patchify(controlnet_conditioning) # Different from FluxDiT + hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_conditioning) # Different from FluxDiT + + controlnet_res_stack = [] + for block, controlnet_block in zip(self.blocks, self.controlnet_blocks): + hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb) + controlnet_res_stack.append(controlnet_block(hidden_states)) + + controlnet_single_res_stack = [] + hidden_states = torch.cat([prompt_emb, hidden_states], dim=1) + for block, controlnet_block in zip(self.single_blocks, self.controlnet_single_blocks): + hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb) + controlnet_single_res_stack.append(controlnet_block(hidden_states[:, prompt_emb.shape[1]:])) + + controlnet_res_stack = self.align_res_stack_to_original_blocks(controlnet_res_stack, 19, hidden_states[:, prompt_emb.shape[1]:]) + controlnet_single_res_stack = self.align_res_stack_to_original_blocks(controlnet_single_res_stack, 38, hidden_states[:, prompt_emb.shape[1]:]) + + return controlnet_res_stack, controlnet_single_res_stack + + + # @staticmethod + # def state_dict_converter(): + # return FluxControlNetStateDictConverter() + + def quantize(self): + def cast_to(weight, dtype=None, device=None, copy=False): + if device is None or weight.device == device: + if not copy: + if dtype is None or weight.dtype == dtype: + return weight + return weight.to(dtype=dtype, copy=copy) + + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight) + return r + + def cast_weight(s, input=None, dtype=None, device=None): + if input is not None: + if dtype is None: + dtype = input.dtype + if device is None: + device = input.device + weight = cast_to(s.weight, dtype, device) + return weight + + def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): + if input is not None: + if dtype is None: + dtype = input.dtype + if bias_dtype is None: + bias_dtype = dtype + if device is None: + device = input.device + bias = None + weight = cast_to(s.weight, dtype, device) + bias = cast_to(s.bias, bias_dtype, device) + return weight, bias + + class quantized_layer: + class QLinear(torch.nn.Linear): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self,input,**kwargs): + weight,bias= cast_bias_weight(self,input) + return torch.nn.functional.linear(input,weight,bias) + + class QRMSNorm(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self,hidden_states,**kwargs): + weight= cast_weight(self.module,hidden_states) + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.module.eps) + hidden_states = hidden_states.to(input_dtype) * weight + return hidden_states + + class QEmbedding(torch.nn.Embedding): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self,input,**kwargs): + weight= cast_weight(self,input) + return torch.nn.functional.embedding( + input, weight, self.padding_idx, self.max_norm, + self.norm_type, self.scale_grad_by_freq, self.sparse) + + def replace_layer(model): + for name, module in model.named_children(): + if isinstance(module,quantized_layer.QRMSNorm): + continue + if isinstance(module, torch.nn.Linear): + with init_weights_on_device(): + new_layer = quantized_layer.QLinear(module.in_features,module.out_features) + new_layer.weight = module.weight + if module.bias is not None: + new_layer.bias = module.bias + setattr(model, name, new_layer) + elif isinstance(module, RMSNorm): + if hasattr(module,"quantized"): + continue + module.quantized= True + new_layer = quantized_layer.QRMSNorm(module) + setattr(model, name, new_layer) + elif isinstance(module,torch.nn.Embedding): + rows, cols = module.weight.shape + new_layer = quantized_layer.QEmbedding( + num_embeddings=rows, + embedding_dim=cols, + _weight=module.weight, + # _freeze=module.freeze, + padding_idx=module.padding_idx, + max_norm=module.max_norm, + norm_type=module.norm_type, + scale_grad_by_freq=module.scale_grad_by_freq, + sparse=module.sparse) + setattr(model, name, new_layer) + else: + replace_layer(module) + + replace_layer(self) + + + +class FluxControlNetStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + hash_value = hash_state_dict_keys(state_dict) + global_rename_dict = { + "context_embedder": "context_embedder", + "x_embedder": "x_embedder", + "time_text_embed.timestep_embedder.linear_1": "time_embedder.timestep_embedder.0", + "time_text_embed.timestep_embedder.linear_2": "time_embedder.timestep_embedder.2", + "time_text_embed.guidance_embedder.linear_1": "guidance_embedder.timestep_embedder.0", + "time_text_embed.guidance_embedder.linear_2": "guidance_embedder.timestep_embedder.2", + "time_text_embed.text_embedder.linear_1": "pooled_text_embedder.0", + "time_text_embed.text_embedder.linear_2": "pooled_text_embedder.2", + "norm_out.linear": "final_norm_out.linear", + "proj_out": "final_proj_out", + } + rename_dict = { + "proj_out": "proj_out", + "norm1.linear": "norm1_a.linear", + "norm1_context.linear": "norm1_b.linear", + "attn.to_q": "attn.a_to_q", + "attn.to_k": "attn.a_to_k", + "attn.to_v": "attn.a_to_v", + "attn.to_out.0": "attn.a_to_out", + "attn.add_q_proj": "attn.b_to_q", + "attn.add_k_proj": "attn.b_to_k", + "attn.add_v_proj": "attn.b_to_v", + "attn.to_add_out": "attn.b_to_out", + "ff.net.0.proj": "ff_a.0", + "ff.net.2": "ff_a.2", + "ff_context.net.0.proj": "ff_b.0", + "ff_context.net.2": "ff_b.2", + "attn.norm_q": "attn.norm_q_a", + "attn.norm_k": "attn.norm_k_a", + "attn.norm_added_q": "attn.norm_q_b", + "attn.norm_added_k": "attn.norm_k_b", + } + rename_dict_single = { + "attn.to_q": "a_to_q", + "attn.to_k": "a_to_k", + "attn.to_v": "a_to_v", + "attn.norm_q": "norm_q_a", + "attn.norm_k": "norm_k_a", + "norm.linear": "norm.linear", + "proj_mlp": "proj_in_besides_attn", + "proj_out": "proj_out", + } + state_dict_ = {} + for name, param in state_dict.items(): + if name.endswith(".weight") or name.endswith(".bias"): + suffix = ".weight" if name.endswith(".weight") else ".bias" + prefix = name[:-len(suffix)] + if prefix in global_rename_dict: + state_dict_[global_rename_dict[prefix] + suffix] = param + elif prefix.startswith("transformer_blocks."): + names = prefix.split(".") + names[0] = "blocks" + middle = ".".join(names[2:]) + if middle in rename_dict: + name_ = ".".join(names[:2] + [rename_dict[middle]] + [suffix[1:]]) + state_dict_[name_] = param + elif prefix.startswith("single_transformer_blocks."): + names = prefix.split(".") + names[0] = "single_blocks" + middle = ".".join(names[2:]) + if middle in rename_dict_single: + name_ = ".".join(names[:2] + [rename_dict_single[middle]] + [suffix[1:]]) + state_dict_[name_] = param + else: + state_dict_[name] = param + else: + state_dict_[name] = param + for name in list(state_dict_.keys()): + if ".proj_in_besides_attn." in name: + name_ = name.replace(".proj_in_besides_attn.", ".to_qkv_mlp.") + param = torch.concat([ + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_q.")], + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_k.")], + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_v.")], + state_dict_[name], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_q.")) + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_k.")) + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_v.")) + state_dict_.pop(name) + for name in list(state_dict_.keys()): + for component in ["a", "b"]: + if f".{component}_to_q." in name: + name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.") + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v.")) + if hash_value == "78d18b9101345ff695f312e7e62538c0": + extra_kwargs = {"num_mode": 10, "mode_dict": {"canny": 0, "tile": 1, "depth": 2, "blur": 3, "pose": 4, "gray": 5, "lq": 6}} + elif hash_value == "b001c89139b5f053c715fe772362dd2a": + extra_kwargs = {"num_single_blocks": 0} + elif hash_value == "52357cb26250681367488a8954c271e8": + extra_kwargs = {"num_joint_blocks": 6, "num_single_blocks": 0, "additional_input_dim": 4} + elif hash_value == "0cfd1740758423a2a854d67c136d1e8c": + extra_kwargs = {"num_joint_blocks": 4, "num_single_blocks": 1} + elif hash_value == "7f9583eb8ba86642abb9a21a4b2c9e16": + extra_kwargs = {"num_joint_blocks": 4, "num_single_blocks": 10} + elif hash_value == "43ad5aaa27dd4ee01b832ed16773fa52": + extra_kwargs = {"num_joint_blocks": 6, "num_single_blocks": 0} + else: + extra_kwargs = {} + return state_dict_, extra_kwargs + + + def from_civitai(self, state_dict): + return self.from_diffusers(state_dict) diff --git a/diffsynth/models/flux_dit.py b/diffsynth/models/flux_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..51a6e7f049d15e764383487c0edbf7da839b5918 --- /dev/null +++ b/diffsynth/models/flux_dit.py @@ -0,0 +1,395 @@ +import torch +from .general_modules import TimestepEmbeddings, AdaLayerNorm, RMSNorm +from einops import rearrange + + +def interact_with_ipadapter(hidden_states, q, ip_k, ip_v, scale=1.0): + batch_size, num_tokens = hidden_states.shape[0:2] + ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v) + ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, num_tokens, -1) + hidden_states = hidden_states + scale * ip_hidden_states + return hidden_states + + +class RoPEEmbedding(torch.nn.Module): + def __init__(self, dim, theta, axes_dim): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + + def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: + assert dim % 2 == 0, "The dimension must be even." + + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + omega = 1.0 / (theta**scale) + + batch_size, seq_length = pos.shape + out = torch.einsum("...n,d->...nd", pos, omega) + cos_out = torch.cos(out) + sin_out = torch.sin(out) + + stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1) + out = stacked_out.view(batch_size, -1, dim // 2, 2, 2) + return out.float() + + + def forward(self, ids): + n_axes = ids.shape[-1] + emb = torch.cat([self.rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3) + return emb.unsqueeze(1) + + + +class FluxJointAttention(torch.nn.Module): + def __init__(self, dim_a, dim_b, num_heads, head_dim, only_out_a=False): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.only_out_a = only_out_a + + self.a_to_qkv = torch.nn.Linear(dim_a, dim_a * 3) + self.b_to_qkv = torch.nn.Linear(dim_b, dim_b * 3) + + self.norm_q_a = RMSNorm(head_dim, eps=1e-6) + self.norm_k_a = RMSNorm(head_dim, eps=1e-6) + self.norm_q_b = RMSNorm(head_dim, eps=1e-6) + self.norm_k_b = RMSNorm(head_dim, eps=1e-6) + + self.a_to_out = torch.nn.Linear(dim_a, dim_a) + if not only_out_a: + self.b_to_out = torch.nn.Linear(dim_b, dim_b) + + + def apply_rope(self, xq, xk, freqs_cis): + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + def forward(self, hidden_states_a, hidden_states_b, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): + batch_size = hidden_states_a.shape[0] + + # Part A + qkv_a = self.a_to_qkv(hidden_states_a) + qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2) + q_a, k_a, v_a = qkv_a.chunk(3, dim=1) + q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a) + + # Part B + qkv_b = self.b_to_qkv(hidden_states_b) + qkv_b = qkv_b.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2) + q_b, k_b, v_b = qkv_b.chunk(3, dim=1) + q_b, k_b = self.norm_q_b(q_b), self.norm_k_b(k_b) + + q = torch.concat([q_b, q_a], dim=2) + k = torch.concat([k_b, k_a], dim=2) + v = torch.concat([v_b, v_a], dim=2) + + q, k = self.apply_rope(q, k, image_rotary_emb) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + hidden_states_b, hidden_states_a = hidden_states[:, :hidden_states_b.shape[1]], hidden_states[:, hidden_states_b.shape[1]:] + if ipadapter_kwargs_list is not None: + hidden_states_a = interact_with_ipadapter(hidden_states_a, q_a, **ipadapter_kwargs_list) + hidden_states_a = self.a_to_out(hidden_states_a) + if self.only_out_a: + return hidden_states_a + else: + hidden_states_b = self.b_to_out(hidden_states_b) + return hidden_states_a, hidden_states_b + + + +class FluxJointTransformerBlock(torch.nn.Module): + def __init__(self, dim, num_attention_heads): + super().__init__() + self.norm1_a = AdaLayerNorm(dim) + self.norm1_b = AdaLayerNorm(dim) + + self.attn = FluxJointAttention(dim, dim, num_attention_heads, dim // num_attention_heads) + + self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_a = torch.nn.Sequential( + torch.nn.Linear(dim, dim*4), + torch.nn.GELU(approximate="tanh"), + torch.nn.Linear(dim*4, dim) + ) + + self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_b = torch.nn.Sequential( + torch.nn.Linear(dim, dim*4), + torch.nn.GELU(approximate="tanh"), + torch.nn.Linear(dim*4, dim) + ) + + + def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): + norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb) + norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb) + + # Attention + attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b, image_rotary_emb, attn_mask, ipadapter_kwargs_list) + + # Part A + hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a + norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a + hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a) + + # Part B + hidden_states_b = hidden_states_b + gate_msa_b * attn_output_b + norm_hidden_states_b = self.norm2_b(hidden_states_b) * (1 + scale_mlp_b) + shift_mlp_b + hidden_states_b = hidden_states_b + gate_mlp_b * self.ff_b(norm_hidden_states_b) + + return hidden_states_a, hidden_states_b + + + +class FluxSingleAttention(torch.nn.Module): + def __init__(self, dim_a, dim_b, num_heads, head_dim): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.a_to_qkv = torch.nn.Linear(dim_a, dim_a * 3) + + self.norm_q_a = RMSNorm(head_dim, eps=1e-6) + self.norm_k_a = RMSNorm(head_dim, eps=1e-6) + + + def apply_rope(self, xq, xk, freqs_cis): + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + + def forward(self, hidden_states, image_rotary_emb): + batch_size = hidden_states.shape[0] + + qkv_a = self.a_to_qkv(hidden_states) + qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2) + q_a, k_a, v = qkv_a.chunk(3, dim=1) + q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a) + + q, k = self.apply_rope(q_a, k_a, image_rotary_emb) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + return hidden_states + + + +class AdaLayerNormSingle(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.silu = torch.nn.SiLU() + self.linear = torch.nn.Linear(dim, 3 * dim, bias=True) + self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa + + + +class FluxSingleTransformerBlock(torch.nn.Module): + def __init__(self, dim, num_attention_heads): + super().__init__() + self.num_heads = num_attention_heads + self.head_dim = dim // num_attention_heads + self.dim = dim + + self.norm = AdaLayerNormSingle(dim) + self.to_qkv_mlp = torch.nn.Linear(dim, dim * (3 + 4)) + self.norm_q_a = RMSNorm(self.head_dim, eps=1e-6) + self.norm_k_a = RMSNorm(self.head_dim, eps=1e-6) + + self.proj_out = torch.nn.Linear(dim * 5, dim) + + + def apply_rope(self, xq, xk, freqs_cis): + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + + def process_attention(self, hidden_states, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): + batch_size = hidden_states.shape[0] + + qkv = hidden_states.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2) + q, k, v = qkv.chunk(3, dim=1) + q, k = self.norm_q_a(q), self.norm_k_a(k) + + q, k = self.apply_rope(q, k, image_rotary_emb) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + if ipadapter_kwargs_list is not None: + hidden_states = interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs_list) + return hidden_states + + + def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): + residual = hidden_states_a + norm_hidden_states, gate = self.norm(hidden_states_a, emb=temb) + hidden_states_a = self.to_qkv_mlp(norm_hidden_states) + attn_output, mlp_hidden_states = hidden_states_a[:, :, :self.dim * 3], hidden_states_a[:, :, self.dim * 3:] + + attn_output = self.process_attention(attn_output, image_rotary_emb, attn_mask, ipadapter_kwargs_list) + mlp_hidden_states = torch.nn.functional.gelu(mlp_hidden_states, approximate="tanh") + + hidden_states_a = torch.cat([attn_output, mlp_hidden_states], dim=2) + hidden_states_a = gate.unsqueeze(1) * self.proj_out(hidden_states_a) + hidden_states_a = residual + hidden_states_a + + return hidden_states_a, hidden_states_b + + + +class AdaLayerNormContinuous(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.silu = torch.nn.SiLU() + self.linear = torch.nn.Linear(dim, dim * 2, bias=True) + self.norm = torch.nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + + def forward(self, x, conditioning): + emb = self.linear(self.silu(conditioning)) + shift, scale = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None] + shift[:, None] + return x + + + +class FluxDiT(torch.nn.Module): + def __init__(self, disable_guidance_embedder=False, input_dim=64, num_blocks=19): + super().__init__() + self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56]) + self.time_embedder = TimestepEmbeddings(256, 3072) + self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072) + self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072)) + self.context_embedder = torch.nn.Linear(4096, 3072) + self.x_embedder = torch.nn.Linear(input_dim, 3072) + + self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_blocks)]) + self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(38)]) + + self.final_norm_out = AdaLayerNormContinuous(3072) + self.final_proj_out = torch.nn.Linear(3072, 64) + + self.input_dim = input_dim + + + def patchify(self, hidden_states): + hidden_states = rearrange(hidden_states, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) + return hidden_states + + + def unpatchify(self, hidden_states, height, width): + hidden_states = rearrange(hidden_states, "B (H W) (C P Q) -> B C (H P) (W Q)", P=2, Q=2, H=height//2, W=width//2) + return hidden_states + + + def prepare_image_ids(self, latents): + batch_size, _, height, width = latents.shape + latent_image_ids = torch.zeros(height // 2, width // 2, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1) + latent_image_ids = latent_image_ids.reshape( + batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + latent_image_ids = latent_image_ids.to(device=latents.device, dtype=latents.dtype) + + return latent_image_ids + + + def construct_mask(self, entity_masks, prompt_seq_len, image_seq_len): + N = len(entity_masks) + batch_size = entity_masks[0].shape[0] + total_seq_len = N * prompt_seq_len + image_seq_len + patched_masks = [self.patchify(entity_masks[i]) for i in range(N)] + attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), dtype=torch.bool).to(device=entity_masks[0].device) + + image_start = N * prompt_seq_len + image_end = N * prompt_seq_len + image_seq_len + # prompt-image mask + for i in range(N): + prompt_start = i * prompt_seq_len + prompt_end = (i + 1) * prompt_seq_len + image_mask = torch.sum(patched_masks[i], dim=-1) > 0 + image_mask = image_mask.unsqueeze(1).repeat(1, prompt_seq_len, 1) + # prompt update with image + attention_mask[:, prompt_start:prompt_end, image_start:image_end] = image_mask + # image update with prompt + attention_mask[:, image_start:image_end, prompt_start:prompt_end] = image_mask.transpose(1, 2) + # prompt-prompt mask + for i in range(N): + for j in range(N): + if i != j: + prompt_start_i = i * prompt_seq_len + prompt_end_i = (i + 1) * prompt_seq_len + prompt_start_j = j * prompt_seq_len + prompt_end_j = (j + 1) * prompt_seq_len + attention_mask[:, prompt_start_i:prompt_end_i, prompt_start_j:prompt_end_j] = False + + attention_mask = attention_mask.float() + attention_mask[attention_mask == 0] = float('-inf') + attention_mask[attention_mask == 1] = 0 + return attention_mask + + + def process_entity_masks(self, hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids, repeat_dim): + max_masks = 0 + attention_mask = None + prompt_embs = [prompt_emb] + if entity_masks is not None: + # entity_masks + batch_size, max_masks = entity_masks.shape[0], entity_masks.shape[1] + entity_masks = entity_masks.repeat(1, 1, repeat_dim, 1, 1) + entity_masks = [entity_masks[:, i, None].squeeze(1) for i in range(max_masks)] + # global mask + global_mask = torch.ones_like(entity_masks[0]).to(device=hidden_states.device, dtype=hidden_states.dtype) + entity_masks = entity_masks + [global_mask] # append global to last + # attention mask + attention_mask = self.construct_mask(entity_masks, prompt_emb.shape[1], hidden_states.shape[1]) + attention_mask = attention_mask.to(device=hidden_states.device, dtype=hidden_states.dtype) + attention_mask = attention_mask.unsqueeze(1) + # embds: n_masks * b * seq * d + local_embs = [entity_prompt_emb[:, i, None].squeeze(1) for i in range(max_masks)] + prompt_embs = local_embs + prompt_embs # append global to last + prompt_embs = [self.context_embedder(prompt_emb) for prompt_emb in prompt_embs] + prompt_emb = torch.cat(prompt_embs, dim=1) + + # positional embedding + text_ids = torch.cat([text_ids] * (max_masks + 1), dim=1) + image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1)) + return prompt_emb, image_rotary_emb, attention_mask + + + def forward( + self, + hidden_states, + timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None, + tiled=False, tile_size=128, tile_stride=64, entity_prompt_emb=None, entity_masks=None, + use_gradient_checkpointing=False, + **kwargs + ): + # (Deprecated) The real forward is in `pipelines.flux_image`. + return None diff --git a/diffsynth/models/flux_infiniteyou.py b/diffsynth/models/flux_infiniteyou.py new file mode 100644 index 0000000000000000000000000000000000000000..861538a4b02fb6a52edee662b6efcd60f78f6916 --- /dev/null +++ b/diffsynth/models/flux_infiniteyou.py @@ -0,0 +1,129 @@ +import math +import torch +import torch.nn as nn + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class InfiniteYouImageProjector(nn.Module): + + def __init__( + self, + dim=1280, + depth=4, + dim_head=64, + heads=20, + num_queries=8, + embedding_dim=512, + output_dim=4096, + ff_mult=4, + ): + super().__init__() + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList([ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ])) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + latents = latents.to(dtype=x.dtype, device=x.device) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) + + @staticmethod + def state_dict_converter(): + return FluxInfiniteYouImageProjectorStateDictConverter() + + +class FluxInfiniteYouImageProjectorStateDictConverter: + + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict['image_proj'] diff --git a/diffsynth/models/flux_ipadapter.py b/diffsynth/models/flux_ipadapter.py new file mode 100644 index 0000000000000000000000000000000000000000..31176fc2c2a508388502b45dc27e4d2218f16eec --- /dev/null +++ b/diffsynth/models/flux_ipadapter.py @@ -0,0 +1,110 @@ +from .general_modules import RMSNorm +from transformers import SiglipVisionModel, SiglipVisionConfig +import torch + + +class SiglipVisionModelSO400M(SiglipVisionModel): + def __init__(self): + config = SiglipVisionConfig( + hidden_size=1152, + image_size=384, + intermediate_size=4304, + model_type="siglip_vision_model", + num_attention_heads=16, + num_hidden_layers=27, + patch_size=14, + architectures=["SiglipModel"], + initializer_factor=1.0, + torch_dtype="float32", + transformers_version="4.37.0.dev0" + ) + super().__init__(config) + +class MLPProjModel(torch.nn.Module): + def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4): + super().__init__() + + self.cross_attention_dim = cross_attention_dim + self.num_tokens = num_tokens + + self.proj = torch.nn.Sequential( + torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2), + torch.nn.GELU(), + torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens), + ) + self.norm = torch.nn.LayerNorm(cross_attention_dim) + + def forward(self, id_embeds): + x = self.proj(id_embeds) + x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) + x = self.norm(x) + return x + +class IpAdapterModule(torch.nn.Module): + def __init__(self, num_attention_heads, attention_head_dim, input_dim): + super().__init__() + self.num_heads = num_attention_heads + self.head_dim = attention_head_dim + output_dim = num_attention_heads * attention_head_dim + self.to_k_ip = torch.nn.Linear(input_dim, output_dim, bias=False) + self.to_v_ip = torch.nn.Linear(input_dim, output_dim, bias=False) + self.norm_added_k = RMSNorm(attention_head_dim, eps=1e-5, elementwise_affine=False) + + + def forward(self, hidden_states): + batch_size = hidden_states.shape[0] + # ip_k + ip_k = self.to_k_ip(hidden_states) + ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + ip_k = self.norm_added_k(ip_k) + # ip_v + ip_v = self.to_v_ip(hidden_states) + ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + return ip_k, ip_v + + +class FluxIpAdapter(torch.nn.Module): + def __init__(self, num_attention_heads=24, attention_head_dim=128, cross_attention_dim=4096, num_tokens=128, num_blocks=57): + super().__init__() + self.ipadapter_modules = torch.nn.ModuleList([IpAdapterModule(num_attention_heads, attention_head_dim, cross_attention_dim) for _ in range(num_blocks)]) + self.image_proj = MLPProjModel(cross_attention_dim=cross_attention_dim, id_embeddings_dim=1152, num_tokens=num_tokens) + self.set_adapter() + + def set_adapter(self): + self.call_block_id = {i:i for i in range(len(self.ipadapter_modules))} + + def forward(self, hidden_states, scale=1.0): + hidden_states = self.image_proj(hidden_states) + hidden_states = hidden_states.view(1, -1, hidden_states.shape[-1]) + ip_kv_dict = {} + for block_id in self.call_block_id: + ipadapter_id = self.call_block_id[block_id] + ip_k, ip_v = self.ipadapter_modules[ipadapter_id](hidden_states) + ip_kv_dict[block_id] = { + "ip_k": ip_k, + "ip_v": ip_v, + "scale": scale + } + return ip_kv_dict + + @staticmethod + def state_dict_converter(): + return FluxIpAdapterStateDictConverter() + + +class FluxIpAdapterStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + state_dict_ = {} + for name in state_dict["ip_adapter"]: + name_ = 'ipadapter_modules.' + name + state_dict_[name_] = state_dict["ip_adapter"][name] + for name in state_dict["image_proj"]: + name_ = "image_proj." + name + state_dict_[name_] = state_dict["image_proj"][name] + return state_dict_ + + def from_civitai(self, state_dict): + return self.from_diffusers(state_dict) diff --git a/diffsynth/models/flux_lora_encoder.py b/diffsynth/models/flux_lora_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..13589b0611f3140479ef4faa3b7a29371caa447b --- /dev/null +++ b/diffsynth/models/flux_lora_encoder.py @@ -0,0 +1,521 @@ +import torch +from einops import rearrange + + +def low_version_attention(query, key, value, attn_bias=None): + scale = 1 / query.shape[-1] ** 0.5 + query = query * scale + attn = torch.matmul(query, key.transpose(-2, -1)) + if attn_bias is not None: + attn = attn + attn_bias + attn = attn.softmax(-1) + return attn @ value + + +class Attention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) + self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out) + + def interact_with_ipadapter(self, hidden_states, q, ip_k, ip_v, scale=1.0): + batch_size = q.shape[0] + ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v) + hidden_states = hidden_states + scale * ip_hidden_states + return hidden_states + + def torch_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + batch_size = encoder_hidden_states.shape[0] + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + + if qkv_preprocessor is not None: + q, k, v = qkv_preprocessor(q, k, v) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + if ipadapter_kwargs is not None: + hidden_states = self.interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + + hidden_states = self.to_out(hidden_states) + + return hidden_states + + def xformers_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = rearrange(q, "b f (n d) -> (b n) f d", n=self.num_heads) + k = rearrange(k, "b f (n d) -> (b n) f d", n=self.num_heads) + v = rearrange(v, "b f (n d) -> (b n) f d", n=self.num_heads) + + if attn_mask is not None: + hidden_states = low_version_attention(q, k, v, attn_bias=attn_mask) + else: + import xformers.ops as xops + hidden_states = xops.memory_efficient_attention(q, k, v) + hidden_states = rearrange(hidden_states, "(b n) f d -> b f (n d)", n=self.num_heads) + + hidden_states = hidden_states.to(q.dtype) + hidden_states = self.to_out(hidden_states) + + return hidden_states + + def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None): + return self.torch_forward(hidden_states, encoder_hidden_states=encoder_hidden_states, attn_mask=attn_mask, ipadapter_kwargs=ipadapter_kwargs, qkv_preprocessor=qkv_preprocessor) + + + + + +class CLIPEncoderLayer(torch.nn.Module): + def __init__(self, embed_dim, intermediate_size, num_heads=12, head_dim=64, use_quick_gelu=True): + super().__init__() + self.attn = Attention(q_dim=embed_dim, num_heads=num_heads, head_dim=head_dim, bias_q=True, bias_kv=True, bias_out=True) + self.layer_norm1 = torch.nn.LayerNorm(embed_dim) + self.layer_norm2 = torch.nn.LayerNorm(embed_dim) + self.fc1 = torch.nn.Linear(embed_dim, intermediate_size) + self.fc2 = torch.nn.Linear(intermediate_size, embed_dim) + + self.use_quick_gelu = use_quick_gelu + + def quickGELU(self, x): + return x * torch.sigmoid(1.702 * x) + + def forward(self, hidden_states, attn_mask=None): + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.attn(hidden_states, attn_mask=attn_mask) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.fc1(hidden_states) + if self.use_quick_gelu: + hidden_states = self.quickGELU(hidden_states) + else: + hidden_states = torch.nn.functional.gelu(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class SDTextEncoder(torch.nn.Module): + def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072): + super().__init__() + + # token_embedding + self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim) + + # position_embeds (This is a fixed tensor) + self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim)) + + # encoders + self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)]) + + # attn_mask + self.attn_mask = self.attention_mask(max_position_embeddings) + + # final_layer_norm + self.final_layer_norm = torch.nn.LayerNorm(embed_dim) + + def attention_mask(self, length): + mask = torch.empty(length, length) + mask.fill_(float("-inf")) + mask.triu_(1) + return mask + + def forward(self, input_ids, clip_skip=1): + embeds = self.token_embedding(input_ids) + self.position_embeds + attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype) + for encoder_id, encoder in enumerate(self.encoders): + embeds = encoder(embeds, attn_mask=attn_mask) + if encoder_id + clip_skip == len(self.encoders): + break + embeds = self.final_layer_norm(embeds) + return embeds + + @staticmethod + def state_dict_converter(): + return SDTextEncoderStateDictConverter() + + +class SDTextEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + rename_dict = { + "text_model.embeddings.token_embedding.weight": "token_embedding.weight", + "text_model.embeddings.position_embedding.weight": "position_embeds", + "text_model.final_layer_norm.weight": "final_layer_norm.weight", + "text_model.final_layer_norm.bias": "final_layer_norm.bias" + } + attn_rename_dict = { + "self_attn.q_proj": "attn.to_q", + "self_attn.k_proj": "attn.to_k", + "self_attn.v_proj": "attn.to_v", + "self_attn.out_proj": "attn.to_out", + "layer_norm1": "layer_norm1", + "layer_norm2": "layer_norm2", + "mlp.fc1": "fc1", + "mlp.fc2": "fc2", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + if name == "text_model.embeddings.position_embedding.weight": + param = param.reshape((1, param.shape[0], param.shape[1])) + state_dict_[rename_dict[name]] = param + elif name.startswith("text_model.encoder.layers."): + param = state_dict[name] + names = name.split(".") + layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1] + name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail]) + state_dict_[name_] = param + return state_dict_ + + def from_civitai(self, state_dict): + rename_dict = { + "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight": "token_embedding.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.bias": "encoders.0.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.weight": "encoders.0.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.bias": "encoders.0.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.weight": "encoders.0.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.bias": "encoders.0.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.weight": "encoders.0.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.bias": "encoders.0.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.weight": "encoders.0.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.bias": "encoders.0.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.weight": "encoders.0.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.bias": "encoders.0.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.weight": "encoders.0.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.bias": "encoders.0.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.weight": "encoders.0.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.bias": "encoders.0.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.weight": "encoders.0.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.bias": "encoders.1.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.weight": "encoders.1.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.bias": "encoders.1.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.weight": "encoders.1.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.bias": "encoders.1.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.weight": "encoders.1.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.bias": "encoders.1.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.weight": "encoders.1.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.bias": "encoders.1.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.weight": "encoders.1.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.bias": "encoders.1.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.weight": "encoders.1.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.bias": "encoders.1.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.weight": "encoders.1.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.bias": "encoders.1.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.weight": "encoders.1.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.bias": "encoders.10.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.weight": "encoders.10.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.bias": "encoders.10.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.weight": "encoders.10.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.bias": "encoders.10.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.weight": "encoders.10.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.bias": "encoders.10.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.weight": "encoders.10.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.bias": "encoders.10.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.weight": "encoders.10.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.bias": "encoders.10.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.weight": "encoders.10.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.bias": "encoders.10.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.weight": "encoders.10.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.bias": "encoders.10.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.weight": "encoders.10.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.bias": "encoders.11.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.weight": "encoders.11.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.bias": "encoders.11.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.weight": "encoders.11.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.bias": "encoders.11.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.weight": "encoders.11.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.bias": "encoders.11.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.weight": "encoders.11.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.bias": "encoders.11.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.weight": "encoders.11.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.bias": "encoders.11.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.weight": "encoders.11.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.bias": "encoders.11.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.weight": "encoders.11.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.bias": "encoders.11.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.weight": "encoders.11.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.bias": "encoders.2.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.weight": "encoders.2.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.bias": "encoders.2.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.weight": "encoders.2.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.bias": "encoders.2.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.weight": "encoders.2.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.bias": "encoders.2.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.weight": "encoders.2.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.bias": "encoders.2.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.weight": "encoders.2.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.bias": "encoders.2.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.weight": "encoders.2.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.bias": "encoders.2.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.weight": "encoders.2.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.bias": "encoders.2.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.weight": "encoders.2.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.bias": "encoders.3.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.weight": "encoders.3.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.bias": "encoders.3.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.weight": "encoders.3.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.bias": "encoders.3.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.weight": "encoders.3.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.bias": "encoders.3.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.weight": "encoders.3.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.bias": "encoders.3.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.weight": "encoders.3.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.bias": "encoders.3.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.weight": "encoders.3.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.bias": "encoders.3.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.weight": "encoders.3.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.bias": "encoders.3.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.weight": "encoders.3.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.bias": "encoders.4.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.weight": "encoders.4.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.bias": "encoders.4.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.weight": "encoders.4.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.bias": "encoders.4.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.weight": "encoders.4.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.bias": "encoders.4.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.weight": "encoders.4.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.bias": "encoders.4.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.weight": "encoders.4.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.bias": "encoders.4.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.weight": "encoders.4.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.bias": "encoders.4.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.weight": "encoders.4.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.bias": "encoders.4.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.weight": "encoders.4.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.bias": "encoders.5.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.weight": "encoders.5.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.bias": "encoders.5.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.weight": "encoders.5.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.bias": "encoders.5.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.weight": "encoders.5.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.bias": "encoders.5.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.weight": "encoders.5.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.bias": "encoders.5.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.weight": "encoders.5.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.bias": "encoders.5.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.weight": "encoders.5.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.bias": "encoders.5.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.weight": "encoders.5.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.bias": "encoders.5.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.weight": "encoders.5.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.bias": "encoders.6.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.weight": "encoders.6.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.bias": "encoders.6.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.weight": "encoders.6.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.bias": "encoders.6.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.weight": "encoders.6.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.bias": "encoders.6.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.weight": "encoders.6.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.bias": "encoders.6.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.weight": "encoders.6.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.bias": "encoders.6.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.weight": "encoders.6.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.bias": "encoders.6.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.weight": "encoders.6.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.bias": "encoders.6.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.weight": "encoders.6.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.bias": "encoders.7.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.weight": "encoders.7.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.bias": "encoders.7.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.weight": "encoders.7.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.bias": "encoders.7.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.weight": "encoders.7.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.bias": "encoders.7.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.weight": "encoders.7.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.bias": "encoders.7.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.weight": "encoders.7.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.bias": "encoders.7.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.weight": "encoders.7.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.bias": "encoders.7.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.weight": "encoders.7.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.bias": "encoders.7.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.weight": "encoders.7.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.bias": "encoders.8.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.weight": "encoders.8.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.bias": "encoders.8.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.weight": "encoders.8.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.bias": "encoders.8.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.weight": "encoders.8.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.bias": "encoders.8.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.weight": "encoders.8.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.bias": "encoders.8.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.weight": "encoders.8.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.bias": "encoders.8.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.weight": "encoders.8.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.bias": "encoders.8.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.weight": "encoders.8.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.bias": "encoders.8.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.weight": "encoders.8.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.bias": "encoders.9.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.weight": "encoders.9.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.bias": "encoders.9.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.weight": "encoders.9.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.bias": "encoders.9.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.weight": "encoders.9.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.bias": "encoders.9.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.weight": "encoders.9.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.bias": "encoders.9.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.weight": "encoders.9.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.bias": "encoders.9.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.weight": "encoders.9.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.bias": "encoders.9.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.weight": "encoders.9.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.bias": "encoders.9.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.weight": "encoders.9.attn.to_v.weight", + "cond_stage_model.transformer.text_model.final_layer_norm.bias": "final_layer_norm.bias", + "cond_stage_model.transformer.text_model.final_layer_norm.weight": "final_layer_norm.weight", + "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight": "position_embeds" + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + if name == "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight": + param = param.reshape((1, param.shape[0], param.shape[1])) + state_dict_[rename_dict[name]] = param + return state_dict_ + + + +class LoRALayerBlock(torch.nn.Module): + def __init__(self, L, dim_in, dim_out): + super().__init__() + self.x = torch.nn.Parameter(torch.randn(1, L, dim_in)) + self.layer_norm = torch.nn.LayerNorm(dim_out) + + def forward(self, lora_A, lora_B): + x = self.x @ lora_A.T @ lora_B.T + x = self.layer_norm(x) + return x + + +class LoRAEmbedder(torch.nn.Module): + def __init__(self, lora_patterns=None, L=1, out_dim=2048): + super().__init__() + if lora_patterns is None: + lora_patterns = self.default_lora_patterns() + + model_dict = {} + for lora_pattern in lora_patterns: + name, dim = lora_pattern["name"], lora_pattern["dim"] + model_dict[name.replace(".", "___")] = LoRALayerBlock(L, dim[0], dim[1]) + self.model_dict = torch.nn.ModuleDict(model_dict) + + proj_dict = {} + for lora_pattern in lora_patterns: + layer_type, dim = lora_pattern["type"], lora_pattern["dim"] + if layer_type not in proj_dict: + proj_dict[layer_type.replace(".", "___")] = torch.nn.Linear(dim[1], out_dim) + self.proj_dict = torch.nn.ModuleDict(proj_dict) + + self.lora_patterns = lora_patterns + + + def default_lora_patterns(self): + lora_patterns = [] + lora_dict = { + "attn.a_to_qkv": (3072, 9216), "attn.a_to_out": (3072, 3072), "ff_a.0": (3072, 12288), "ff_a.2": (12288, 3072), "norm1_a.linear": (3072, 18432), + "attn.b_to_qkv": (3072, 9216), "attn.b_to_out": (3072, 3072), "ff_b.0": (3072, 12288), "ff_b.2": (12288, 3072), "norm1_b.linear": (3072, 18432), + } + for i in range(19): + for suffix in lora_dict: + lora_patterns.append({ + "name": f"blocks.{i}.{suffix}", + "dim": lora_dict[suffix], + "type": suffix, + }) + lora_dict = {"to_qkv_mlp": (3072, 21504), "proj_out": (15360, 3072), "norm.linear": (3072, 9216)} + for i in range(38): + for suffix in lora_dict: + lora_patterns.append({ + "name": f"single_blocks.{i}.{suffix}", + "dim": lora_dict[suffix], + "type": suffix, + }) + return lora_patterns + + def forward(self, lora): + lora_emb = [] + for lora_pattern in self.lora_patterns: + name, layer_type = lora_pattern["name"], lora_pattern["type"] + lora_A = lora[name + ".lora_A.weight"] + lora_B = lora[name + ".lora_B.weight"] + lora_out = self.model_dict[name.replace(".", "___")](lora_A, lora_B) + lora_out = self.proj_dict[layer_type.replace(".", "___")](lora_out) + lora_emb.append(lora_out) + lora_emb = torch.concat(lora_emb, dim=1) + return lora_emb + + +class FluxLoRAEncoder(torch.nn.Module): + def __init__(self, embed_dim=4096, encoder_intermediate_size=8192, num_encoder_layers=1, num_embeds_per_lora=16, num_special_embeds=1): + super().__init__() + self.num_embeds_per_lora = num_embeds_per_lora + # embedder + self.embedder = LoRAEmbedder(L=num_embeds_per_lora, out_dim=embed_dim) + + # encoders + self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size, num_heads=32, head_dim=128) for _ in range(num_encoder_layers)]) + + # special embedding + self.special_embeds = torch.nn.Parameter(torch.randn(1, num_special_embeds, embed_dim)) + self.num_special_embeds = num_special_embeds + + # final layer + self.final_layer_norm = torch.nn.LayerNorm(embed_dim) + self.final_linear = torch.nn.Linear(embed_dim, embed_dim) + + def forward(self, lora): + lora_embeds = self.embedder(lora) + special_embeds = self.special_embeds.to(dtype=lora_embeds.dtype, device=lora_embeds.device) + embeds = torch.concat([special_embeds, lora_embeds], dim=1) + for encoder_id, encoder in enumerate(self.encoders): + embeds = encoder(embeds) + embeds = embeds[:, :self.num_special_embeds] + embeds = self.final_layer_norm(embeds) + embeds = self.final_linear(embeds) + return embeds + + @staticmethod + def state_dict_converter(): + return FluxLoRAEncoderStateDictConverter() + + +class FluxLoRAEncoderStateDictConverter: + def from_civitai(self, state_dict): + return state_dict diff --git a/diffsynth/models/flux_lora_patcher.py b/diffsynth/models/flux_lora_patcher.py new file mode 100644 index 0000000000000000000000000000000000000000..5d8fc8cea03bbbc658d8e1869432d903b6ae7ce9 --- /dev/null +++ b/diffsynth/models/flux_lora_patcher.py @@ -0,0 +1,306 @@ +import torch, math +from ..core.loader import load_state_dict +from typing import Union + +class GeneralLoRALoader: + def __init__(self, device="cpu", torch_dtype=torch.float32): + self.device = device + self.torch_dtype = torch_dtype + + + def get_name_dict(self, lora_state_dict): + lora_name_dict = {} + for key in lora_state_dict: + if ".lora_B." not in key: + continue + keys = key.split(".") + if len(keys) > keys.index("lora_B") + 2: + keys.pop(keys.index("lora_B") + 1) + keys.pop(keys.index("lora_B")) + if keys[0] == "diffusion_model": + keys.pop(0) + keys.pop(-1) + target_name = ".".join(keys) + lora_name_dict[target_name] = (key, key.replace(".lora_B.", ".lora_A.")) + return lora_name_dict + + + def load(self, model: torch.nn.Module, state_dict_lora, alpha=1.0): + updated_num = 0 + lora_name_dict = self.get_name_dict(state_dict_lora) + for name, module in model.named_modules(): + if name in lora_name_dict: + weight_up = state_dict_lora[lora_name_dict[name][0]].to(device=self.device, dtype=self.torch_dtype) + weight_down = state_dict_lora[lora_name_dict[name][1]].to(device=self.device, dtype=self.torch_dtype) + if len(weight_up.shape) == 4: + weight_up = weight_up.squeeze(3).squeeze(2) + weight_down = weight_down.squeeze(3).squeeze(2) + weight_lora = alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) + else: + weight_lora = alpha * torch.mm(weight_up, weight_down) + state_dict = module.state_dict() + state_dict["weight"] = state_dict["weight"].to(device=self.device, dtype=self.torch_dtype) + weight_lora + module.load_state_dict(state_dict) + updated_num += 1 + print(f"{updated_num} tensors are updated by LoRA.") + +class FluxLoRALoader(GeneralLoRALoader): + def __init__(self, device="cpu", torch_dtype=torch.float32): + super().__init__(device=device, torch_dtype=torch_dtype) + + self.diffusers_rename_dict = { + "transformer.single_transformer_blocks.blockid.attn.to_k.lora_A.weight":"single_blocks.blockid.a_to_k.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.attn.to_k.lora_B.weight":"single_blocks.blockid.a_to_k.lora_B.default.weight", + "transformer.single_transformer_blocks.blockid.attn.to_q.lora_A.weight":"single_blocks.blockid.a_to_q.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.attn.to_q.lora_B.weight":"single_blocks.blockid.a_to_q.lora_B.default.weight", + "transformer.single_transformer_blocks.blockid.attn.to_v.lora_A.weight":"single_blocks.blockid.a_to_v.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.attn.to_v.lora_B.weight":"single_blocks.blockid.a_to_v.lora_B.default.weight", + "transformer.single_transformer_blocks.blockid.norm.linear.lora_A.weight":"single_blocks.blockid.norm.linear.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.norm.linear.lora_B.weight":"single_blocks.blockid.norm.linear.lora_B.default.weight", + "transformer.single_transformer_blocks.blockid.proj_mlp.lora_A.weight":"single_blocks.blockid.proj_in_besides_attn.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.proj_mlp.lora_B.weight":"single_blocks.blockid.proj_in_besides_attn.lora_B.default.weight", + "transformer.single_transformer_blocks.blockid.proj_out.lora_A.weight":"single_blocks.blockid.proj_out.lora_A.default.weight", + "transformer.single_transformer_blocks.blockid.proj_out.lora_B.weight":"single_blocks.blockid.proj_out.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_A.weight":"blocks.blockid.attn.b_to_k.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_B.weight":"blocks.blockid.attn.b_to_k.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_A.weight":"blocks.blockid.attn.b_to_q.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_B.weight":"blocks.blockid.attn.b_to_q.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_A.weight":"blocks.blockid.attn.b_to_v.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_B.weight":"blocks.blockid.attn.b_to_v.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.to_add_out.lora_A.weight":"blocks.blockid.attn.b_to_out.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.to_add_out.lora_B.weight":"blocks.blockid.attn.b_to_out.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.to_k.lora_A.weight":"blocks.blockid.attn.a_to_k.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.to_k.lora_B.weight":"blocks.blockid.attn.a_to_k.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.to_out.0.lora_A.weight":"blocks.blockid.attn.a_to_out.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.to_out.0.lora_B.weight":"blocks.blockid.attn.a_to_out.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.to_q.lora_A.weight":"blocks.blockid.attn.a_to_q.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.to_q.lora_B.weight":"blocks.blockid.attn.a_to_q.lora_B.default.weight", + "transformer.transformer_blocks.blockid.attn.to_v.lora_A.weight":"blocks.blockid.attn.a_to_v.lora_A.default.weight", + "transformer.transformer_blocks.blockid.attn.to_v.lora_B.weight":"blocks.blockid.attn.a_to_v.lora_B.default.weight", + "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_A.weight":"blocks.blockid.ff_a.0.lora_A.default.weight", + "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_B.weight":"blocks.blockid.ff_a.0.lora_B.default.weight", + "transformer.transformer_blocks.blockid.ff.net.2.lora_A.weight":"blocks.blockid.ff_a.2.lora_A.default.weight", + "transformer.transformer_blocks.blockid.ff.net.2.lora_B.weight":"blocks.blockid.ff_a.2.lora_B.default.weight", + "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_A.weight":"blocks.blockid.ff_b.0.lora_A.default.weight", + "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_B.weight":"blocks.blockid.ff_b.0.lora_B.default.weight", + "transformer.transformer_blocks.blockid.ff_context.net.2.lora_A.weight":"blocks.blockid.ff_b.2.lora_A.default.weight", + "transformer.transformer_blocks.blockid.ff_context.net.2.lora_B.weight":"blocks.blockid.ff_b.2.lora_B.default.weight", + "transformer.transformer_blocks.blockid.norm1.linear.lora_A.weight":"blocks.blockid.norm1_a.linear.lora_A.default.weight", + "transformer.transformer_blocks.blockid.norm1.linear.lora_B.weight":"blocks.blockid.norm1_a.linear.lora_B.default.weight", + "transformer.transformer_blocks.blockid.norm1_context.linear.lora_A.weight":"blocks.blockid.norm1_b.linear.lora_A.default.weight", + "transformer.transformer_blocks.blockid.norm1_context.linear.lora_B.weight":"blocks.blockid.norm1_b.linear.lora_B.default.weight", + } + + self.civitai_rename_dict = { + "lora_unet_double_blocks_blockid_img_mod_lin.lora_down.weight": "blocks.blockid.norm1_a.linear.lora_A.default.weight", + "lora_unet_double_blocks_blockid_img_mod_lin.lora_up.weight": "blocks.blockid.norm1_a.linear.lora_B.default.weight", + "lora_unet_double_blocks_blockid_txt_mod_lin.lora_down.weight": "blocks.blockid.norm1_b.linear.lora_A.default.weight", + "lora_unet_double_blocks_blockid_txt_mod_lin.lora_up.weight": "blocks.blockid.norm1_b.linear.lora_B.default.weight", + "lora_unet_double_blocks_blockid_img_attn_qkv.lora_down.weight": "blocks.blockid.attn.a_to_qkv.lora_A.default.weight", + "lora_unet_double_blocks_blockid_img_attn_qkv.lora_up.weight": "blocks.blockid.attn.a_to_qkv.lora_B.default.weight", + "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_down.weight": "blocks.blockid.attn.b_to_qkv.lora_A.default.weight", + "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_up.weight": "blocks.blockid.attn.b_to_qkv.lora_B.default.weight", + "lora_unet_double_blocks_blockid_img_attn_proj.lora_down.weight": "blocks.blockid.attn.a_to_out.lora_A.default.weight", + "lora_unet_double_blocks_blockid_img_attn_proj.lora_up.weight": "blocks.blockid.attn.a_to_out.lora_B.default.weight", + "lora_unet_double_blocks_blockid_txt_attn_proj.lora_down.weight": "blocks.blockid.attn.b_to_out.lora_A.default.weight", + "lora_unet_double_blocks_blockid_txt_attn_proj.lora_up.weight": "blocks.blockid.attn.b_to_out.lora_B.default.weight", + "lora_unet_double_blocks_blockid_img_mlp_0.lora_down.weight": "blocks.blockid.ff_a.0.lora_A.default.weight", + "lora_unet_double_blocks_blockid_img_mlp_0.lora_up.weight": "blocks.blockid.ff_a.0.lora_B.default.weight", + "lora_unet_double_blocks_blockid_img_mlp_2.lora_down.weight": "blocks.blockid.ff_a.2.lora_A.default.weight", + "lora_unet_double_blocks_blockid_img_mlp_2.lora_up.weight": "blocks.blockid.ff_a.2.lora_B.default.weight", + "lora_unet_double_blocks_blockid_txt_mlp_0.lora_down.weight": "blocks.blockid.ff_b.0.lora_A.default.weight", + "lora_unet_double_blocks_blockid_txt_mlp_0.lora_up.weight": "blocks.blockid.ff_b.0.lora_B.default.weight", + "lora_unet_double_blocks_blockid_txt_mlp_2.lora_down.weight": "blocks.blockid.ff_b.2.lora_A.default.weight", + "lora_unet_double_blocks_blockid_txt_mlp_2.lora_up.weight": "blocks.blockid.ff_b.2.lora_B.default.weight", + "lora_unet_single_blocks_blockid_modulation_lin.lora_down.weight": "single_blocks.blockid.norm.linear.lora_A.default.weight", + "lora_unet_single_blocks_blockid_modulation_lin.lora_up.weight": "single_blocks.blockid.norm.linear.lora_B.default.weight", + "lora_unet_single_blocks_blockid_linear1.lora_down.weight": "single_blocks.blockid.to_qkv_mlp.lora_A.default.weight", + "lora_unet_single_blocks_blockid_linear1.lora_up.weight": "single_blocks.blockid.to_qkv_mlp.lora_B.default.weight", + "lora_unet_single_blocks_blockid_linear2.lora_down.weight": "single_blocks.blockid.proj_out.lora_A.default.weight", + "lora_unet_single_blocks_blockid_linear2.lora_up.weight": "single_blocks.blockid.proj_out.lora_B.default.weight", + } + + def load(self, model: torch.nn.Module, state_dict_lora, alpha=1.0): + super().load(model, state_dict_lora, alpha) + + + def convert_state_dict(self,state_dict): + + def guess_block_id(name,model_resource): + if model_resource == 'civitai': + names = name.split("_") + for i in names: + if i.isdigit(): + return i, name.replace(f"_{i}_", "_blockid_") + if model_resource == 'diffusers': + names = name.split(".") + for i in names: + if i.isdigit(): + return i, name.replace(f"transformer_blocks.{i}.", "transformer_blocks.blockid.") + return None, None + + def guess_resource(state_dict): + for k in state_dict: + if "lora_unet_" in k: + return 'civitai' + elif k.startswith("transformer."): + return 'diffusers' + else: + None + + model_resource = guess_resource(state_dict) + if model_resource is None: + return state_dict + + rename_dict = self.diffusers_rename_dict if model_resource == 'diffusers' else self.civitai_rename_dict + def guess_alpha(state_dict): + for name, param in state_dict.items(): + if ".alpha" in name: + for suffix in [".lora_down.weight", ".lora_A.weight"]: + name_ = name.replace(".alpha", suffix) + if name_ in state_dict: + lora_alpha = param.item() / state_dict[name_].shape[0] + lora_alpha = math.sqrt(lora_alpha) + return lora_alpha + + return 1 + + alpha = guess_alpha(state_dict) + + state_dict_ = {} + for name, param in state_dict.items(): + block_id, source_name = guess_block_id(name,model_resource) + if alpha != 1: + param *= alpha + if source_name in rename_dict: + target_name = rename_dict[source_name] + target_name = target_name.replace(".blockid.", f".{block_id}.") + state_dict_[target_name] = param + else: + state_dict_[name] = param + + if model_resource == 'diffusers': + for name in list(state_dict_.keys()): + if "single_blocks." in name and ".a_to_q." in name: + mlp = state_dict_.get(name.replace(".a_to_q.", ".proj_in_besides_attn."), None) + if mlp is None: + dim = 4 + if 'lora_A' in name: + dim = 1 + mlp = torch.zeros(dim * state_dict_[name].shape[0], + *state_dict_[name].shape[1:], + dtype=state_dict_[name].dtype) + else: + state_dict_.pop(name.replace(".a_to_q.", ".proj_in_besides_attn.")) + if 'lora_A' in name: + param = torch.concat([ + state_dict_.pop(name), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")), + mlp, + ], dim=0) + elif 'lora_B' in name: + d, r = state_dict_[name].shape + param = torch.zeros((3*d+mlp.shape[0], 3*r+mlp.shape[1]), dtype=state_dict_[name].dtype, device=state_dict_[name].device) + param[:d, :r] = state_dict_.pop(name) + param[d:2*d, r:2*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")) + param[2*d:3*d, 2*r:3*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")) + param[3*d:, 3*r:] = mlp + else: + param = torch.concat([ + state_dict_.pop(name), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")), + mlp, + ], dim=0) + name_ = name.replace(".a_to_q.", ".to_qkv_mlp.") + state_dict_[name_] = param + for name in list(state_dict_.keys()): + for component in ["a", "b"]: + if f".{component}_to_q." in name: + name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.") + concat_dim = 0 + if 'lora_A' in name: + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + elif 'lora_B' in name: + origin = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")] + d, r = origin.shape + # print(d, r) + param = torch.zeros((3*d, 3*r), dtype=origin.dtype, device=origin.device) + param[:d, :r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")] + param[d:2*d, r:2*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")] + param[2*d:3*d, 2*r:3*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")] + else: + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v.")) + return state_dict_ + + +class LoraMerger(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.weight_base = torch.nn.Parameter(torch.randn((dim,))) + self.weight_lora = torch.nn.Parameter(torch.randn((dim,))) + self.weight_cross = torch.nn.Parameter(torch.randn((dim,))) + self.weight_out = torch.nn.Parameter(torch.ones((dim,))) + self.bias = torch.nn.Parameter(torch.randn((dim,))) + self.activation = torch.nn.Sigmoid() + self.norm_base = torch.nn.LayerNorm(dim, eps=1e-5) + self.norm_lora = torch.nn.LayerNorm(dim, eps=1e-5) + + def forward(self, base_output, lora_outputs): + norm_base_output = self.norm_base(base_output) + norm_lora_outputs = self.norm_lora(lora_outputs) + gate = self.activation( + norm_base_output * self.weight_base \ + + norm_lora_outputs * self.weight_lora \ + + norm_base_output * norm_lora_outputs * self.weight_cross + self.bias + ) + output = base_output + (self.weight_out * gate * lora_outputs).sum(dim=0) + return output + +class FluxLoraPatcher(torch.nn.Module): + def __init__(self, lora_patterns=None): + super().__init__() + if lora_patterns is None: + lora_patterns = self.default_lora_patterns() + model_dict = {} + for lora_pattern in lora_patterns: + name, dim = lora_pattern["name"], lora_pattern["dim"] + model_dict[name.replace(".", "___")] = LoraMerger(dim) + self.model_dict = torch.nn.ModuleDict(model_dict) + + def default_lora_patterns(self): + lora_patterns = [] + lora_dict = { + "attn.a_to_qkv": 9216, "attn.a_to_out": 3072, "ff_a.0": 12288, "ff_a.2": 3072, "norm1_a.linear": 18432, + "attn.b_to_qkv": 9216, "attn.b_to_out": 3072, "ff_b.0": 12288, "ff_b.2": 3072, "norm1_b.linear": 18432, + } + for i in range(19): + for suffix in lora_dict: + lora_patterns.append({ + "name": f"blocks.{i}.{suffix}", + "dim": lora_dict[suffix] + }) + lora_dict = {"to_qkv_mlp": 21504, "proj_out": 3072, "norm.linear": 9216} + for i in range(38): + for suffix in lora_dict: + lora_patterns.append({ + "name": f"single_blocks.{i}.{suffix}", + "dim": lora_dict[suffix] + }) + return lora_patterns + + def forward(self, base_output, lora_outputs, name): + return self.model_dict[name.replace(".", "___")](base_output, lora_outputs) diff --git a/diffsynth/models/flux_text_encoder_clip.py b/diffsynth/models/flux_text_encoder_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..1425423ce6d1df946198a16a7e96078ab8fed807 --- /dev/null +++ b/diffsynth/models/flux_text_encoder_clip.py @@ -0,0 +1,112 @@ +import torch + + +class Attention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) + self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out) + + def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + batch_size = encoder_hidden_states.shape[0] + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + + hidden_states = self.to_out(hidden_states) + + return hidden_states + + +class CLIPEncoderLayer(torch.nn.Module): + def __init__(self, embed_dim, intermediate_size, num_heads=12, head_dim=64, use_quick_gelu=True): + super().__init__() + self.attn = Attention(q_dim=embed_dim, num_heads=num_heads, head_dim=head_dim, bias_q=True, bias_kv=True, bias_out=True) + self.layer_norm1 = torch.nn.LayerNorm(embed_dim) + self.layer_norm2 = torch.nn.LayerNorm(embed_dim) + self.fc1 = torch.nn.Linear(embed_dim, intermediate_size) + self.fc2 = torch.nn.Linear(intermediate_size, embed_dim) + + self.use_quick_gelu = use_quick_gelu + + def quickGELU(self, x): + return x * torch.sigmoid(1.702 * x) + + def forward(self, hidden_states, attn_mask=None): + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.attn(hidden_states, attn_mask=attn_mask) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.fc1(hidden_states) + if self.use_quick_gelu: + hidden_states = self.quickGELU(hidden_states) + else: + hidden_states = torch.nn.functional.gelu(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class FluxTextEncoderClip(torch.nn.Module): + def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072): + super().__init__() + + # token_embedding + self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim) + + # position_embeds (This is a fixed tensor) + self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim)) + + # encoders + self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)]) + + # attn_mask + self.attn_mask = self.attention_mask(max_position_embeddings) + + # final_layer_norm + self.final_layer_norm = torch.nn.LayerNorm(embed_dim) + + def attention_mask(self, length): + mask = torch.empty(length, length) + mask.fill_(float("-inf")) + mask.triu_(1) + return mask + + def forward(self, input_ids, clip_skip=2, extra_mask=None): + embeds = self.token_embedding(input_ids) + embeds = embeds + self.position_embeds.to(dtype=embeds.dtype, device=input_ids.device) + attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype) + if extra_mask is not None: + attn_mask[:, extra_mask[0]==0] = float("-inf") + for encoder_id, encoder in enumerate(self.encoders): + embeds = encoder(embeds, attn_mask=attn_mask) + if encoder_id + clip_skip == len(self.encoders): + hidden_states = embeds + embeds = self.final_layer_norm(embeds) + pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)] + return pooled_embeds, hidden_states diff --git a/diffsynth/models/flux_text_encoder_t5.py b/diffsynth/models/flux_text_encoder_t5.py new file mode 100644 index 0000000000000000000000000000000000000000..ee72e4a89b2089b62c6ea86c4aef91755ee9ee9a --- /dev/null +++ b/diffsynth/models/flux_text_encoder_t5.py @@ -0,0 +1,43 @@ +import torch +from transformers import T5EncoderModel, T5Config + + +class FluxTextEncoderT5(T5EncoderModel): + def __init__(self): + config = T5Config(**{ + "architectures": [ + "T5EncoderModel" + ], + "classifier_dropout": 0.0, + "d_ff": 10240, + "d_kv": 64, + "d_model": 4096, + "decoder_start_token_id": 0, + "dense_act_fn": "gelu_new", + "dropout_rate": 0.1, + "dtype": "bfloat16", + "eos_token_id": 1, + "feed_forward_proj": "gated-gelu", + "initializer_factor": 1.0, + "is_encoder_decoder": True, + "is_gated_act": True, + "layer_norm_epsilon": 1e-06, + "model_type": "t5", + "num_decoder_layers": 24, + "num_heads": 64, + "num_layers": 24, + "output_past": True, + "pad_token_id": 0, + "relative_attention_max_distance": 128, + "relative_attention_num_buckets": 32, + "tie_word_embeddings": False, + "transformers_version": "4.57.1", + "use_cache": True, + "vocab_size": 32128 + }) + super().__init__(config) + + def forward(self, input_ids): + outputs = super().forward(input_ids=input_ids) + prompt_emb = outputs.last_hidden_state + return prompt_emb diff --git a/diffsynth/models/flux_vae.py b/diffsynth/models/flux_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..5eabeaee6ad54f0b1c02f1b24cd2ccd84d0238bf --- /dev/null +++ b/diffsynth/models/flux_vae.py @@ -0,0 +1,451 @@ +import torch +from einops import rearrange, repeat + + +class TileWorker: + def __init__(self): + pass + + + def mask(self, height, width, border_width): + # Create a mask with shape (height, width). + # The centre area is filled with 1, and the border line is filled with values in range (0, 1]. + x = torch.arange(height).repeat(width, 1).T + y = torch.arange(width).repeat(height, 1) + mask = torch.stack([x + 1, height - x, y + 1, width - y]).min(dim=0).values + mask = (mask / border_width).clip(0, 1) + return mask + + + def tile(self, model_input, tile_size, tile_stride, tile_device, tile_dtype): + # Convert a tensor (b, c, h, w) to (b, c, tile_size, tile_size, tile_num) + batch_size, channel, _, _ = model_input.shape + model_input = model_input.to(device=tile_device, dtype=tile_dtype) + unfold_operator = torch.nn.Unfold( + kernel_size=(tile_size, tile_size), + stride=(tile_stride, tile_stride) + ) + model_input = unfold_operator(model_input) + model_input = model_input.view((batch_size, channel, tile_size, tile_size, -1)) + + return model_input + + + def tiled_inference(self, forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype): + # Call y=forward_fn(x) for each tile + tile_num = model_input.shape[-1] + model_output_stack = [] + + for tile_id in range(0, tile_num, tile_batch_size): + + # process input + tile_id_ = min(tile_id + tile_batch_size, tile_num) + x = model_input[:, :, :, :, tile_id: tile_id_] + x = x.to(device=inference_device, dtype=inference_dtype) + x = rearrange(x, "b c h w n -> (n b) c h w") + + # process output + y = forward_fn(x) + y = rearrange(y, "(n b) c h w -> b c h w n", n=tile_id_-tile_id) + y = y.to(device=tile_device, dtype=tile_dtype) + model_output_stack.append(y) + + model_output = torch.concat(model_output_stack, dim=-1) + return model_output + + + def io_scale(self, model_output, tile_size): + # Determine the size modification happened in forward_fn + # We only consider the same scale on height and width. + io_scale = model_output.shape[2] / tile_size + return io_scale + + + def untile(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype): + # The reversed function of tile + mask = self.mask(tile_size, tile_size, border_width) + mask = mask.to(device=tile_device, dtype=tile_dtype) + mask = rearrange(mask, "h w -> 1 1 h w 1") + model_output = model_output * mask + + fold_operator = torch.nn.Fold( + output_size=(height, width), + kernel_size=(tile_size, tile_size), + stride=(tile_stride, tile_stride) + ) + mask = repeat(mask[0, 0, :, :, 0], "h w -> 1 (h w) n", n=model_output.shape[-1]) + model_output = rearrange(model_output, "b c h w n -> b (c h w) n") + model_output = fold_operator(model_output) / fold_operator(mask) + + return model_output + + + def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_batch_size=1, tile_device="cpu", tile_dtype=torch.float32, border_width=None): + # Prepare + inference_device, inference_dtype = model_input.device, model_input.dtype + height, width = model_input.shape[2], model_input.shape[3] + border_width = int(tile_stride*0.5) if border_width is None else border_width + + # tile + model_input = self.tile(model_input, tile_size, tile_stride, tile_device, tile_dtype) + + # inference + model_output = self.tiled_inference(forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype) + + # resize + io_scale = self.io_scale(model_output, tile_size) + height, width = int(height*io_scale), int(width*io_scale) + tile_size, tile_stride = int(tile_size*io_scale), int(tile_stride*io_scale) + border_width = int(border_width*io_scale) + + # untile + model_output = self.untile(model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype) + + # Done! + model_output = model_output.to(device=inference_device, dtype=inference_dtype) + return model_output + + +class ConvAttention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Conv2d(q_dim, dim_inner, kernel_size=(1, 1), bias=bias_q) + self.to_k = torch.nn.Conv2d(kv_dim, dim_inner, kernel_size=(1, 1), bias=bias_kv) + self.to_v = torch.nn.Conv2d(kv_dim, dim_inner, kernel_size=(1, 1), bias=bias_kv) + self.to_out = torch.nn.Conv2d(dim_inner, q_dim, kernel_size=(1, 1), bias=bias_out) + + def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + batch_size = encoder_hidden_states.shape[0] + + conv_input = rearrange(hidden_states, "B L C -> B C L 1") + q = self.to_q(conv_input) + q = rearrange(q[:, :, :, 0], "B C L -> B L C") + conv_input = rearrange(encoder_hidden_states, "B L C -> B C L 1") + k = self.to_k(conv_input) + v = self.to_v(conv_input) + k = rearrange(k[:, :, :, 0], "B C L -> B L C") + v = rearrange(v[:, :, :, 0], "B C L -> B L C") + + q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + + conv_input = rearrange(hidden_states, "B L C -> B C L 1") + hidden_states = self.to_out(conv_input) + hidden_states = rearrange(hidden_states[:, :, :, 0], "B C L -> B L C") + + return hidden_states + + +class Attention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) + self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out) + + def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + batch_size = encoder_hidden_states.shape[0] + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + + hidden_states = self.to_out(hidden_states) + + return hidden_states + + +class VAEAttentionBlock(torch.nn.Module): + + def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, norm_num_groups=32, eps=1e-5, use_conv_attention=True): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + + self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True) + + if use_conv_attention: + self.transformer_blocks = torch.nn.ModuleList([ + ConvAttention( + inner_dim, + num_attention_heads, + attention_head_dim, + bias_q=True, + bias_kv=True, + bias_out=True + ) + for d in range(num_layers) + ]) + else: + self.transformer_blocks = torch.nn.ModuleList([ + Attention( + inner_dim, + num_attention_heads, + attention_head_dim, + bias_q=True, + bias_kv=True, + bias_out=True + ) + for d in range(num_layers) + ]) + + def forward(self, hidden_states, time_emb, text_emb, res_stack): + batch, _, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) + + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + + hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() + hidden_states = hidden_states + residual + + return hidden_states, time_emb, text_emb, res_stack + + +class ResnetBlock(torch.nn.Module): + def __init__(self, in_channels, out_channels, temb_channels=None, groups=32, eps=1e-5): + super().__init__() + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + if temb_channels is not None: + self.time_emb_proj = torch.nn.Linear(temb_channels, out_channels) + self.norm2 = torch.nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) + self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.nonlinearity = torch.nn.SiLU() + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=True) + + def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): + x = hidden_states + x = self.norm1(x) + x = self.nonlinearity(x) + x = self.conv1(x) + if time_emb is not None: + emb = self.nonlinearity(time_emb) + emb = self.time_emb_proj(emb)[:, :, None, None] + x = x + emb + x = self.norm2(x) + x = self.nonlinearity(x) + x = self.conv2(x) + if self.conv_shortcut is not None: + hidden_states = self.conv_shortcut(hidden_states) + hidden_states = hidden_states + x + return hidden_states, time_emb, text_emb, res_stack + + +class UpSampler(torch.nn.Module): + def __init__(self, channels): + super().__init__() + self.conv = torch.nn.Conv2d(channels, channels, 3, padding=1) + + def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): + hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest") + hidden_states = self.conv(hidden_states) + return hidden_states, time_emb, text_emb, res_stack + + +class DownSampler(torch.nn.Module): + def __init__(self, channels, padding=1, extra_padding=False): + super().__init__() + self.conv = torch.nn.Conv2d(channels, channels, 3, stride=2, padding=padding) + self.extra_padding = extra_padding + + def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): + if self.extra_padding: + hidden_states = torch.nn.functional.pad(hidden_states, (0, 1, 0, 1), mode="constant", value=0) + hidden_states = self.conv(hidden_states) + return hidden_states, time_emb, text_emb, res_stack + + +class FluxVAEDecoder(torch.nn.Module): + def __init__(self, use_conv_attention=True): + super().__init__() + self.scaling_factor = 0.3611 + self.shift_factor = 0.1159 + self.conv_in = torch.nn.Conv2d(16, 512, kernel_size=3, padding=1) # Different from SD 1.x + + self.blocks = torch.nn.ModuleList([ + # UNetMidBlock2D + ResnetBlock(512, 512, eps=1e-6), + VAEAttentionBlock(1, 512, 512, 1, eps=1e-6, use_conv_attention=use_conv_attention), + ResnetBlock(512, 512, eps=1e-6), + # UpDecoderBlock2D + ResnetBlock(512, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + UpSampler(512), + # UpDecoderBlock2D + ResnetBlock(512, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + UpSampler(512), + # UpDecoderBlock2D + ResnetBlock(512, 256, eps=1e-6), + ResnetBlock(256, 256, eps=1e-6), + ResnetBlock(256, 256, eps=1e-6), + UpSampler(256), + # UpDecoderBlock2D + ResnetBlock(256, 128, eps=1e-6), + ResnetBlock(128, 128, eps=1e-6), + ResnetBlock(128, 128, eps=1e-6), + ]) + + self.conv_norm_out = torch.nn.GroupNorm(num_channels=128, num_groups=32, eps=1e-6) + self.conv_act = torch.nn.SiLU() + self.conv_out = torch.nn.Conv2d(128, 3, kernel_size=3, padding=1) + + def tiled_forward(self, sample, tile_size=64, tile_stride=32): + hidden_states = TileWorker().tiled_forward( + lambda x: self.forward(x), + sample, + tile_size, + tile_stride, + tile_device=sample.device, + tile_dtype=sample.dtype + ) + return hidden_states + + def forward(self, sample, tiled=False, tile_size=64, tile_stride=32, **kwargs): + # For VAE Decoder, we do not need to apply the tiler on each layer. + if tiled: + return self.tiled_forward(sample, tile_size=tile_size, tile_stride=tile_stride) + + # 1. pre-process + hidden_states = sample / self.scaling_factor + self.shift_factor + hidden_states = self.conv_in(hidden_states) + time_emb = None + text_emb = None + res_stack = None + + # 2. blocks + for i, block in enumerate(self.blocks): + hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack) + + # 3. output + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class FluxVAEEncoder(torch.nn.Module): + def __init__(self, use_conv_attention=True): + super().__init__() + self.scaling_factor = 0.3611 + self.shift_factor = 0.1159 + self.conv_in = torch.nn.Conv2d(3, 128, kernel_size=3, padding=1) + + self.blocks = torch.nn.ModuleList([ + # DownEncoderBlock2D + ResnetBlock(128, 128, eps=1e-6), + ResnetBlock(128, 128, eps=1e-6), + DownSampler(128, padding=0, extra_padding=True), + # DownEncoderBlock2D + ResnetBlock(128, 256, eps=1e-6), + ResnetBlock(256, 256, eps=1e-6), + DownSampler(256, padding=0, extra_padding=True), + # DownEncoderBlock2D + ResnetBlock(256, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + DownSampler(512, padding=0, extra_padding=True), + # DownEncoderBlock2D + ResnetBlock(512, 512, eps=1e-6), + ResnetBlock(512, 512, eps=1e-6), + # UNetMidBlock2D + ResnetBlock(512, 512, eps=1e-6), + VAEAttentionBlock(1, 512, 512, 1, eps=1e-6, use_conv_attention=use_conv_attention), + ResnetBlock(512, 512, eps=1e-6), + ]) + + self.conv_norm_out = torch.nn.GroupNorm(num_channels=512, num_groups=32, eps=1e-6) + self.conv_act = torch.nn.SiLU() + self.conv_out = torch.nn.Conv2d(512, 32, kernel_size=3, padding=1) + + def tiled_forward(self, sample, tile_size=64, tile_stride=32): + hidden_states = TileWorker().tiled_forward( + lambda x: self.forward(x), + sample, + tile_size, + tile_stride, + tile_device=sample.device, + tile_dtype=sample.dtype + ) + return hidden_states + + def forward(self, sample, tiled=False, tile_size=64, tile_stride=32, **kwargs): + # For VAE Decoder, we do not need to apply the tiler on each layer. + if tiled: + return self.tiled_forward(sample, tile_size=tile_size, tile_stride=tile_stride) + + # 1. pre-process + hidden_states = self.conv_in(sample) + time_emb = None + text_emb = None + res_stack = None + + # 2. blocks + for i, block in enumerate(self.blocks): + hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack) + + # 3. output + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + hidden_states = hidden_states[:, :16] + hidden_states = (hidden_states - self.shift_factor) * self.scaling_factor + + return hidden_states + + def encode_video(self, sample, batch_size=8): + B = sample.shape[0] + hidden_states = [] + + for i in range(0, sample.shape[2], batch_size): + + j = min(i + batch_size, sample.shape[2]) + sample_batch = rearrange(sample[:,:,i:j], "B C T H W -> (B T) C H W") + + hidden_states_batch = self(sample_batch) + hidden_states_batch = rearrange(hidden_states_batch, "(B T) C H W -> B C T H W", B=B) + + hidden_states.append(hidden_states_batch) + + hidden_states = torch.concat(hidden_states, dim=2) + return hidden_states diff --git a/diffsynth/models/flux_value_control.py b/diffsynth/models/flux_value_control.py new file mode 100644 index 0000000000000000000000000000000000000000..549dbc93b41343a42266af11584e2e7d39a17cd6 --- /dev/null +++ b/diffsynth/models/flux_value_control.py @@ -0,0 +1,56 @@ +import torch +from .general_modules import TemporalTimesteps + + +class MultiValueEncoder(torch.nn.Module): + def __init__(self, encoders=()): + super().__init__() + if not isinstance(encoders, list): + encoders = [encoders] + self.encoders = torch.nn.ModuleList(encoders) + + def __call__(self, values, dtype): + emb = [] + for encoder, value in zip(self.encoders, values): + if value is not None: + value = value.unsqueeze(0) + emb.append(encoder(value, dtype)) + emb = torch.concat(emb, dim=0) + return emb + + +class SingleValueEncoder(torch.nn.Module): + def __init__(self, dim_in=256, dim_out=4096, prefer_len=32, computation_device=None): + super().__init__() + self.prefer_len = prefer_len + self.prefer_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device) + self.prefer_value_embedder = torch.nn.Sequential( + torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out) + ) + self.positional_embedding = torch.nn.Parameter( + torch.randn(self.prefer_len, dim_out) + ) + + def forward(self, value, dtype): + value = value * 1000 + emb = self.prefer_proj(value).to(dtype) + emb = self.prefer_value_embedder(emb).squeeze(0) + base_embeddings = emb.expand(self.prefer_len, -1) + positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device) + learned_embeddings = base_embeddings + positional_embedding + return learned_embeddings + + @staticmethod + def state_dict_converter(): + return SingleValueEncoderStateDictConverter() + + +class SingleValueEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + return state_dict diff --git a/diffsynth/models/general_modules.py b/diffsynth/models/general_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..216247c6c1f1c9fbba1e6a7f9631cfeec2684743 --- /dev/null +++ b/diffsynth/models/general_modules.py @@ -0,0 +1,139 @@ +import torch, math + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, + computation_device = None, + align_dtype_to_timestep = False, +): + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device if computation_device is None else computation_device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent).to(timesteps.device) + if align_dtype_to_timestep: + emb = emb.to(timesteps.dtype) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TemporalTimesteps(torch.nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, computation_device = None, scale=1, align_dtype_to_timestep=False): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.computation_device = computation_device + self.scale = scale + self.align_dtype_to_timestep = align_dtype_to_timestep + + def forward(self, timesteps): + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + computation_device=self.computation_device, + scale=self.scale, + align_dtype_to_timestep=self.align_dtype_to_timestep, + ) + return t_emb + + +class DiffusersCompatibleTimestepProj(torch.nn.Module): + def __init__(self, dim_in, dim_out): + super().__init__() + self.linear_1 = torch.nn.Linear(dim_in, dim_out) + self.act = torch.nn.SiLU() + self.linear_2 = torch.nn.Linear(dim_out, dim_out) + + def forward(self, x): + x = self.linear_1(x) + x = self.act(x) + x = self.linear_2(x) + return x + + +class TimestepEmbeddings(torch.nn.Module): + def __init__(self, dim_in, dim_out, computation_device=None, diffusers_compatible_format=False, scale=1, align_dtype_to_timestep=False): + super().__init__() + self.time_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device, scale=scale, align_dtype_to_timestep=align_dtype_to_timestep) + if diffusers_compatible_format: + self.timestep_embedder = DiffusersCompatibleTimestepProj(dim_in, dim_out) + else: + self.timestep_embedder = torch.nn.Sequential( + torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out) + ) + + def forward(self, timestep, dtype): + time_emb = self.time_proj(timestep).to(dtype) + time_emb = self.timestep_embedder(time_emb) + return time_emb + + +class RMSNorm(torch.nn.Module): + def __init__(self, dim, eps, elementwise_affine=True): + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = torch.nn.Parameter(torch.ones((dim,))) + else: + self.weight = None + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + hidden_states = hidden_states.to(input_dtype) + if self.weight is not None: + hidden_states = hidden_states * self.weight + return hidden_states + + +class AdaLayerNorm(torch.nn.Module): + def __init__(self, dim, single=False, dual=False): + super().__init__() + self.single = single + self.dual = dual + self.linear = torch.nn.Linear(dim, dim * [[6, 2][single], 9][dual]) + self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(torch.nn.functional.silu(emb)) + if self.single: + scale, shift = emb.unsqueeze(1).chunk(2, dim=2) + x = self.norm(x) * (1 + scale) + shift + return x + elif self.dual: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = emb.unsqueeze(1).chunk(9, dim=2) + norm_x = self.norm(x) + x = norm_x * (1 + scale_msa) + shift_msa + norm_x2 = norm_x * (1 + scale_msa2) + shift_msa2 + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_x2, gate_msa2 + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.unsqueeze(1).chunk(6, dim=2) + x = self.norm(x) * (1 + scale_msa) + shift_msa + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp diff --git a/diffsynth/models/longcat_video_dit.py b/diffsynth/models/longcat_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..6d657238793ef1b42f28be86c707283bd2fdff4f --- /dev/null +++ b/diffsynth/models/longcat_video_dit.py @@ -0,0 +1,901 @@ +from typing import List, Optional, Tuple + +import math +import torch +import torch.nn as nn +import torch.amp as amp + +import numpy as np +import torch.nn.functional as F +from einops import rearrange, repeat +from .wan_video_dit import flash_attention +from ..core.gradient import gradient_checkpoint_forward + + +class RMSNorm_FP32(torch.nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +def broadcat(tensors, dim=-1): + num_tensors = len(tensors) + shape_lens = set(list(map(lambda t: len(t.shape), tensors))) + assert len(shape_lens) == 1, "tensors must all have the same number of dimensions" + shape_len = list(shape_lens)[0] + dim = (dim + shape_len) if dim < 0 else dim + dims = list(zip(*map(lambda t: list(t.shape), tensors))) + expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] + assert all( + [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)] + ), "invalid dimensions for broadcastable concatentation" + max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) + expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) + expanded_dims.insert(dim, (dim, dims[dim])) + expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) + tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) + return torch.cat(tensors, dim=dim) + + +def rotate_half(x): + x = rearrange(x, "... (d r) -> ... d r", r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, "... d r -> ... (d r)") + + +class RotaryPositionalEmbedding(nn.Module): + + def __init__(self, + head_dim, + cp_split_hw=None + ): + """Rotary positional embedding for 3D + Reference : https://blog.eleuther.ai/rotary-embeddings/ + Paper: https://arxiv.org/pdf/2104.09864.pdf + Args: + dim: Dimension of embedding + base: Base value for exponential + """ + super().__init__() + self.head_dim = head_dim + assert self.head_dim % 8 == 0, 'Dim must be a multiply of 8 for 3D RoPE.' + self.cp_split_hw = cp_split_hw + # We take the assumption that the longest side of grid will not larger than 512, i.e, 512 * 8 = 4098 input pixels + self.base = 10000 + self.freqs_dict = {} + + def register_grid_size(self, grid_size): + if grid_size not in self.freqs_dict: + self.freqs_dict.update({ + grid_size: self.precompute_freqs_cis_3d(grid_size) + }) + + def precompute_freqs_cis_3d(self, grid_size): + num_frames, height, width = grid_size + dim_t = self.head_dim - 4 * (self.head_dim // 6) + dim_h = 2 * (self.head_dim // 6) + dim_w = 2 * (self.head_dim // 6) + freqs_t = 1.0 / (self.base ** (torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t)) + freqs_h = 1.0 / (self.base ** (torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h)) + freqs_w = 1.0 / (self.base ** (torch.arange(0, dim_w, 2)[: (dim_w // 2)].float() / dim_w)) + grid_t = np.linspace(0, num_frames, num_frames, endpoint=False, dtype=np.float32) + grid_h = np.linspace(0, height, height, endpoint=False, dtype=np.float32) + grid_w = np.linspace(0, width, width, endpoint=False, dtype=np.float32) + grid_t = torch.from_numpy(grid_t).float() + grid_h = torch.from_numpy(grid_h).float() + grid_w = torch.from_numpy(grid_w).float() + freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t) + freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h) + freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w) + freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2) + freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2) + freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2) + freqs = broadcat((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1) + # (T H W D) + freqs = rearrange(freqs, "T H W D -> (T H W) D") + # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1: + # with torch.no_grad(): + # freqs = rearrange(freqs, "(T H W) D -> T H W D", T=num_frames, H=height, W=width) + # freqs = context_parallel_util.split_cp_2d(freqs, seq_dim_hw=(1, 2), split_hw=self.cp_split_hw) + # freqs = rearrange(freqs, "T H W D -> (T H W) D") + + return freqs + + def forward(self, q, k, grid_size): + """3D RoPE. + + Args: + query: [B, head, seq, head_dim] + key: [B, head, seq, head_dim] + Returns: + query and key with the same shape as input. + """ + + if grid_size not in self.freqs_dict: + self.register_grid_size(grid_size) + + freqs_cis = self.freqs_dict[grid_size].to(q.device) + q_, k_ = q.float(), k.float() + freqs_cis = freqs_cis.float().to(q.device) + cos, sin = freqs_cis.cos(), freqs_cis.sin() + cos, sin = rearrange(cos, 'n d -> 1 1 n d'), rearrange(sin, 'n d -> 1 1 n d') + q_ = (q_ * cos) + (rotate_half(q_) * sin) + k_ = (k_ * cos) + (rotate_half(k_) * sin) + + return q_.type_as(q), k_.type_as(k) + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params: dict = None, + cp_split_hw: Optional[List[int]] = None + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.enable_flashattn3 = enable_flashattn3 + self.enable_flashattn2 = enable_flashattn2 + self.enable_xformers = enable_xformers + self.enable_bsa = enable_bsa + self.bsa_params = bsa_params + self.cp_split_hw = cp_split_hw + + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.proj = nn.Linear(dim, dim) + + self.rope_3d = RotaryPositionalEmbedding( + self.head_dim, + cp_split_hw=cp_split_hw + ) + + def _process_attn(self, q, k, v, shape): + q = rearrange(q, "B H S D -> B S (H D)") + k = rearrange(k, "B H S D -> B S (H D)") + v = rearrange(v, "B H S D -> B S (H D)") + x = flash_attention(q, k, v, num_heads=self.num_heads) + x = rearrange(x, "B S (H D) -> B H S D", H=self.num_heads) + return x + + def forward(self, x: torch.Tensor, shape=None, num_cond_latents=None, return_kv=False) -> torch.Tensor: + """ + """ + B, N, C = x.shape + qkv = self.qkv(x) + + qkv_shape = (B, N, 3, self.num_heads, self.head_dim) + qkv = qkv.view(qkv_shape).permute((2, 0, 3, 1, 4)) # [3, B, H, N, D] + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + if return_kv: + k_cache, v_cache = k.clone(), v.clone() + + q, k = self.rope_3d(q, k, shape) + + # cond mode + if num_cond_latents is not None and num_cond_latents > 0: + num_cond_latents_thw = num_cond_latents * (N // shape[0]) + # process the condition tokens + q_cond = q[:, :, :num_cond_latents_thw].contiguous() + k_cond = k[:, :, :num_cond_latents_thw].contiguous() + v_cond = v[:, :, :num_cond_latents_thw].contiguous() + x_cond = self._process_attn(q_cond, k_cond, v_cond, shape) + # process the noise tokens + q_noise = q[:, :, num_cond_latents_thw:].contiguous() + x_noise = self._process_attn(q_noise, k, v, shape) + # merge x_cond and x_noise + x = torch.cat([x_cond, x_noise], dim=2).contiguous() + else: + x = self._process_attn(q, k, v, shape) + + x_output_shape = (B, N, C) + x = x.transpose(1, 2) # [B, H, N, D] --> [B, N, H, D] + x = x.reshape(x_output_shape) # [B, N, H, D] --> [B, N, C] + x = self.proj(x) + + if return_kv: + return x, (k_cache, v_cache) + else: + return x + + def forward_with_kv_cache(self, x: torch.Tensor, shape=None, num_cond_latents=None, kv_cache=None) -> torch.Tensor: + """ + """ + B, N, C = x.shape + qkv = self.qkv(x) + + qkv_shape = (B, N, 3, self.num_heads, self.head_dim) + qkv = qkv.view(qkv_shape).permute((2, 0, 3, 1, 4)) # [3, B, H, N, D] + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + T, H, W = shape + k_cache, v_cache = kv_cache + assert k_cache.shape[0] == v_cache.shape[0] and k_cache.shape[0] in [1, B] + if k_cache.shape[0] == 1: + k_cache = k_cache.repeat(B, 1, 1, 1) + v_cache = v_cache.repeat(B, 1, 1, 1) + + if num_cond_latents is not None and num_cond_latents > 0: + k_full = torch.cat([k_cache, k], dim=2).contiguous() + v_full = torch.cat([v_cache, v], dim=2).contiguous() + q_padding = torch.cat([torch.empty_like(k_cache), q], dim=2).contiguous() + q_padding, k_full = self.rope_3d(q_padding, k_full, (T + num_cond_latents, H, W)) + q = q_padding[:, :, -N:].contiguous() + + x = self._process_attn(q, k_full, v_full, shape) + + x_output_shape = (B, N, C) + x = x.transpose(1, 2) # [B, H, N, D] --> [B, N, H, D] + x = x.reshape(x_output_shape) # [B, N, H, D] --> [B, N, C] + x = self.proj(x) + + return x + + +class MultiHeadCrossAttention(nn.Module): + def __init__( + self, + dim, + num_heads, + enable_flashattn3=False, + enable_flashattn2=False, + enable_xformers=False, + ): + super(MultiHeadCrossAttention, self).__init__() + assert dim % num_heads == 0, "d_model must be divisible by num_heads" + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q_linear = nn.Linear(dim, dim) + self.kv_linear = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + + self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6) + + self.enable_flashattn3 = enable_flashattn3 + self.enable_flashattn2 = enable_flashattn2 + self.enable_xformers = enable_xformers + + def _process_cross_attn(self, x, cond, kv_seqlen): + B, N, C = x.shape + assert C == self.dim and cond.shape[2] == self.dim + + q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim) + kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim) + k, v = kv.unbind(2) + + q, k = self.q_norm(q), self.k_norm(k) + + q = rearrange(q, "B S H D -> B S (H D)") + k = rearrange(k, "B S H D -> B S (H D)") + v = rearrange(v, "B S H D -> B S (H D)") + x = flash_attention(q, k, v, num_heads=self.num_heads) + + x = x.view(B, -1, C) + x = self.proj(x) + return x + + def forward(self, x, cond, kv_seqlen, num_cond_latents=None, shape=None): + """ + x: [B, N, C] + cond: [B, M, C] + """ + if num_cond_latents is None or num_cond_latents == 0: + return self._process_cross_attn(x, cond, kv_seqlen) + else: + B, N, C = x.shape + if num_cond_latents is not None and num_cond_latents > 0: + assert shape is not None, "SHOULD pass in the shape" + num_cond_latents_thw = num_cond_latents * (N // shape[0]) + x_noise = x[:, num_cond_latents_thw:] # [B, N_noise, C] + output_noise = self._process_cross_attn(x_noise, cond, kv_seqlen) # [B, N_noise, C] + output = torch.cat([ + torch.zeros((B, num_cond_latents_thw, C), dtype=output_noise.dtype, device=output_noise.device), + output_noise + ], dim=1).contiguous() + else: + raise NotImplementedError + + return output + + +class LayerNorm_FP32(nn.LayerNorm): + def __init__(self, dim, eps, elementwise_affine): + super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + out = F.layer_norm( + inputs.float(), + self.normalized_shape, + None if self.weight is None else self.weight.float(), + None if self.bias is None else self.bias.float() , + self.eps + ).to(origin_dtype) + return out + + +def modulate_fp32(norm_func, x, shift, scale): + # Suppose x is (B, N, D), shift is (B, -1, D), scale is (B, -1, D) + # ensure the modulation params be fp32 + assert shift.dtype == torch.float32, scale.dtype == torch.float32 + dtype = x.dtype + x = norm_func(x.to(torch.float32)) + x = x * (scale + 1) + shift + x = x.to(dtype) + return x + + +class FinalLayer_FP32(nn.Module): + """ + The final layer of DiT. + """ + + def __init__(self, hidden_size, num_patch, out_channels, adaln_tembed_dim): + super().__init__() + self.hidden_size = hidden_size + self.num_patch = num_patch + self.out_channels = out_channels + self.adaln_tembed_dim = adaln_tembed_dim + + self.norm_final = LayerNorm_FP32(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, num_patch * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(adaln_tembed_dim, 2 * hidden_size, bias=True)) + + def forward(self, x, t, latent_shape): + # timestep shape: [B, T, C] + assert t.dtype == torch.float32 + B, N, C = x.shape + T, _, _ = latent_shape + + with amp.autocast('cuda', dtype=torch.float32): + shift, scale = self.adaLN_modulation(t).unsqueeze(2).chunk(2, dim=-1) # [B, T, 1, C] + x = modulate_fp32(self.norm_final, x.view(B, T, -1, C), shift, scale).view(B, N, C) + x = self.linear(x) + return x + + +class FeedForwardSwiGLU(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + multiple_of: int = 256, + ffn_dim_multiplier: Optional[float] = None, + ): + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + # custom dim factor multiplier + if ffn_dim_multiplier is not None: + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + self.dim = dim + self.hidden_dim = hidden_dim + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, t_embed_dim, frequency_embedding_size=256): + super().__init__() + self.t_embed_dim = t_embed_dim + self.frequency_embedding_size = frequency_embedding_size + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, t_embed_dim, bias=True), + nn.SiLU(), + nn.Linear(t_embed_dim, t_embed_dim, bias=True), + ) + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) + freqs = freqs.to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t, dtype): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + if t_freq.dtype != dtype: + t_freq = t_freq.to(dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +class CaptionEmbedder(nn.Module): + """ + Embeds class labels into vector representations. + """ + + def __init__(self, in_channels, hidden_size): + super().__init__() + self.in_channels = in_channels + self.hidden_size = hidden_size + self.y_proj = nn.Sequential( + nn.Linear(in_channels, hidden_size, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + def forward(self, caption): + B, _, N, C = caption.shape + caption = self.y_proj(caption) + return caption + + +class PatchEmbed3D(nn.Module): + """Video to Patch Embedding. + + Args: + patch_size (int): Patch token size. Default: (2,4,4). + in_chans (int): Number of input video channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__( + self, + patch_size=(2, 4, 4), + in_chans=3, + embed_dim=96, + norm_layer=None, + flatten=True, + ): + super().__init__() + self.patch_size = patch_size + self.flatten = flatten + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, D, H, W = x.size() + if W % self.patch_size[2] != 0: + x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2])) + if H % self.patch_size[1] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1])) + if D % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0])) + + B, C, T, H, W = x.shape + x = self.proj(x) # (B C T H W) + if self.norm is not None: + D, Wh, Ww = x.size(2), x.size(3), x.size(4) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC + return x + + +class LongCatSingleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: int, + adaln_tembed_dim: int, + enable_flashattn3: bool = False, + enable_flashattn2: bool = False, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params=None, + cp_split_hw=None + ): + super().__init__() + + self.hidden_size = hidden_size + + # scale and gate modulation + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(adaln_tembed_dim, 6 * hidden_size, bias=True) + ) + + self.mod_norm_attn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False) + self.mod_norm_ffn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False) + self.pre_crs_attn_norm = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=True) + + self.attn = Attention( + dim=hidden_size, + num_heads=num_heads, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + enable_bsa=enable_bsa, + bsa_params=bsa_params, + cp_split_hw=cp_split_hw + ) + self.cross_attn = MultiHeadCrossAttention( + dim=hidden_size, + num_heads=num_heads, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + ) + self.ffn = FeedForwardSwiGLU(dim=hidden_size, hidden_dim=int(hidden_size * mlp_ratio)) + + def forward(self, x, y, t, y_seqlen, latent_shape, num_cond_latents=None, return_kv=False, kv_cache=None, skip_crs_attn=False): + """ + x: [B, N, C] + y: [1, N_valid_tokens, C] + t: [B, T, C_t] + y_seqlen: [B]; type of a list + latent_shape: latent shape of a single item + """ + x_dtype = x.dtype + + B, N, C = x.shape + T, _, _ = latent_shape # S != T*H*W in case of CP split on H*W. + + # compute modulation params in fp32 + with amp.autocast(device_type='cuda', dtype=torch.float32): + shift_msa, scale_msa, gate_msa, \ + shift_mlp, scale_mlp, gate_mlp = \ + self.adaLN_modulation(t).unsqueeze(2).chunk(6, dim=-1) # [B, T, 1, C] + + # self attn with modulation + x_m = modulate_fp32(self.mod_norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa).view(B, N, C) + + if kv_cache is not None: + kv_cache = (kv_cache[0].to(x.device), kv_cache[1].to(x.device)) + attn_outputs = self.attn.forward_with_kv_cache(x_m, shape=latent_shape, num_cond_latents=num_cond_latents, kv_cache=kv_cache) + else: + attn_outputs = self.attn(x_m, shape=latent_shape, num_cond_latents=num_cond_latents, return_kv=return_kv) + + if return_kv: + x_s, kv_cache = attn_outputs + else: + x_s = attn_outputs + + with amp.autocast(device_type='cuda', dtype=torch.float32): + x = x + (gate_msa * x_s.view(B, -1, N//T, C)).view(B, -1, C) # [B, N, C] + x = x.to(x_dtype) + + # cross attn + if not skip_crs_attn: + if kv_cache is not None: + num_cond_latents = None + x = x + self.cross_attn(self.pre_crs_attn_norm(x), y, y_seqlen, num_cond_latents=num_cond_latents, shape=latent_shape) + + # ffn with modulation + x_m = modulate_fp32(self.mod_norm_ffn, x.view(B, -1, N//T, C), shift_mlp, scale_mlp).view(B, -1, C) + x_s = self.ffn(x_m) + with amp.autocast(device_type='cuda', dtype=torch.float32): + x = x + (gate_mlp * x_s.view(B, -1, N//T, C)).view(B, -1, C) # [B, N, C] + x = x.to(x_dtype) + + if return_kv: + return x, kv_cache + else: + return x + + +class LongCatVideoTransformer3DModel(torch.nn.Module): + def __init__( + self, + in_channels: int = 16, + out_channels: int = 16, + hidden_size: int = 4096, + depth: int = 48, + num_heads: int = 32, + caption_channels: int = 4096, + mlp_ratio: int = 4, + adaln_tembed_dim: int = 512, + frequency_embedding_size: int = 256, + # default params + patch_size: Tuple[int] = (1, 2, 2), + # attention config + enable_flashattn3: bool = False, + enable_flashattn2: bool = True, + enable_xformers: bool = False, + enable_bsa: bool = False, + bsa_params: dict = {'sparsity': 0.9375, 'chunk_3d_shape_q': [4, 4, 4], 'chunk_3d_shape_k': [4, 4, 4]}, + cp_split_hw: Optional[List[int]] = [1, 1], + text_tokens_zero_pad: bool = True, + ) -> None: + super().__init__() + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels + self.cp_split_hw = cp_split_hw + + self.x_embedder = PatchEmbed3D(patch_size, in_channels, hidden_size) + self.t_embedder = TimestepEmbedder(t_embed_dim=adaln_tembed_dim, frequency_embedding_size=frequency_embedding_size) + self.y_embedder = CaptionEmbedder( + in_channels=caption_channels, + hidden_size=hidden_size, + ) + + self.blocks = nn.ModuleList( + [ + LongCatSingleStreamBlock( + hidden_size=hidden_size, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + adaln_tembed_dim=adaln_tembed_dim, + enable_flashattn3=enable_flashattn3, + enable_flashattn2=enable_flashattn2, + enable_xformers=enable_xformers, + enable_bsa=enable_bsa, + bsa_params=bsa_params, + cp_split_hw=cp_split_hw + ) + for i in range(depth) + ] + ) + + self.final_layer = FinalLayer_FP32( + hidden_size, + np.prod(self.patch_size), + out_channels, + adaln_tembed_dim, + ) + + self.gradient_checkpointing = False + self.text_tokens_zero_pad = text_tokens_zero_pad + + self.lora_dict = {} + self.active_loras = [] + + def enable_loras(self, lora_key_list=[]): + self.disable_all_loras() + + module_loras = {} # {module_name: [lora1, lora2, ...]} + model_device = next(self.parameters()).device + model_dtype = next(self.parameters()).dtype + + for lora_key in lora_key_list: + if lora_key in self.lora_dict: + for lora in self.lora_dict[lora_key].loras: + lora.to(model_device, dtype=model_dtype, non_blocking=True) + module_name = lora.lora_name.replace("lora___lorahyphen___", "").replace("___lorahyphen___", ".") + if module_name not in module_loras: + module_loras[module_name] = [] + module_loras[module_name].append(lora) + self.active_loras.append(lora_key) + + for module_name, loras in module_loras.items(): + module = self._get_module_by_name(module_name) + if not hasattr(module, 'org_forward'): + module.org_forward = module.forward + module.forward = self._create_multi_lora_forward(module, loras) + + def _create_multi_lora_forward(self, module, loras): + def multi_lora_forward(x, *args, **kwargs): + weight_dtype = x.dtype + org_output = module.org_forward(x, *args, **kwargs) + + total_lora_output = 0 + for lora in loras: + if lora.use_lora: + lx = lora.lora_down(x.to(lora.lora_down.weight.dtype)) + lx = lora.lora_up(lx) + lora_output = lx.to(weight_dtype) * lora.multiplier * lora.alpha_scale + total_lora_output += lora_output + + return org_output + total_lora_output + + return multi_lora_forward + + def _get_module_by_name(self, module_name): + try: + module = self + for part in module_name.split('.'): + module = getattr(module, part) + return module + except AttributeError as e: + raise ValueError(f"Cannot find module: {module_name}, error: {e}") + + def disable_all_loras(self): + for name, module in self.named_modules(): + if hasattr(module, 'org_forward'): + module.forward = module.org_forward + delattr(module, 'org_forward') + + for lora_key, lora_network in self.lora_dict.items(): + for lora in lora_network.loras: + lora.to("cpu") + + self.active_loras.clear() + + def enable_bsa(self,): + for block in self.blocks: + block.attn.enable_bsa = True + + def disable_bsa(self,): + for block in self.blocks: + block.attn.enable_bsa = False + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask=None, + num_cond_latents=0, + return_kv=False, + kv_cache_dict={}, + skip_crs_attn=False, + offload_kv_cache=False, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + ): + + B, _, T, H, W = hidden_states.shape + + N_t = T // self.patch_size[0] + N_h = H // self.patch_size[1] + N_w = W // self.patch_size[2] + + assert self.patch_size[0]==1, "Currently, 3D x_embedder should not compress the temporal dimension." + + # expand the shape of timestep from [B] to [B, T] + if len(timestep.shape) == 1: + timestep = timestep.unsqueeze(1).expand(-1, N_t).clone() # [B, T] + timestep[:, :num_cond_latents] = 0 + + dtype = hidden_states.dtype + hidden_states = hidden_states.to(dtype) + timestep = timestep.to(dtype) + encoder_hidden_states = encoder_hidden_states.to(dtype) + + hidden_states = self.x_embedder(hidden_states) # [B, N, C] + + with amp.autocast(device_type='cuda', dtype=torch.float32): + t = self.t_embedder(timestep.float().flatten(), dtype=torch.float32).reshape(B, N_t, -1) # [B, T, C_t] + + encoder_hidden_states = self.y_embedder(encoder_hidden_states) # [B, 1, N_token, C] + + if self.text_tokens_zero_pad and encoder_attention_mask is not None: + encoder_hidden_states = encoder_hidden_states * encoder_attention_mask[:, None, :, None] + encoder_attention_mask = (encoder_attention_mask * 0 + 1).to(encoder_attention_mask.dtype) + + if encoder_attention_mask is not None: + encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1) + encoder_hidden_states = encoder_hidden_states.squeeze(1).masked_select(encoder_attention_mask.unsqueeze(-1) != 0).view(1, -1, hidden_states.shape[-1]) # [1, N_valid_tokens, C] + y_seqlens = encoder_attention_mask.sum(dim=1).tolist() # [B] + else: + y_seqlens = [encoder_hidden_states.shape[2]] * encoder_hidden_states.shape[0] + encoder_hidden_states = encoder_hidden_states.squeeze(1).view(1, -1, hidden_states.shape[-1]) + + # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1: + # hidden_states = rearrange(hidden_states, "B (T H W) C -> B T H W C", T=N_t, H=N_h, W=N_w) + # hidden_states = context_parallel_util.split_cp_2d(hidden_states, seq_dim_hw=(2, 3), split_hw=self.cp_split_hw) + # hidden_states = rearrange(hidden_states, "B T H W C -> B (T H W) C") + + # blocks + kv_cache_dict_ret = {} + for i, block in enumerate(self.blocks): + block_outputs = gradient_checkpoint_forward( + block, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + x=hidden_states, + y=encoder_hidden_states, + t=t, + y_seqlen=y_seqlens, + latent_shape=(N_t, N_h, N_w), + num_cond_latents=num_cond_latents, + return_kv=return_kv, + kv_cache=kv_cache_dict.get(i, None), + skip_crs_attn=skip_crs_attn, + ) + + if return_kv: + hidden_states, kv_cache = block_outputs + if offload_kv_cache: + kv_cache_dict_ret[i] = (kv_cache[0].cpu(), kv_cache[1].cpu()) + else: + kv_cache_dict_ret[i] = (kv_cache[0].contiguous(), kv_cache[1].contiguous()) + else: + hidden_states = block_outputs + + hidden_states = self.final_layer(hidden_states, t, (N_t, N_h, N_w)) # [B, N, C=T_p*H_p*W_p*C_out] + + # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1: + # hidden_states = context_parallel_util.gather_cp_2d(hidden_states, shape=(N_t, N_h, N_w), split_hw=self.cp_split_hw) + + hidden_states = self.unpatchify(hidden_states, N_t, N_h, N_w) # [B, C_out, H, W] + + # cast to float32 for better accuracy + hidden_states = hidden_states.to(torch.float32) + + if return_kv: + return hidden_states, kv_cache_dict_ret + else: + return hidden_states + + + def unpatchify(self, x, N_t, N_h, N_w): + """ + Args: + x (torch.Tensor): of shape [B, N, C] + + Return: + x (torch.Tensor): of shape [B, C_out, T, H, W] + """ + T_p, H_p, W_p = self.patch_size + x = rearrange( + x, + "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)", + N_t=N_t, + N_h=N_h, + N_w=N_w, + T_p=T_p, + H_p=H_p, + W_p=W_p, + C_out=self.out_channels, + ) + return x + + @staticmethod + def state_dict_converter(): + return LongCatVideoTransformer3DModelDictConverter() + + +class LongCatVideoTransformer3DModelDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + return state_dict + diff --git a/diffsynth/models/model_loader.py b/diffsynth/models/model_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..16d72ddba8abb44dda24a03122d9fff79dd42660 --- /dev/null +++ b/diffsynth/models/model_loader.py @@ -0,0 +1,111 @@ +from ..core.loader import load_model, hash_model_file +from ..core.vram import AutoWrappedModule +from ..configs import MODEL_CONFIGS, VRAM_MANAGEMENT_MODULE_MAPS +import importlib, json, torch + + +class ModelPool: + def __init__(self): + self.model = [] + self.model_name = [] + self.model_path = [] + + def import_model_class(self, model_class): + split = model_class.rfind(".") + model_resource, model_class = model_class[:split], model_class[split+1:] + model_class = importlib.import_module(model_resource).__getattribute__(model_class) + return model_class + + def need_to_enable_vram_management(self, vram_config): + return vram_config["offload_dtype"] is not None and vram_config["offload_device"] is not None + + def fetch_module_map(self, model_class, vram_config): + if self.need_to_enable_vram_management(vram_config): + if model_class in VRAM_MANAGEMENT_MODULE_MAPS: + module_map = {self.import_model_class(source): self.import_model_class(target) for source, target in VRAM_MANAGEMENT_MODULE_MAPS[model_class].items()} + else: + module_map = {self.import_model_class(model_class): AutoWrappedModule} + else: + module_map = None + return module_map + + def load_model_file(self, config, path, vram_config, vram_limit=None): + model_class = self.import_model_class(config["model_class"]) + model_config = config.get("extra_kwargs", {}) + if "state_dict_converter" in config: + state_dict_converter = self.import_model_class(config["state_dict_converter"]) + else: + state_dict_converter = None + module_map = self.fetch_module_map(config["model_class"], vram_config) + model = load_model( + model_class, path, model_config, + vram_config["computation_dtype"], vram_config["computation_device"], + state_dict_converter, + use_disk_map=True, + vram_config=vram_config, module_map=module_map, vram_limit=vram_limit, + ) + return model + + def default_vram_config(self): + vram_config = { + "offload_dtype": None, + "offload_device": None, + "onload_dtype": torch.bfloat16, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cpu", + "computation_dtype": torch.bfloat16, + "computation_device": "cpu", + } + return vram_config + + def auto_load_model(self, path, vram_config=None, vram_limit=None, clear_parameters=False): + print(f"Loading models from: {json.dumps(path, indent=4)}") + if vram_config is None: + vram_config = self.default_vram_config() + model_hash = hash_model_file(path) + loaded = False + for config in MODEL_CONFIGS: + if config["model_hash"] == model_hash: + model = self.load_model_file(config, path, vram_config, vram_limit=vram_limit) + if clear_parameters: self.clear_parameters(model) + self.model.append(model) + model_name = config["model_name"] + self.model_name.append(model_name) + self.model_path.append(path) + model_info = {"model_name": model_name, "model_class": config["model_class"], "extra_kwargs": config.get("extra_kwargs")} + print(f"Loaded model: {json.dumps(model_info, indent=4)}") + loaded = True + if not loaded: + raise ValueError(f"Cannot detect the model type. File: {path}. Model hash: {model_hash}") + + def fetch_model(self, model_name, index=None): + fetched_models = [] + fetched_model_paths = [] + for model, model_path, model_name_ in zip(self.model, self.model_path, self.model_name): + if model_name == model_name_: + fetched_models.append(model) + fetched_model_paths.append(model_path) + if len(fetched_models) == 0: + print(f"No {model_name} models available. This is not an error.") + model = None + elif len(fetched_models) == 1: + print(f"Using {model_name} from {json.dumps(fetched_model_paths[0], indent=4)}.") + model = fetched_models[0] + else: + if index is None: + model = fetched_models[0] + print(f"More than one {model_name} models are loaded: {fetched_model_paths}. Using {model_name} from {json.dumps(fetched_model_paths[0], indent=4)}.") + elif isinstance(index, int): + model = fetched_models[:index] + print(f"More than one {model_name} models are loaded: {fetched_model_paths}. Using {model_name} from {json.dumps(fetched_model_paths[:index], indent=4)}.") + else: + model = fetched_models + print(f"More than one {model_name} models are loaded: {fetched_model_paths}. Using {model_name} from {json.dumps(fetched_model_paths, indent=4)}.") + return model + + def clear_parameters(self, model: torch.nn.Module): + for name, module in model.named_children(): + self.clear_parameters(module) + for name, param in model.named_parameters(recurse=False): + setattr(model, name, None) diff --git a/diffsynth/models/nexus_gen.py b/diffsynth/models/nexus_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..011039842312b01a0bbb69b999bc868902736e9a --- /dev/null +++ b/diffsynth/models/nexus_gen.py @@ -0,0 +1,161 @@ +import torch +from PIL import Image + + +class NexusGenAutoregressiveModel(torch.nn.Module): + def __init__(self, max_length=1024, max_pixels=262640): + super(NexusGenAutoregressiveModel, self).__init__() + from .nexus_gen_ar_model import Qwen2_5_VLForConditionalGeneration + from transformers import Qwen2_5_VLConfig + self.max_length = max_length + self.max_pixels = max_pixels + model_config = Qwen2_5_VLConfig(**{ + "_name_or_path": "DiffSynth-Studio/Nexus-GenV2", + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_qwen2_5_vl.Qwen2_5_VLConfig", + "AutoModel": "modeling_qwen2_5_vl.Qwen2_5_VLModel", + "AutoModelForCausalLM": "modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration" + }, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "image_token_id": 151655, + "initializer_range": 0.02, + "intermediate_size": 18944, + "max_position_embeddings": 128000, + "max_window_layers": 28, + "model_type": "qwen2_5_vl", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "pad_token_id": 151643, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_section": [ + 16, + 24, + 24 + ], + "rope_type": "default", + "type": "default" + }, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.49.0", + "use_cache": False, + "use_sliding_window": False, + "video_token_id": 151656, + "vision_config": { + "hidden_size": 1280, + "in_chans": 3, + "model_type": "qwen2_5_vl", + "spatial_patch_size": 14, + "tokens_per_second": 2, + "torch_dtype": "bfloat16" + }, + "vision_end_token_id": 151653, + "vision_start_token_id": 151652, + "vision_token_id": 151654, + "vocab_size": 152064 + }) + self.model = Qwen2_5_VLForConditionalGeneration(model_config) + self.processor = None + + + def load_processor(self, path): + from .nexus_gen_ar_model import Qwen2_5_VLProcessor + self.processor = Qwen2_5_VLProcessor.from_pretrained(path) + + + @staticmethod + def state_dict_converter(): + return NexusGenAutoregressiveModelStateDictConverter() + + def bound_image(self, image, max_pixels=262640): + from qwen_vl_utils import smart_resize + resized_height, resized_width = smart_resize( + image.height, + image.width, + max_pixels=max_pixels, + ) + return image.resize((resized_width, resized_height)) + + def get_editing_msg(self, instruction): + if '' not in instruction: + instruction = ' ' + instruction + messages = [{"role":"user", "content":instruction}, {"role":"assistant", "content":"Here is the image: "}] + return messages + + def get_generation_msg(self, instruction): + instruction = "Generate an image according to the following description: {}".format(instruction) + messages = [{"role":"user", "content":instruction}, {"role":"assistant", "content":"Here is an image based on the description: "}] + return messages + + def forward(self, instruction, ref_image=None, num_img_tokens=81): + """ + Generate target embeddings for the given instruction and reference image. + """ + if ref_image is not None: + messages = self.get_editing_msg(instruction) + images = [self.bound_image(ref_image)] + [Image.new(mode='RGB', size=(252, 252), color=(255, 255, 255))] + output_image_embeddings = self.get_target_embeddings(images, messages, self.processor, self.model, num_img_tokens) + else: + messages = self.get_generation_msg(instruction) + images = [Image.new(mode='RGB', size=(252, 252), color=(255, 255, 255))] + output_image_embeddings = self.get_target_embeddings(images, messages, self.processor, self.model, num_img_tokens) + + return output_image_embeddings + + def get_target_embeddings(self, images, messages, processor, model, num_img_tokens=81): + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + text = text.replace('', '<|vision_start|><|image_pad|><|vision_end|>') + inputs = processor( + text=[text], + images=images, + padding=True, + return_tensors="pt", + ) + inputs = inputs.to(model.device) + + input_embeds = model.model.embed_tokens(inputs['input_ids']) + image_embeds = model.visual(inputs['pixel_values'], grid_thw=inputs['image_grid_thw']) + ground_truth_image_embeds = image_embeds[-num_img_tokens:] + input_image_embeds = image_embeds[:-num_img_tokens] + + image_mask = inputs['input_ids'] == model.config.image_token_id + indices = image_mask.cumsum(dim=1) + input_image_mask = torch.logical_and(indices <= (image_embeds.shape[0] - ground_truth_image_embeds.shape[0]), image_mask) + gt_image_mask = torch.logical_and(image_mask, ~input_image_mask) + input_image_mask = input_image_mask.unsqueeze(-1).expand_as(input_embeds) + input_embeds = input_embeds.masked_scatter(input_image_mask, input_image_embeds) + + image_prefill_embeds = model.image_prefill_embeds( + torch.arange(81, device=model.device).long() + ) + input_embeds = input_embeds.masked_scatter(gt_image_mask.unsqueeze(-1).expand_as(input_embeds), image_prefill_embeds) + + position_ids, _ = model.get_rope_index( + inputs['input_ids'], + inputs['image_grid_thw'], + attention_mask=inputs['attention_mask']) + position_ids = position_ids.contiguous() + outputs = model(inputs_embeds=input_embeds, position_ids=position_ids, attention_mask=inputs['attention_mask'], return_dict=True) + output_image_embeddings = outputs.image_embeddings[:, :-1, :] + output_image_embeddings = output_image_embeddings[gt_image_mask[:, 1:]] + return output_image_embeddings, input_image_embeds, inputs['image_grid_thw'] + + +class NexusGenAutoregressiveModelStateDictConverter: + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict = {"model." + key: value for key, value in state_dict.items()} + return state_dict diff --git a/diffsynth/models/nexus_gen_ar_model.py b/diffsynth/models/nexus_gen_ar_model.py new file mode 100644 index 0000000000000000000000000000000000000000..d5a29735e1ca71fcbf7847928a309e7799718ec7 --- /dev/null +++ b/diffsynth/models/nexus_gen_ar_model.py @@ -0,0 +1,1143 @@ +import os +import re +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch.nn import CrossEntropyLoss + +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin, LogitsProcessorList, StoppingCriteriaList, GenerationConfig, GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput +from transformers.utils import add_start_docstrings_to_model_forward, logging, replace_return_docstrings +from transformers.modeling_outputs import ModelOutput +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VisionTransformerPretrainedModel, + Qwen2_5_VLModel, + Qwen2_5_VLPreTrainedModel, + QWEN2_5_VL_INPUTS_DOCSTRING, + ) + +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_utils import ImageInput, VideoInput +from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput + +GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput] + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "Qwen2_5_VLConfig" + + +@dataclass +class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): + """ + Base class for Qwen2_5_VL causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + image_embeddings: torch.FloatTensor = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + + +class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.model = Qwen2_5_VLModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vision_head = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.rope_deltas = None # cache rope_deltas here + self.image_prefill_embeds = nn.Embedding(81, config.hidden_size) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. + + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. + + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] + + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embedding for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + @add_start_docstrings_to_model_forward(QWEN2_5_VL_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=Qwen2_5_VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + image_embeddings: Optional[torch.Tensor] = None, + token_loss_weight: Optional[float] = 0.1, + img_loss_weight: Optional[float] = 1.0, + ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": "What is shown in this image?"}, + ], + }, + ] + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos]) + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # test feature + inputs_embeds = self.model.embed_tokens(input_ids) + # for image encoding and training + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + # position_ids [3, B, L] + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + image_embeds = self.vision_head(hidden_states) + + loss = None + if labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + # prepare labels for logits + logits_labels = labels.clone().detach() + image_tokens = (labels == self.config.image_token_id) + logits_labels[image_tokens] = -100 + + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = logits_labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) * token_loss_weight + + shift_image_tokens_2d = (labels[..., 1:].contiguous() == self.config.image_token_id) # (B, L-1) + shifted_image_embeds = image_embeds[:, :-1, :].contiguous() # (B, L-1, D) + masked_image_embeds = shifted_image_embeds[shift_image_tokens_2d] # (num_image_tokens, D) + + mse_loss_fct = nn.MSELoss() + mse_loss_fct = mse_loss_fct.to(shift_logits.device) + if image_embeddings is None: + image_embeddings = torch.zeros_like(masked_image_embeds) + img_loss = mse_loss_fct(masked_image_embeds, image_embeddings) + + cos_sim = torch.cosine_similarity( + masked_image_embeds, + image_embeddings, + dim=-1 + ) + cos_loss = (1 - cos_sim).mean() + img_loss = 0.5 * img_loss + 0.5 * cos_loss + # fix nan for empty image tokens + if image_embeddings.size(0) == 0: + img_loss = img_loss.nan_to_num(0.0) + # combine the loss + loss = loss + img_loss_weight * img_loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2_5_VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + image_embeddings=image_embeds, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + + + def _sample( + self, + input_ids: torch.LongTensor, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + streamer: Optional["BaseStreamer"], + **model_kwargs, + ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and + can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`: + A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + pad_token_id = generation_config._pad_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + max_length = generation_config.max_length + has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) + do_sample = generation_config.do_sample + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # keep track of which sequences are already finished + batch_size, cur_len = input_ids.shape + this_peer_finished = False + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + model_forward = self.__call__ + if isinstance(model_kwargs.get("past_key_values"), Cache): + is_compileable = model_kwargs["past_key_values"].is_compileable and self._supports_static_cache + is_compileable = is_compileable and not self.generation_config.disable_compile + if is_compileable and ( + self.device.type == "cuda" or generation_config.compile_config._compile_all_devices + ): + os.environ["TOKENIZERS_PARALLELISM"] = "0" + model_forward = self.get_compiled_call(generation_config.compile_config) + + is_prefill = True + is_sampling_img = input_ids[:, -1] == self.config.vision_start_token_id + generation_image_grid_thw = model_kwargs.pop("generation_image_grid_thw", self.get_default_image_grid_thw()) + num_img_tokens = self.get_num_image_tokens(generation_image_grid_thw) + output_image_embeddings = [] + while self._has_unfinished_sequences( + this_peer_finished, synced_gpus, device=input_ids.device, cur_len=cur_len, max_length=max_length + ): + # prepare model inputs + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # prepare prefilled embeds + model_inputs.update(self.prepare_prefilled_image_embeds(len(output_image_embeddings), num_img_tokens, is_sampling_img, **model_kwargs)) + + # parse position_ids from model_kwargs + model_inputs.update(self.prepare_image_position_ids(input_ids, generation_image_grid_thw, is_sampling_img, **model_kwargs)) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + if is_prefill: + outputs = self(**model_inputs, return_dict=True) + is_prefill = False + else: + outputs = model_forward(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + # TODO: support batch image sampling + if bool(is_sampling_img) and len(output_image_embeddings) < num_img_tokens: + output_image_embeddings.append(outputs.image_embeddings[:, -1, :].unsqueeze(1)) + + if synced_gpus and this_peer_finished: + continue + # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration + # (the clone itself is always small) + next_token_logits = outputs.logits[:, -1, :].clone().float() + next_token_logits = next_token_logits.to(input_ids.device) + + # do not sample token + next_token_logits[:, self.config.vision_end_token_id] = -float('inf') + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (next_token_scores,) + if output_logits: + raw_logits += (next_token_logits,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + # token selection + if do_sample: + probs = nn.functional.softmax(next_token_scores, dim=-1) + # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + # while not bool(is_sampling_img) and torch.any(next_tokens == self.config.vision_end_token_id): + # probs[:, self.config.vision_end_token_id] = 0 + # next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(next_token_scores, dim=-1) + + # finished sentences should have their next token be a padding token + if has_eos_stopping_criteria: + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + #TODO: support batch image sample + if num_img_tokens is not None: + cur_img_tokens = (input_ids == self.config.vision_start_token_id).flip(dims=[1]).float().argmax(dim=1) + # check whether is sampling images + is_end_img = torch.logical_and(cur_img_tokens == num_img_tokens, is_sampling_img) + is_sampling_img = torch.logical_and(is_sampling_img, cur_img_tokens < num_img_tokens) + next_tokens[is_sampling_img] = self.config.image_token_id + # check whether to end sampling images + next_tokens[is_end_img] = self.config.vision_end_token_id + else: + # check whether to end sampling images + is_sampling_img = torch.logical_and(is_sampling_img, (next_tokens != self.config.vision_end_token_id)) + # replace the next token with the image token if is sampling image + next_tokens[is_sampling_img] = self.config.image_token_id + # check whether to start sampling images + is_sampling_img = torch.logical_or(is_sampling_img, (next_tokens == self.config.vision_start_token_id)) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + + if streamer is not None: + streamer.put(next_tokens.cpu()) + + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + cur_len += 1 + + # This is needed to properly delete outputs.logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + del outputs + + if streamer is not None: + streamer.end() + + # output the image embeddings + output_image_embeddings = torch.cat(output_image_embeddings, dim=1) if len(output_image_embeddings) > 0 else None + + if return_dict_in_generate: + return GenerateDecoderOnlyAll2AllOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + output_image_embeddings=output_image_embeddings, + ) + else: + return input_ids + + + def prepare_prefilled_image_embeds(self, cur_image_tokens, num_img_tokens, is_sampling_img, **model_kwargs): + if cur_image_tokens == 0 or cur_image_tokens > num_img_tokens or not bool(is_sampling_img): + return {} + # TODO: support batch image sample + image_idx = torch.tensor([cur_image_tokens-1]).to(self.device).long().unsqueeze(0) + inputs_embeds = self.image_prefill_embeds(image_idx) + return {"inputs_embeds": inputs_embeds} + + + def get_default_image_grid_thw(self,): + return torch.tensor([[1, 18, 18]]).to(self.device) + + + def get_num_image_tokens(self, image_grid_thw): + return int(torch.prod(image_grid_thw, dim=1).sum() // 4) + + + def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]): + num_img_tokens = model_kwargs.pop("generation_image_grid_thw", None) + super()._validate_model_kwargs(model_kwargs) + model_kwargs["generation_image_grid_thw"] = num_img_tokens + + def prepare_image_position_ids(self, input_ids, generation_image_grid_thw, is_sampling_img, **model_kwargs): + # Overwritten -- prepare position_ids for image tokens + cur_img_tokens = int((input_ids == self.config.vision_start_token_id).flip(dims=[1]).float().argmax(dim=1)) + # TODO: support batch image sample + if cur_img_tokens > 0 and bool(is_sampling_img): + image_grid_thw = generation_image_grid_thw + if model_kwargs.get('image_grid_thw') is not None: + image_grid_thw = torch.cat([model_kwargs.get('image_grid_thw'), image_grid_thw]) + remaining_img_tokens = self.get_num_image_tokens(generation_image_grid_thw) - cur_img_tokens + padding_ids = input_ids.new_full((1, remaining_img_tokens), fill_value=self.config.image_token_id) + padded_ids = torch.cat([input_ids, padding_ids], dim=1) + position_ids, _ = self.get_rope_index(padded_ids, image_grid_thw, None, None) + if model_kwargs.get("use_cache", True): + position_ids = position_ids[:, :, input_ids.shape[1] - 1].unsqueeze(-1) + else: + position_ids = position_ids[:, :, :input_ids.shape[1]] + return {"position_ids": position_ids} + return {} + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + image_embeddings=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + position_ids=position_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + use_cache=use_cache, + **kwargs, + ) + + # Qwen2-5-VL position_ids are prepared with rope_deltas in forward + model_inputs["position_ids"] = None + + if cache_position[0] != 0: + model_inputs["pixel_values"] = None + model_inputs["pixel_values_videos"] = None + return model_inputs + + def _get_image_nums_and_video_nums( + self, + input_ids: Optional[torch.LongTensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Returns: + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) + """ + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + + vision_start_mask = input_ids == vision_start_token_id + vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) + image_mask = input_ids == image_token_id + video_mask = input_ids == video_token_id + image_nums = torch.sum(vision_first_mask & image_mask, dim=1) + video_nums = torch.sum(vision_first_mask & video_mask, dim=1) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + + def _expand_dict_for_generation_visual(dict_to_expand): + image_grid_thw = model_kwargs.get("image_grid_thw", None) + video_grid_thw = model_kwargs.get("video_grid_thw", None) + image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids) + + def _repeat_interleave_samples(x, lengths, repeat_times): + samples = torch.split(x, lengths) + repeat_args = [repeat_times] + [1] * (x.dim() - 1) + result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # split images into samples + samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "image_grid_thw": + # get the num of images for each sample + lengths = list(image_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "pixel_values_videos": + samples = torch.split(video_grid_thw, list(video_nums)) + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "video_grid_thw": + lengths = list(video_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "second_per_grid_ts": + if not isinstance(dict_to_expand[key], list): + raise TypeError( + f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead." + ) + tensor = torch.tensor(dict_to_expand[key]) + lengths = list(video_nums) + tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size) + dict_to_expand[key] = tensor.tolist() + return dict_to_expand + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if ( + key != "cache_position" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + return dict_to_expand + + # input_ids is required for expanding visual inputs + # If input_ids is unavailable, visual inputs will not be used; therefore, there is no need to expand visual inputs. + if input_ids is not None and input_ids.numel() != 0: + model_kwargs = _expand_dict_for_generation_visual(model_kwargs) + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs + + +__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel"] + + + +class Qwen2_5_VLVideosProcessorKwargs(VideosKwargs, total=False): + fps: Union[List[float], float] + + +class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False): + videos_kwargs: Qwen2_5_VLVideosProcessorKwargs + _defaults = { + "text_kwargs": { + "padding": False, + }, + "videos_kwargs": {"fps": 2.0}, + } + + +class Qwen2_5_VLProcessor(ProcessorMixin): + r""" + Constructs a Qwen2.5-VL processor which wraps a Qwen2.5-VL image processor and a Qwen2 tokenizer into a single processor. + [`Qwen2_5_VLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the + [`~Qwen2_5_VLProcessor.__call__`] and [`~Qwen2_5_VLProcessor.decode`] for more information. + Args: + image_processor ([`Qwen2VLImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`Qwen2TokenizerFast`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages + in a chat into a tokenizable string. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + + image_processor_class = "AutoImageProcessor" + tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast") + + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): + self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token + self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + def __call__( + self, + images: ImageInput = None, + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, + videos: VideoInput = None, + **kwargs: Unpack[Qwen2_5_VLProcessorKwargs], + ) -> BatchFeature: + """ + Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` + and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode + the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to + Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. Both channels-first and channels-last formats are supported. + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`): + The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch + tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`. + - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. + - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. + - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. + """ + output_kwargs = self._merge_kwargs( + Qwen2_5_VLProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + if images is not None: + image_inputs = self.image_processor(images=images, videos=None, **output_kwargs["images_kwargs"]) + image_grid_thw = image_inputs["image_grid_thw"] + else: + image_inputs = {} + image_grid_thw = None + + if videos is not None: + videos_inputs = self.image_processor(images=None, videos=videos, **output_kwargs["images_kwargs"]) + video_grid_thw = videos_inputs["video_grid_thw"] + + fps = output_kwargs["videos_kwargs"].pop("fps", 2.0) + if isinstance(fps, (int, float)): + second_per_grid_ts = [self.image_processor.temporal_patch_size / fps] * len(video_grid_thw) + elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw): + second_per_grid_ts = [self.image_processor.temporal_patch_size / tmp for tmp in fps] + else: + raise ValueError( + f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number." + ) + videos_inputs.update({"second_per_grid_ts": second_per_grid_ts}) + + else: + videos_inputs = {} + video_grid_thw = None + + if not isinstance(text, list): + text = [text] + + if image_grid_thw is not None: + merge_length = self.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + text[i] = text[i].replace( + self.image_token, + "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length), + 1, + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + + if video_grid_thw is not None: + merge_length = self.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.video_token in text[i]: + text[i] = text[i].replace( + self.video_token, + "<|placeholder|>" * (video_grid_thw[index].prod() // merge_length), + 1, + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.video_token) + + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def batch_decode_all2all(self, *args, **kwargs): + """ + This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + decoded = self.tokenizer.batch_decode(*args, **kwargs) + pattern = r'<\|vision_start\|>.*?<\|vision_end\|>' + decoded_with_image_tag = [re.sub(pattern, '', d, flags=re.DOTALL) for d in decoded] + decoded_with_image_tag = [re.sub(r'<\|im_end\|>', '', d) for d in decoded_with_image_tag] + return decoded_with_image_tag + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + def post_process_image_text_to_text( + self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs + ): + """ + Post-process the output of the model to decode the text. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + or `(sequence_length,)`. + skip_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + Clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `List[str]`: The decoded text. + """ + return self.tokenizer.batch_decode( + generated_outputs, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + return names_from_processor + ["second_per_grid_ts"] + + +__all__ = ["Qwen2_5_VLProcessor"] diff --git a/diffsynth/models/nexus_gen_projector.py b/diffsynth/models/nexus_gen_projector.py new file mode 100644 index 0000000000000000000000000000000000000000..d69b3e1bfd50fc3b9c098f7775afae4020f9b320 --- /dev/null +++ b/diffsynth/models/nexus_gen_projector.py @@ -0,0 +1,417 @@ +import math +import torch +import torch.nn as nn +from typing import Optional, Tuple + + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + mrope_section = mrope_section * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + from transformers.modeling_rope_utils import _compute_default_rope_parameters + self.rope_init_fn = _compute_default_rope_parameters + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len, **self.rope_kwargs + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block. In contrast to other models, Qwen2_5_VL has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class Qwen2_5_VLAttention(nn.Module): + def __init__(self, config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.is_causal = True + self.attention_dropout = config.attention_dropout + self.rope_scaling = config.rope_scaling + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + # Fix precision issues in Qwen2-VL float16 inference + # Replace inf values with zeros in attention weights to prevent NaN propagation + if query_states.dtype == torch.float16: + attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights) + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output + + +class Qwen2MLP(nn.Module): + def __init__(self, config): + super().__init__() + from transformers.activations import ACT2FN + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen2_5_VLDecoderLayer(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Qwen2_5_VLAttention(config, layer_idx) + + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class NexusGenImageEmbeddingMerger(nn.Module): + def __init__(self, num_layers=1, out_channel=4096, expand_ratio=4, device='cpu'): + super().__init__() + from transformers import Qwen2_5_VLConfig + from transformers.activations import ACT2FN + config = Qwen2_5_VLConfig(**{ + "_name_or_path": "DiffSynth-Studio/Nexus-GenV2", + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_qwen2_5_vl.Qwen2_5_VLConfig", + "AutoModel": "modeling_qwen2_5_vl.Qwen2_5_VLModel", + "AutoModelForCausalLM": "modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration" + }, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "image_token_id": 151655, + "initializer_range": 0.02, + "intermediate_size": 18944, + "max_position_embeddings": 128000, + "max_window_layers": 28, + "model_type": "qwen2_5_vl", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "pad_token_id": 151643, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_section": [ + 16, + 24, + 24 + ], + "rope_type": "default", + "type": "default" + }, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.49.0", + "use_cache": False, + "use_sliding_window": False, + "video_token_id": 151656, + "vision_config": { + "hidden_size": 1280, + "in_chans": 3, + "model_type": "qwen2_5_vl", + "spatial_patch_size": 14, + "tokens_per_second": 2, + "torch_dtype": "bfloat16" + }, + "vision_end_token_id": 151653, + "vision_start_token_id": 151652, + "vision_token_id": 151654, + "vocab_size": 152064 + }) + self.config = config + self.num_layers = num_layers + self.layers = nn.ModuleList([Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(num_layers)]) + self.projector = nn.Sequential(Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps), + nn.Linear(config.hidden_size, out_channel * expand_ratio), + Qwen2RMSNorm(out_channel * expand_ratio, eps=config.rms_norm_eps), + ACT2FN[config.hidden_act], nn.Linear(out_channel * expand_ratio, out_channel), + Qwen2RMSNorm(out_channel, eps=config.rms_norm_eps)) + self.base_grid = torch.tensor([[1, 72, 72]], device=device) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config, device=device) + + def get_position_ids(self, image_grid_thw): + """ + Generates position ids for the input embeddings grid. + modified from the qwen2_vl mrope. + """ + batch_size = image_grid_thw.shape[0] + spatial_merge_size = self.config.vision_config.spatial_merge_size + t, h, w = ( + image_grid_thw[0][0], + image_grid_thw[0][1], + image_grid_thw[0][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + scale_h = self.base_grid[0][1].item() / h.item() + scale_w = self.base_grid[0][2].item() / w.item() + + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + time_tensor = expanded_range * self.config.vision_config.tokens_per_second + t_index = time_tensor.long().flatten().to(image_grid_thw.device) + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten().to(image_grid_thw.device) * scale_h + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten().to(image_grid_thw.device) * scale_w + # 3, B, L + position_ids = torch.stack([t_index, h_index, w_index]).unsqueeze(0).repeat(batch_size, 1, 1).permute(1, 0, 2) + return position_ids + + def forward(self, embeds, embeds_grid, ref_embeds=None, ref_embeds_grid=None): + position_ids = self.get_position_ids(embeds_grid) + hidden_states = embeds + if ref_embeds is not None: + position_ids_ref_embeds = self.get_position_ids(ref_embeds_grid) + position_ids = torch.cat((position_ids, position_ids_ref_embeds), dim=-1) + hidden_states = torch.cat((embeds, ref_embeds), dim=1) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings) + + hidden_states = self.projector(hidden_states) + return hidden_states + + @staticmethod + def state_dict_converter(): + return NexusGenMergerStateDictConverter() + + +class NexusGenMergerStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + merger_state_dict = {key.replace("embedding_merger.", ""): value for key, value in state_dict.items() if key.startswith('embedding_merger.')} + return merger_state_dict + + +class NexusGenAdapter(nn.Module): + """ + Adapter for Nexus-Gen generation decoder. + """ + def __init__(self, input_dim=3584, output_dim=4096): + super(NexusGenAdapter, self).__init__() + self.adapter = nn.Sequential(nn.Linear(input_dim, output_dim), + nn.LayerNorm(output_dim), nn.ReLU(), + nn.Linear(output_dim, output_dim), + nn.LayerNorm(output_dim)) + + def forward(self, x): + return self.adapter(x) + + @staticmethod + def state_dict_converter(): + return NexusGenAdapterStateDictConverter() + + +class NexusGenAdapterStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + adapter_state_dict = {key: value for key, value in state_dict.items() if key.startswith('adapter.')} + return adapter_state_dict diff --git a/diffsynth/models/qwen_image_controlnet.py b/diffsynth/models/qwen_image_controlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..6ce40809065b3eac020d2b1da29101681a44764a --- /dev/null +++ b/diffsynth/models/qwen_image_controlnet.py @@ -0,0 +1,56 @@ +import torch +import torch.nn as nn +from .general_modules import RMSNorm + + +class BlockWiseControlBlock(torch.nn.Module): + # [linear, gelu, linear] + def __init__(self, dim: int = 3072): + super().__init__() + self.x_rms = RMSNorm(dim, eps=1e-6) + self.y_rms = RMSNorm(dim, eps=1e-6) + self.input_proj = nn.Linear(dim, dim) + self.act = nn.GELU() + self.output_proj = nn.Linear(dim, dim) + + def forward(self, x, y): + x, y = self.x_rms(x), self.y_rms(y) + x = self.input_proj(x + y) + x = self.act(x) + x = self.output_proj(x) + return x + + def init_weights(self): + # zero initialize output_proj + nn.init.zeros_(self.output_proj.weight) + nn.init.zeros_(self.output_proj.bias) + + +class QwenImageBlockWiseControlNet(torch.nn.Module): + def __init__( + self, + num_layers: int = 60, + in_dim: int = 64, + additional_in_dim: int = 0, + dim: int = 3072, + ): + super().__init__() + self.img_in = nn.Linear(in_dim + additional_in_dim, dim) + self.controlnet_blocks = nn.ModuleList( + [ + BlockWiseControlBlock(dim) + for _ in range(num_layers) + ] + ) + + def init_weight(self): + nn.init.zeros_(self.img_in.weight) + nn.init.zeros_(self.img_in.bias) + for block in self.controlnet_blocks: + block.init_weights() + + def process_controlnet_conditioning(self, controlnet_conditioning): + return self.img_in(controlnet_conditioning) + + def blockwise_forward(self, img, controlnet_conditioning, block_id): + return self.controlnet_blocks[block_id](img, controlnet_conditioning) diff --git a/diffsynth/models/qwen_image_dit.py b/diffsynth/models/qwen_image_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..ed3d9adb0dc337a92483d80865dffe9103f08538 --- /dev/null +++ b/diffsynth/models/qwen_image_dit.py @@ -0,0 +1,536 @@ +import torch, math +import torch.nn as nn +from typing import Tuple, Optional, Union, List +from einops import rearrange +from .general_modules import TimestepEmbeddings, RMSNorm, AdaLayerNorm + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + + +def qwen_image_flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, attention_mask = None, enable_fp8_attention: bool = False): + if FLASH_ATTN_3_AVAILABLE and attention_mask is None: + if not enable_fp8_attention: + q = rearrange(q, "b n s d -> b s n d", n=num_heads) + k = rearrange(k, "b n s d -> b s n d", n=num_heads) + v = rearrange(v, "b n s d -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x, tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + else: + origin_dtype = q.dtype + q_std, k_std, v_std = q.std(), k.std(), v.std() + q, k, v = (q / q_std).to(torch.float8_e4m3fn), (k / k_std).to(torch.float8_e4m3fn), (v / v_std).to(torch.float8_e4m3fn) + q = rearrange(q, "b n s d -> b s n d", n=num_heads) + k = rearrange(k, "b n s d -> b s n d", n=num_heads) + v = rearrange(v, "b n s d -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=q_std * k_std / math.sqrt(q.size(-1))) + if isinstance(x, tuple): + x = x[0] + x = x.to(origin_dtype) * v_std + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + else: + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +class ApproximateGELU(nn.Module): + def __init__(self, dim_in: int, dim_out: int, bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + return x * torch.sigmoid(1.702 * x) + +def apply_rotary_emb_qwen( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] +): + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + return x_out.type_as(x) + + +class QwenEmbedRope(nn.Module): + def __init__(self, theta: int, axes_dim: list[int], scale_rope=False): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + pos_index = torch.arange(4096) + neg_index = torch.arange(4096).flip(0) * -1 - 1 + self.pos_freqs = torch.cat([ + self.rope_params(pos_index, self.axes_dim[0], self.theta), + self.rope_params(pos_index, self.axes_dim[1], self.theta), + self.rope_params(pos_index, self.axes_dim[2], self.theta), + ], dim=1) + self.neg_freqs = torch.cat([ + self.rope_params(neg_index, self.axes_dim[0], self.theta), + self.rope_params(neg_index, self.axes_dim[1], self.theta), + self.rope_params(neg_index, self.axes_dim[2], self.theta), + ], dim=1) + self.rope_cache = {} + self.scale_rope = scale_rope + + def rope_params(self, index, dim, theta=10000): + """ + Args: + index: [0, 1, 2, 3] 1D Tensor representing the position index of the token + """ + assert dim % 2 == 0 + freqs = torch.outer( + index, + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)) + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + + def _expand_pos_freqs_if_needed(self, video_fhw, txt_seq_lens): + if isinstance(video_fhw, list): + video_fhw = tuple(max([i[j] for i in video_fhw]) for j in range(3)) + _, height, width = video_fhw + if self.scale_rope: + max_vid_index = max(height // 2, width // 2) + else: + max_vid_index = max(height, width) + required_len = max_vid_index + max(txt_seq_lens) + cur_max_len = self.pos_freqs.shape[0] + if required_len <= cur_max_len: + return + + new_max_len = math.ceil(required_len / 512) * 512 + pos_index = torch.arange(new_max_len) + neg_index = torch.arange(new_max_len).flip(0) * -1 - 1 + self.pos_freqs = torch.cat([ + self.rope_params(pos_index, self.axes_dim[0], self.theta), + self.rope_params(pos_index, self.axes_dim[1], self.theta), + self.rope_params(pos_index, self.axes_dim[2], self.theta), + ], dim=1) + self.neg_freqs = torch.cat([ + self.rope_params(neg_index, self.axes_dim[0], self.theta), + self.rope_params(neg_index, self.axes_dim[1], self.theta), + self.rope_params(neg_index, self.axes_dim[2], self.theta), + ], dim=1) + return + + + def forward(self, video_fhw, txt_seq_lens, device): + self._expand_pos_freqs_if_needed(video_fhw, txt_seq_lens) + if self.pos_freqs.device != device: + self.pos_freqs = self.pos_freqs.to(device) + self.neg_freqs = self.neg_freqs.to(device) + + vid_freqs = [] + max_vid_index = 0 + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + rope_key = f"{idx}_{height}_{width}" + + if rope_key not in self.rope_cache: + seq_lens = frame * height * width + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + if self.scale_rope: + freqs_height = torch.cat( + [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0 + ) + freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0) + freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1) + + else: + freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1) + + freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1) + self.rope_cache[rope_key] = freqs.clone().contiguous() + vid_freqs.append(self.rope_cache[rope_key]) + + if self.scale_rope: + max_vid_index = max(height // 2, width // 2, max_vid_index) + else: + max_vid_index = max(height, width, max_vid_index) + + max_len = max(txt_seq_lens) + txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + vid_freqs = torch.cat(vid_freqs, dim=0) + + return vid_freqs, txt_freqs + + + def forward_sampling(self, video_fhw, txt_seq_lens, device): + self._expand_pos_freqs_if_needed(video_fhw, txt_seq_lens) + if self.pos_freqs.device != device: + self.pos_freqs = self.pos_freqs.to(device) + self.neg_freqs = self.neg_freqs.to(device) + + vid_freqs = [] + max_vid_index = 0 + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + rope_key = f"{idx}_{height}_{width}" + if idx > 0 and f"{0}_{height}_{width}" not in self.rope_cache: + frame_0, height_0, width_0 = video_fhw[0] + + rope_key_0 = f"0_{height_0}_{width_0}" + spatial_freqs_0 = self.rope_cache[rope_key_0].reshape(frame_0, height_0, width_0, -1) + h_indices = torch.linspace(0, height_0 - 1, height).long() + w_indices = torch.linspace(0, width_0 - 1, width).long() + h_grid, w_grid = torch.meshgrid(h_indices, w_indices, indexing='ij') + sampled_rope = spatial_freqs_0[:, h_grid, w_grid, :] + + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + sampled_rope[:, :, :, :freqs_frame.shape[-1]] = freqs_frame + + seq_lens = frame * height * width + self.rope_cache[rope_key] = sampled_rope.reshape(seq_lens, -1).clone() + if rope_key not in self.rope_cache: + seq_lens = frame * height * width + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + if self.scale_rope: + freqs_height = torch.cat( + [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0 + ) + freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0) + freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1) + + else: + freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1) + + freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1) + self.rope_cache[rope_key] = freqs.clone() + vid_freqs.append(self.rope_cache[rope_key].contiguous()) + + if self.scale_rope: + max_vid_index = max(height // 2, width // 2, max_vid_index) + else: + max_vid_index = max(height, width, max_vid_index) + + max_len = max(txt_seq_lens) + txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + vid_freqs = torch.cat(vid_freqs, dim=0) + + return vid_freqs, txt_freqs + + +class QwenFeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + dropout: float = 0.0, + ): + super().__init__() + inner_dim = int(dim * 4) + self.net = nn.ModuleList([]) + self.net.append(ApproximateGELU(dim, inner_dim)) + self.net.append(nn.Dropout(dropout)) + self.net.append(nn.Linear(inner_dim, dim_out)) + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + +class QwenDoubleStreamAttention(nn.Module): + def __init__( + self, + dim_a, + dim_b, + num_heads, + head_dim, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = nn.Linear(dim_a, dim_a) + self.to_k = nn.Linear(dim_a, dim_a) + self.to_v = nn.Linear(dim_a, dim_a) + self.norm_q = RMSNorm(head_dim, eps=1e-6) + self.norm_k = RMSNorm(head_dim, eps=1e-6) + + self.add_q_proj = nn.Linear(dim_b, dim_b) + self.add_k_proj = nn.Linear(dim_b, dim_b) + self.add_v_proj = nn.Linear(dim_b, dim_b) + self.norm_added_q = RMSNorm(head_dim, eps=1e-6) + self.norm_added_k = RMSNorm(head_dim, eps=1e-6) + + self.to_out = torch.nn.Sequential(nn.Linear(dim_a, dim_a)) + self.to_add_out = nn.Linear(dim_b, dim_b) + + def forward( + self, + image: torch.FloatTensor, + text: torch.FloatTensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + enable_fp8_attention: bool = False, + ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + img_q, img_k, img_v = self.to_q(image), self.to_k(image), self.to_v(image) + txt_q, txt_k, txt_v = self.add_q_proj(text), self.add_k_proj(text), self.add_v_proj(text) + seq_txt = txt_q.shape[1] + + img_q = rearrange(img_q, 'b s (h d) -> b h s d', h=self.num_heads) + img_k = rearrange(img_k, 'b s (h d) -> b h s d', h=self.num_heads) + img_v = rearrange(img_v, 'b s (h d) -> b h s d', h=self.num_heads) + + txt_q = rearrange(txt_q, 'b s (h d) -> b h s d', h=self.num_heads) + txt_k = rearrange(txt_k, 'b s (h d) -> b h s d', h=self.num_heads) + txt_v = rearrange(txt_v, 'b s (h d) -> b h s d', h=self.num_heads) + + img_q, img_k = self.norm_q(img_q), self.norm_k(img_k) + txt_q, txt_k = self.norm_added_q(txt_q), self.norm_added_k(txt_k) + + if image_rotary_emb is not None: + img_freqs, txt_freqs = image_rotary_emb + img_q = apply_rotary_emb_qwen(img_q, img_freqs) + img_k = apply_rotary_emb_qwen(img_k, img_freqs) + txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs) + txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs) + + joint_q = torch.cat([txt_q, img_q], dim=2) + joint_k = torch.cat([txt_k, img_k], dim=2) + joint_v = torch.cat([txt_v, img_v], dim=2) + + joint_attn_out = qwen_image_flash_attention(joint_q, joint_k, joint_v, num_heads=joint_q.shape[1], attention_mask=attention_mask, enable_fp8_attention=enable_fp8_attention).to(joint_q.dtype) + + txt_attn_output = joint_attn_out[:, :seq_txt, :] + img_attn_output = joint_attn_out[:, seq_txt:, :] + + img_attn_output = self.to_out(img_attn_output) + txt_attn_output = self.to_add_out(txt_attn_output) + + return img_attn_output, txt_attn_output + + +class QwenImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + ): + super().__init__() + + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + self.img_mod = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, 6 * dim), + ) + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = QwenDoubleStreamAttention( + dim_a=dim, + dim_b=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + ) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenFeedForward(dim=dim, dim_out=dim) + + self.txt_mod = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, 6 * dim, bias=True), + ) + self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_mlp = QwenFeedForward(dim=dim, dim_out=dim) + + def _modulate(self, x, mod_params): + shift, scale, gate = mod_params.chunk(3, dim=-1) + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1) + + def forward( + self, + image: torch.Tensor, + text: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + enable_fp8_attention = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + + img_mod_attn, img_mod_mlp = self.img_mod(temb).chunk(2, dim=-1) # [B, 3*dim] each + txt_mod_attn, txt_mod_mlp = self.txt_mod(temb).chunk(2, dim=-1) # [B, 3*dim] each + + img_normed = self.img_norm1(image) + img_modulated, img_gate = self._modulate(img_normed, img_mod_attn) + + txt_normed = self.txt_norm1(text) + txt_modulated, txt_gate = self._modulate(txt_normed, txt_mod_attn) + + img_attn_out, txt_attn_out = self.attn( + image=img_modulated, + text=txt_modulated, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + ) + + image = image + img_gate * img_attn_out + text = text + txt_gate * txt_attn_out + + img_normed_2 = self.img_norm2(image) + img_modulated_2, img_gate_2 = self._modulate(img_normed_2, img_mod_mlp) + + txt_normed_2 = self.txt_norm2(text) + txt_modulated_2, txt_gate_2 = self._modulate(txt_normed_2, txt_mod_mlp) + + img_mlp_out = self.img_mlp(img_modulated_2) + txt_mlp_out = self.txt_mlp(txt_modulated_2) + + image = image + img_gate_2 * img_mlp_out + text = text + txt_gate_2 * txt_mlp_out + + return text, image + + +class QwenImageDiT(torch.nn.Module): + def __init__( + self, + num_layers: int = 60, + ): + super().__init__() + + self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=[16,56,56], scale_rope=True) + + self.time_text_embed = TimestepEmbeddings(256, 3072, diffusers_compatible_format=True, scale=1000, align_dtype_to_timestep=True) + self.txt_norm = RMSNorm(3584, eps=1e-6) + + self.img_in = nn.Linear(64, 3072) + self.txt_in = nn.Linear(3584, 3072) + + print(5 / 0) + print(sb) + + self.transformer_blocks = nn.ModuleList( + [ + QwenImageTransformerBlock( + dim=3072, + num_attention_heads=24, + attention_head_dim=128, + ) + for _ in range(num_layers) + ] + ) + self.norm_out = AdaLayerNorm(3072, single=True) + self.proj_out = nn.Linear(3072, 64) + + + def process_entity_masks(self, latents, prompt_emb, prompt_emb_mask, entity_prompt_emb, entity_prompt_emb_mask, entity_masks, height, width, image, img_shapes): + # prompt_emb + all_prompt_emb = entity_prompt_emb + [prompt_emb] + all_prompt_emb = [self.txt_in(self.txt_norm(local_prompt_emb)) for local_prompt_emb in all_prompt_emb] + all_prompt_emb = torch.cat(all_prompt_emb, dim=1) + + # image_rotary_emb + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + entity_seq_lens = [emb_mask.sum(dim=1).tolist() for emb_mask in entity_prompt_emb_mask] + entity_rotary_emb = [self.pos_embed(img_shapes, entity_seq_len, device=latents.device)[1] for entity_seq_len in entity_seq_lens] + txt_rotary_emb = torch.cat(entity_rotary_emb + [image_rotary_emb[1]], dim=0) + image_rotary_emb = (image_rotary_emb[0], txt_rotary_emb) + + # attention_mask + repeat_dim = latents.shape[1] + max_masks = entity_masks.shape[1] + entity_masks = entity_masks.repeat(1, 1, repeat_dim, 1, 1) + entity_masks = [entity_masks[:, i, None].squeeze(1) for i in range(max_masks)] + global_mask = torch.ones_like(entity_masks[0]).to(device=latents.device, dtype=latents.dtype) + entity_masks = entity_masks + [global_mask] + + N = len(entity_masks) + batch_size = entity_masks[0].shape[0] + seq_lens = [mask_.sum(dim=1).item() for mask_ in entity_prompt_emb_mask] + [prompt_emb_mask.sum(dim=1).item()] + total_seq_len = sum(seq_lens) + image.shape[1] + patched_masks = [] + for i in range(N): + patched_mask = rearrange(entity_masks[i], "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + patched_masks.append(patched_mask) + attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), dtype=torch.bool).to(device=entity_masks[0].device) + + # prompt-image attention mask + image_start = sum(seq_lens) + image_end = total_seq_len + cumsum = [0] + single_image_seq = image_end - image_start + for length in seq_lens: + cumsum.append(cumsum[-1] + length) + for i in range(N): + prompt_start = cumsum[i] + prompt_end = cumsum[i+1] + image_mask = torch.sum(patched_masks[i], dim=-1) > 0 + image_mask = image_mask.unsqueeze(1).repeat(1, seq_lens[i], 1) + # repeat image mask to match the single image sequence length + repeat_time = single_image_seq // image_mask.shape[-1] + image_mask = image_mask.repeat(1, 1, repeat_time) + # prompt update with image + attention_mask[:, prompt_start:prompt_end, image_start:image_end] = image_mask + # image update with prompt + attention_mask[:, image_start:image_end, prompt_start:prompt_end] = image_mask.transpose(1, 2) + # prompt-prompt attention mask, let the prompt tokens not attend to each other + for i in range(N): + for j in range(N): + if i == j: + continue + start_i, end_i = cumsum[i], cumsum[i+1] + start_j, end_j = cumsum[j], cumsum[j+1] + attention_mask[:, start_i:end_i, start_j:end_j] = False + + attention_mask = attention_mask.float() + attention_mask[attention_mask == 0] = float('-inf') + attention_mask[attention_mask == 1] = 0 + attention_mask = attention_mask.to(device=latents.device, dtype=latents.dtype).unsqueeze(1) + + return all_prompt_emb, image_rotary_emb, attention_mask + + + def forward( + self, + latents=None, + timestep=None, + prompt_emb=None, + prompt_emb_mask=None, + height=None, + width=None, + ): + img_shapes = [(latents.shape[0], latents.shape[2]//2, latents.shape[3]//2)] + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + + image = rearrange(latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + image = self.img_in(image) + text = self.txt_in(self.txt_norm(prompt_emb)) + + conditioning = self.time_text_embed(timestep, image.dtype) + + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + + for block in self.transformer_blocks: + text, image = block( + image=image, + text=text, + temb=conditioning, + image_rotary_emb=image_rotary_emb, + ) + + image = self.norm_out(image, conditioning) + image = self.proj_out(image) + + latents = rearrange(image, "B (H W) (C P Q) -> B C (H P) (W Q)", H=height//16, W=width//16, P=2, Q=2) + return image diff --git a/diffsynth/models/qwen_image_text_encoder.py b/diffsynth/models/qwen_image_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..f19d2d8ae2a61fd9cc45414a6bac18d28e4edcc9 --- /dev/null +++ b/diffsynth/models/qwen_image_text_encoder.py @@ -0,0 +1,190 @@ +import torch +from typing import Optional, Union + + +class QwenImageTextEncoder(torch.nn.Module): + def __init__(self): + super().__init__() + from transformers import Qwen2_5_VLConfig, Qwen2_5_VLModel + config = Qwen2_5_VLConfig(**{ + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "image_token_id": 151655, + "initializer_range": 0.02, + "intermediate_size": 18944, + "max_position_embeddings": 128000, + "max_window_layers": 28, + "model_type": "qwen2_5_vl", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_section": [ + 16, + 24, + 24 + ], + "rope_type": "default", + "type": "default" + }, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "text_config": { + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "image_token_id": None, + "initializer_range": 0.02, + "intermediate_size": 18944, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 128000, + "max_window_layers": 28, + "model_type": "qwen2_5_vl_text", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_section": [ + 16, + 24, + 24 + ], + "rope_type": "default", + "type": "default" + }, + "rope_theta": 1000000.0, + "sliding_window": None, + "torch_dtype": "float32", + "use_cache": True, + "use_sliding_window": False, + "video_token_id": None, + "vision_end_token_id": 151653, + "vision_start_token_id": 151652, + "vision_token_id": 151654, + "vocab_size": 152064 + }, + "tie_word_embeddings": False, + "torch_dtype": "float32", + "transformers_version": "4.54.0", + "use_cache": True, + "use_sliding_window": False, + "video_token_id": 151656, + "vision_config": { + "depth": 32, + "fullatt_block_indexes": [ + 7, + 15, + 23, + 31 + ], + "hidden_act": "silu", + "hidden_size": 1280, + "in_channels": 3, + "in_chans": 3, + "initializer_range": 0.02, + "intermediate_size": 3420, + "model_type": "qwen2_5_vl", + "num_heads": 16, + "out_hidden_size": 3584, + "patch_size": 14, + "spatial_merge_size": 2, + "spatial_patch_size": 14, + "temporal_patch_size": 2, + "tokens_per_second": 2, + "torch_dtype": "float32", + "window_size": 112 + }, + "vision_end_token_id": 151653, + "vision_start_token_id": 151652, + "vision_token_id": 151654, + "vocab_size": 152064 + }) + self.model = Qwen2_5_VLModel(config) + self.lm_head = torch.nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.config = config + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs, + ): + output_attentions = False + output_hidden_states = True + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **kwargs, + ) + return outputs.hidden_states diff --git a/diffsynth/models/qwen_image_vae.py b/diffsynth/models/qwen_image_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..cb047131058c703da701ca9417f71a3eba94d1ea --- /dev/null +++ b/diffsynth/models/qwen_image_vae.py @@ -0,0 +1,723 @@ +import torch +from typing import List, Optional, Tuple, Union +from torch import nn + + +CACHE_T = 2 + +class QwenImageCausalConv3d(torch.nn.Conv3d): + r""" + A custom 3D causal convolution layer with feature caching support. + + This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature + caching for efficient inference. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int or tuple, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + + # Set up causal padding + self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = torch.nn.functional.pad(x, padding) + return super().forward(x) + + + +class QwenImageRMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return torch.nn.functional.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + + +class QwenImageResidualBlock(nn.Module): + r""" + A custom residual block module. + + Args: + in_dim (int): Number of input channels. + out_dim (int): Number of output channels. + dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. + non_linearity (str, optional): Type of non-linearity to use. Default is "silu". + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + non_linearity: str = "silu", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = torch.nn.SiLU() + + # layers + self.norm1 = QwenImageRMS_norm(in_dim, images=False) + self.conv1 = QwenImageCausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = QwenImageRMS_norm(out_dim, images=False) + self.dropout = nn.Dropout(dropout) + self.conv2 = QwenImageCausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = QwenImageCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + # Apply shortcut connection + h = self.conv_shortcut(x) + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + # Second normalization and activation + x = self.norm2(x) + x = self.nonlinearity(x) + + # Dropout + x = self.dropout(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv2(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv2(x) + + # Add residual connection + return x + h + + + +class QwenImageAttentionBlock(nn.Module): + r""" + Causal self-attention with a single head. + + Args: + dim (int): The number of channels in the input tensor. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = QwenImageRMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + def forward(self, x): + identity = x + batch_size, channels, time, height, width = x.size() + + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + x = self.norm(x) + + # compute query, key, value + qkv = self.to_qkv(x) + qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) + qkv = qkv.permute(0, 1, 3, 2).contiguous() + q, k, v = qkv.chunk(3, dim=-1) + + # apply attention + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + + # output projection + x = self.proj(x) + + # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] + x = x.view(batch_size, time, channels, height, width) + x = x.permute(0, 2, 1, 3, 4) + + return x + identity + + + +class QwenImageUpsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x): + return super().forward(x.float()).type_as(x) + + + +class QwenImageResample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + - 'downsample2d': 2D downsampling with zero-padding and convolution. + - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str) -> None: + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) + self.time_conv = QwenImageCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = QwenImageCausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + # cache last frame of last two chunk + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + + +class QwenImageMidBlock(nn.Module): + """ + Middle block for WanVAE encoder and decoder. + + Args: + dim (int): Number of input/output channels. + dropout (float): Dropout rate. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [QwenImageResidualBlock(dim, dim, dropout, non_linearity)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImageAttentionBlock(dim)) + resnets.append(QwenImageResidualBlock(dim, dim, dropout, non_linearity)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + # First residual block + x = self.resnets[0](x, feat_cache, feat_idx) + + # Process through attention and residual blocks + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + x = attn(x) + + x = resnet(x, feat_cache, feat_idx) + + return x + + + +class QwenImageEncoder3d(nn.Module): + r""" + A 3D encoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_downsample (list of bool): Whether to downsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + non_linearity: str = "silu", + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.nonlinearity = torch.nn.SiLU() + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = QwenImageCausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + self.down_blocks = torch.nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + self.down_blocks.append(QwenImageResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + self.down_blocks.append(QwenImageAttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + self.down_blocks.append(QwenImageResample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = QwenImageMidBlock(out_dim, dropout, non_linearity, num_layers=1) + + # output blocks + self.norm_out = QwenImageRMS_norm(out_dim, images=False) + self.conv_out = QwenImageCausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## downsamples + for layer in self.down_blocks: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + x = self.mid_block(x, feat_cache, feat_idx) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + + +class QwenImageUpBlock(nn.Module): + """ + A block that handles upsampling for the WanVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: Optional[str] = None, + non_linearity: str = "silu", + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # Create layers list + resnets = [] + # Add residual blocks and attention if needed + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImageResidualBlock(current_dim, out_dim, dropout, non_linearity)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + self.upsamplers = None + if upsample_mode is not None: + self.upsamplers = nn.ModuleList([QwenImageResample(out_dim, mode=upsample_mode)]) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache, feat_idx) + else: + x = resnet(x) + + if self.upsamplers is not None: + if feat_cache is not None: + x = self.upsamplers[0](x, feat_cache, feat_idx) + else: + x = self.upsamplers[0](x) + return x + + + +class QwenImageDecoder3d(nn.Module): + r""" + A 3D decoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_upsample (list of bool): Whether to upsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + non_linearity: str = "silu", + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + self.nonlinearity = torch.nn.SiLU() + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + # init block + self.conv_in = QwenImageCausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = QwenImageMidBlock(dims[0], dropout, non_linearity, num_layers=1) + + # upsample blocks + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i > 0: + in_dim = in_dim // 2 + + # Determine if we need upsampling + upsample_mode = None + if i != len(dim_mult) - 1: + upsample_mode = "upsample3d" if temperal_upsample[i] else "upsample2d" + + # Create and add the upsampling block + up_block = QwenImageUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + non_linearity=non_linearity, + ) + self.up_blocks.append(up_block) + + # Update scale for next iteration + if upsample_mode is not None: + scale *= 2.0 + + # output blocks + self.norm_out = QwenImageRMS_norm(out_dim, images=False) + self.conv_out = QwenImageCausalConv3d(out_dim, 3, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## middle + x = self.mid_block(x, feat_cache, feat_idx) + + ## upsamples + for up_block in self.up_blocks: + x = up_block(x, feat_cache, feat_idx) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + + +class QwenImageVAE(torch.nn.Module): + def __init__( + self, + base_dim: int = 96, + z_dim: int = 16, + dim_mult: Tuple[int] = [1, 2, 4, 4], + num_res_blocks: int = 2, + attn_scales: List[float] = [], + temperal_downsample: List[bool] = [False, True, True], + dropout: float = 0.0, + ) -> None: + super().__init__() + + self.z_dim = z_dim + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + self.encoder = QwenImageEncoder3d( + base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.quant_conv = QwenImageCausalConv3d(z_dim * 2, z_dim * 2, 1) + self.post_quant_conv = QwenImageCausalConv3d(z_dim, z_dim, 1) + + self.decoder = QwenImageDecoder3d( + base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + ) + + mean = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ] + std = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ] + self.mean = torch.tensor(mean).view(1, 16, 1, 1, 1) + self.std = 1 / torch.tensor(std).view(1, 16, 1, 1, 1) + + def encode(self, x, **kwargs): + x = x.unsqueeze(2) + x = self.encoder(x) + x = self.quant_conv(x) + x = x[:, :16] + mean, std = self.mean.to(dtype=x.dtype, device=x.device), self.std.to(dtype=x.dtype, device=x.device) + x = (x - mean) * std + x = x.squeeze(2) + return x + + def decode(self, x, **kwargs): + x = x.unsqueeze(2) + mean, std = self.mean.to(dtype=x.dtype, device=x.device), self.std.to(dtype=x.dtype, device=x.device) + x = x / std + mean + x = self.post_quant_conv(x) + x = self.decoder(x) + x = x.squeeze(2) + return x diff --git a/diffsynth/models/sd_text_encoder.py b/diffsynth/models/sd_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b0a1171c265a43048c1623bd0c2375f4fc3f5e5d --- /dev/null +++ b/diffsynth/models/sd_text_encoder.py @@ -0,0 +1,412 @@ +import torch +from .attention import Attention +from einops import rearrange + + +def low_version_attention(query, key, value, attn_bias=None): + scale = 1 / query.shape[-1] ** 0.5 + query = query * scale + attn = torch.matmul(query, key.transpose(-2, -1)) + if attn_bias is not None: + attn = attn + attn_bias + attn = attn.softmax(-1) + return attn @ value + + +class Attention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) + self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out) + + def interact_with_ipadapter(self, hidden_states, q, ip_k, ip_v, scale=1.0): + batch_size = q.shape[0] + ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v) + hidden_states = hidden_states + scale * ip_hidden_states + return hidden_states + + def torch_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + batch_size = encoder_hidden_states.shape[0] + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) + + if qkv_preprocessor is not None: + q, k, v = qkv_preprocessor(q, k, v) + + hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + if ipadapter_kwargs is not None: + hidden_states = self.interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) + hidden_states = hidden_states.to(q.dtype) + + hidden_states = self.to_out(hidden_states) + + return hidden_states + + def xformers_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + q = self.to_q(hidden_states) + k = self.to_k(encoder_hidden_states) + v = self.to_v(encoder_hidden_states) + + q = rearrange(q, "b f (n d) -> (b n) f d", n=self.num_heads) + k = rearrange(k, "b f (n d) -> (b n) f d", n=self.num_heads) + v = rearrange(v, "b f (n d) -> (b n) f d", n=self.num_heads) + + if attn_mask is not None: + hidden_states = low_version_attention(q, k, v, attn_bias=attn_mask) + else: + import xformers.ops as xops + hidden_states = xops.memory_efficient_attention(q, k, v) + hidden_states = rearrange(hidden_states, "(b n) f d -> b f (n d)", n=self.num_heads) + + hidden_states = hidden_states.to(q.dtype) + hidden_states = self.to_out(hidden_states) + + return hidden_states + + def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None): + return self.torch_forward(hidden_states, encoder_hidden_states=encoder_hidden_states, attn_mask=attn_mask, ipadapter_kwargs=ipadapter_kwargs, qkv_preprocessor=qkv_preprocessor) + + + + + +class CLIPEncoderLayer(torch.nn.Module): + def __init__(self, embed_dim, intermediate_size, num_heads=12, head_dim=64, use_quick_gelu=True): + super().__init__() + self.attn = Attention(q_dim=embed_dim, num_heads=num_heads, head_dim=head_dim, bias_q=True, bias_kv=True, bias_out=True) + self.layer_norm1 = torch.nn.LayerNorm(embed_dim) + self.layer_norm2 = torch.nn.LayerNorm(embed_dim) + self.fc1 = torch.nn.Linear(embed_dim, intermediate_size) + self.fc2 = torch.nn.Linear(intermediate_size, embed_dim) + + self.use_quick_gelu = use_quick_gelu + + def quickGELU(self, x): + return x * torch.sigmoid(1.702 * x) + + def forward(self, hidden_states, attn_mask=None): + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.attn(hidden_states, attn_mask=attn_mask) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.fc1(hidden_states) + if self.use_quick_gelu: + hidden_states = self.quickGELU(hidden_states) + else: + hidden_states = torch.nn.functional.gelu(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class SDTextEncoder(torch.nn.Module): + def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072): + super().__init__() + + # token_embedding + self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim) + + # position_embeds (This is a fixed tensor) + self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim)) + + # encoders + self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)]) + + # attn_mask + self.attn_mask = self.attention_mask(max_position_embeddings) + + # final_layer_norm + self.final_layer_norm = torch.nn.LayerNorm(embed_dim) + + def attention_mask(self, length): + mask = torch.empty(length, length) + mask.fill_(float("-inf")) + mask.triu_(1) + return mask + + def forward(self, input_ids, clip_skip=1): + embeds = self.token_embedding(input_ids) + self.position_embeds + attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype) + for encoder_id, encoder in enumerate(self.encoders): + embeds = encoder(embeds, attn_mask=attn_mask) + if encoder_id + clip_skip == len(self.encoders): + break + embeds = self.final_layer_norm(embeds) + return embeds + + @staticmethod + def state_dict_converter(): + return SDTextEncoderStateDictConverter() + + +class SDTextEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + rename_dict = { + "text_model.embeddings.token_embedding.weight": "token_embedding.weight", + "text_model.embeddings.position_embedding.weight": "position_embeds", + "text_model.final_layer_norm.weight": "final_layer_norm.weight", + "text_model.final_layer_norm.bias": "final_layer_norm.bias" + } + attn_rename_dict = { + "self_attn.q_proj": "attn.to_q", + "self_attn.k_proj": "attn.to_k", + "self_attn.v_proj": "attn.to_v", + "self_attn.out_proj": "attn.to_out", + "layer_norm1": "layer_norm1", + "layer_norm2": "layer_norm2", + "mlp.fc1": "fc1", + "mlp.fc2": "fc2", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + if name == "text_model.embeddings.position_embedding.weight": + param = param.reshape((1, param.shape[0], param.shape[1])) + state_dict_[rename_dict[name]] = param + elif name.startswith("text_model.encoder.layers."): + param = state_dict[name] + names = name.split(".") + layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1] + name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail]) + state_dict_[name_] = param + return state_dict_ + + def from_civitai(self, state_dict): + rename_dict = { + "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight": "token_embedding.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.bias": "encoders.0.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.weight": "encoders.0.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.bias": "encoders.0.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.weight": "encoders.0.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.bias": "encoders.0.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.weight": "encoders.0.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.bias": "encoders.0.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.weight": "encoders.0.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.bias": "encoders.0.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.weight": "encoders.0.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.bias": "encoders.0.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.weight": "encoders.0.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.bias": "encoders.0.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.weight": "encoders.0.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.bias": "encoders.0.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.weight": "encoders.0.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.bias": "encoders.1.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.weight": "encoders.1.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.bias": "encoders.1.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.weight": "encoders.1.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.bias": "encoders.1.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.weight": "encoders.1.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.bias": "encoders.1.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.weight": "encoders.1.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.bias": "encoders.1.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.weight": "encoders.1.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.bias": "encoders.1.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.weight": "encoders.1.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.bias": "encoders.1.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.weight": "encoders.1.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.bias": "encoders.1.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.weight": "encoders.1.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.bias": "encoders.10.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.weight": "encoders.10.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.bias": "encoders.10.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.weight": "encoders.10.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.bias": "encoders.10.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.weight": "encoders.10.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.bias": "encoders.10.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.weight": "encoders.10.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.bias": "encoders.10.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.weight": "encoders.10.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.bias": "encoders.10.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.weight": "encoders.10.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.bias": "encoders.10.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.weight": "encoders.10.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.bias": "encoders.10.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.weight": "encoders.10.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.bias": "encoders.11.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.weight": "encoders.11.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.bias": "encoders.11.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.weight": "encoders.11.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.bias": "encoders.11.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.weight": "encoders.11.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.bias": "encoders.11.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.weight": "encoders.11.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.bias": "encoders.11.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.weight": "encoders.11.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.bias": "encoders.11.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.weight": "encoders.11.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.bias": "encoders.11.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.weight": "encoders.11.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.bias": "encoders.11.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.weight": "encoders.11.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.bias": "encoders.2.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.weight": "encoders.2.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.bias": "encoders.2.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.weight": "encoders.2.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.bias": "encoders.2.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.weight": "encoders.2.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.bias": "encoders.2.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.weight": "encoders.2.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.bias": "encoders.2.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.weight": "encoders.2.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.bias": "encoders.2.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.weight": "encoders.2.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.bias": "encoders.2.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.weight": "encoders.2.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.bias": "encoders.2.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.weight": "encoders.2.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.bias": "encoders.3.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.weight": "encoders.3.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.bias": "encoders.3.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.weight": "encoders.3.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.bias": "encoders.3.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.weight": "encoders.3.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.bias": "encoders.3.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.weight": "encoders.3.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.bias": "encoders.3.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.weight": "encoders.3.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.bias": "encoders.3.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.weight": "encoders.3.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.bias": "encoders.3.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.weight": "encoders.3.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.bias": "encoders.3.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.weight": "encoders.3.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.bias": "encoders.4.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.weight": "encoders.4.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.bias": "encoders.4.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.weight": "encoders.4.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.bias": "encoders.4.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.weight": "encoders.4.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.bias": "encoders.4.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.weight": "encoders.4.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.bias": "encoders.4.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.weight": "encoders.4.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.bias": "encoders.4.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.weight": "encoders.4.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.bias": "encoders.4.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.weight": "encoders.4.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.bias": "encoders.4.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.weight": "encoders.4.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.bias": "encoders.5.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.weight": "encoders.5.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.bias": "encoders.5.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.weight": "encoders.5.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.bias": "encoders.5.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.weight": "encoders.5.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.bias": "encoders.5.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.weight": "encoders.5.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.bias": "encoders.5.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.weight": "encoders.5.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.bias": "encoders.5.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.weight": "encoders.5.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.bias": "encoders.5.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.weight": "encoders.5.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.bias": "encoders.5.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.weight": "encoders.5.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.bias": "encoders.6.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.weight": "encoders.6.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.bias": "encoders.6.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.weight": "encoders.6.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.bias": "encoders.6.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.weight": "encoders.6.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.bias": "encoders.6.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.weight": "encoders.6.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.bias": "encoders.6.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.weight": "encoders.6.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.bias": "encoders.6.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.weight": "encoders.6.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.bias": "encoders.6.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.weight": "encoders.6.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.bias": "encoders.6.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.weight": "encoders.6.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.bias": "encoders.7.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.weight": "encoders.7.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.bias": "encoders.7.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.weight": "encoders.7.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.bias": "encoders.7.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.weight": "encoders.7.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.bias": "encoders.7.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.weight": "encoders.7.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.bias": "encoders.7.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.weight": "encoders.7.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.bias": "encoders.7.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.weight": "encoders.7.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.bias": "encoders.7.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.weight": "encoders.7.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.bias": "encoders.7.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.weight": "encoders.7.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.bias": "encoders.8.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.weight": "encoders.8.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.bias": "encoders.8.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.weight": "encoders.8.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.bias": "encoders.8.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.weight": "encoders.8.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.bias": "encoders.8.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.weight": "encoders.8.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.bias": "encoders.8.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.weight": "encoders.8.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.bias": "encoders.8.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.weight": "encoders.8.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.bias": "encoders.8.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.weight": "encoders.8.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.bias": "encoders.8.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.weight": "encoders.8.attn.to_v.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.bias": "encoders.9.layer_norm1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.weight": "encoders.9.layer_norm1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.bias": "encoders.9.layer_norm2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.weight": "encoders.9.layer_norm2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.bias": "encoders.9.fc1.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.weight": "encoders.9.fc1.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.bias": "encoders.9.fc2.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.weight": "encoders.9.fc2.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.bias": "encoders.9.attn.to_k.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.weight": "encoders.9.attn.to_k.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.bias": "encoders.9.attn.to_out.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.weight": "encoders.9.attn.to_out.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.bias": "encoders.9.attn.to_q.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.weight": "encoders.9.attn.to_q.weight", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.bias": "encoders.9.attn.to_v.bias", + "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.weight": "encoders.9.attn.to_v.weight", + "cond_stage_model.transformer.text_model.final_layer_norm.bias": "final_layer_norm.bias", + "cond_stage_model.transformer.text_model.final_layer_norm.weight": "final_layer_norm.weight", + "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight": "position_embeds" + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + if name == "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight": + param = param.reshape((1, param.shape[0], param.shape[1])) + state_dict_[rename_dict[name]] = param + return state_dict_ diff --git a/diffsynth/models/step1x_connector.py b/diffsynth/models/step1x_connector.py new file mode 100644 index 0000000000000000000000000000000000000000..225c8fbcb54f8daebf48656141e9a8998d002fd8 --- /dev/null +++ b/diffsynth/models/step1x_connector.py @@ -0,0 +1,663 @@ +from typing import Optional + +import torch, math +import torch.nn +from einops import rearrange +from torch import nn +from functools import partial +from einops import rearrange + + + +def attention(q, k, v, attn_mask, mode="torch"): + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + x = rearrange(x, "b n s d -> b s (n d)") + return x + + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = (bias, bias) + drop_probs = (drop, drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer( + in_channels, hidden_channels, bias=bias[0], device=device, dtype=dtype + ) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = ( + norm_layer(hidden_channels, device=device, dtype=dtype) + if norm_layer is not None + else nn.Identity() + ) + self.fc2 = linear_layer( + hidden_channels, out_features, bias=bias[1], device=device, dtype=dtype + ) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear( + frequency_embedding_size, hidden_size, bias=True, **factory_kwargs + ), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) # type: ignore + nn.init.normal_(self.mlp[2].weight, std=0.02) # type: ignore + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding( + t, self.frequency_embedding_size, self.max_period + ).type(t.dtype) # type: ignore + t_emb = self.mlp(t_freq) + return t_emb + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") + +class IndividualTokenRefinerBlock(torch.nn.Module): + def __init__( + self, + hidden_size, + heads_num, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + need_CA: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.need_CA = need_CA + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + self.self_attn_qkv = nn.Linear( + hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs + ) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_proj = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + self.norm2 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + + if self.need_CA: + self.cross_attnblock=CrossAttnBlock(hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs,) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + y: torch.Tensor = None, + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, mode="torch", attn_mask=attn_mask) + + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + if self.need_CA: + x = self.cross_attnblock(x, c, attn_mask, y) + + # FFN Layer + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + + return x + + + + +class CrossAttnBlock(torch.nn.Module): + def __init__( + self, + hidden_size, + heads_num, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + + self.norm1 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + self.norm1_2 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + self.self_attn_q = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.self_attn_kv = nn.Linear( + hidden_size, hidden_size*2, bias=qkv_bias, **factory_kwargs + ) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_proj = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + self.norm2 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + act_layer = get_activation_layer(act_type) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + y: torch.Tensor=None, + + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + norm_y = self.norm1_2(y) + q = self.self_attn_q(norm_x) + q = rearrange(q, "B L (H D) -> B L H D", H=self.heads_num) + kv = self.self_attn_kv(norm_y) + k, v = rearrange(kv, "B L (K H D) -> K B L H D", K=2, H=self.heads_num) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, mode="torch", attn_mask=attn_mask) + + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + return x + + + +class IndividualTokenRefiner(torch.nn.Module): + def __init__( + self, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + need_CA:bool=False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.need_CA = need_CA + self.blocks = nn.ModuleList( + [ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + need_CA=self.need_CA, + **factory_kwargs, + ) + for _ in range(depth) + ] + ) + + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + y:torch.Tensor=None, + ): + self_attn_mask = None + if mask is not None: + batch_size = mask.shape[0] + seq_len = mask.shape[1] + mask = mask.to(x.device) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat( + 1, 1, seq_len, 1 + ) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + # batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + # avoids self-attention weight being NaN for padding tokens + self_attn_mask[:, :, :, 0] = True + + + for block in self.blocks: + x = block(x, c, self_attn_mask,y) + + return x + + +class SingleTokenRefiner(torch.nn.Module): + """ + A single token refiner block for llm text embedding refine. + """ + def __init__( + self, + in_channels, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + need_CA:bool=False, + attn_mode: str = "torch", + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.attn_mode = attn_mode + self.need_CA = need_CA + assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner." + + self.input_embedder = nn.Linear( + in_channels, hidden_size, bias=True, **factory_kwargs + ) + if self.need_CA: + self.input_embedder_CA = nn.Linear( + in_channels, hidden_size, bias=True, **factory_kwargs + ) + + act_layer = get_activation_layer(act_type) + # Build timestep embedding layer + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + # Build context embedding layer + self.c_embedder = TextProjection( + in_channels, hidden_size, act_layer, **factory_kwargs + ) + + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + need_CA=need_CA, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + y: torch.LongTensor=None, + ): + timestep_aware_representations = self.t_embedder(t) + + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.unsqueeze(-1) # [b, s1, 1] + context_aware_representations = (x * mask_float).sum( + dim=1 + ) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder(context_aware_representations) + c = timestep_aware_representations + context_aware_representations + + x = self.input_embedder(x) + if self.need_CA: + y = self.input_embedder_CA(y) + x = self.individual_token_refiner(x, c, mask, y) + else: + x = self.individual_token_refiner(x, c, mask) + + return x + + +class Qwen2Connector(torch.nn.Module): + def __init__( + self, + # biclip_dim=1024, + in_channels=3584, + hidden_size=4096, + heads_num=32, + depth=2, + need_CA=False, + device=None, + dtype=torch.bfloat16, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype":dtype} + + self.S =SingleTokenRefiner(in_channels=in_channels,hidden_size=hidden_size,heads_num=heads_num,depth=depth,need_CA=need_CA,**factory_kwargs) + self.global_proj_out=nn.Linear(in_channels,768) + + self.scale_factor = nn.Parameter(torch.zeros(1)) + with torch.no_grad(): + self.scale_factor.data += -(1 - 0.09) + + def forward(self, x,t,mask): + mask_float = mask.unsqueeze(-1) # [b, s1, 1] + x_mean = (x * mask_float).sum( + dim=1 + ) / mask_float.sum(dim=1) * (1 + self.scale_factor.to(dtype=x.dtype, device=x.device)) + + global_out=self.global_proj_out(x_mean) + encoder_hidden_states = self.S(x,t,mask) + return encoder_hidden_states,global_out diff --git a/diffsynth/models/step1x_text_encoder.py b/diffsynth/models/step1x_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..d0fe22157e0938b0821c5a7dc5d0376d2bc35b6a --- /dev/null +++ b/diffsynth/models/step1x_text_encoder.py @@ -0,0 +1,194 @@ +import torch +from typing import Optional, Union +from .qwen_image_text_encoder import QwenImageTextEncoder + + +class Step1xEditEmbedder(torch.nn.Module): + def __init__(self, model: QwenImageTextEncoder, processor, max_length=640, dtype=torch.bfloat16, device="cuda"): + super().__init__() + self.max_length = max_length + self.dtype = dtype + self.device = device + + Qwen25VL_7b_PREFIX = '''Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt: +- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes. +- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n +Here are examples of how to transform or refine prompts: +- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers. +- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n +Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations: +User Prompt:''' + + self.prefix = Qwen25VL_7b_PREFIX + self.model = model + self.processor = processor + + def model_forward( + self, + model: QwenImageTextEncoder, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs, + ): + output_attentions = output_attentions if output_attentions is not None else model.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else model.config.output_hidden_states + ) + + outputs = model.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **kwargs, + ) + return outputs.hidden_states + + def forward(self, caption, ref_images): + text_list = caption + embs = torch.zeros( + len(text_list), + self.max_length, + self.model.config.hidden_size, + dtype=torch.bfloat16, + device=torch.cuda.current_device(), + ) + masks = torch.zeros( + len(text_list), + self.max_length, + dtype=torch.long, + device=torch.cuda.current_device(), + ) + + def split_string(s): + s = s.replace("“", '"').replace("”", '"').replace("'", '''"''') # use english quotes + result = [] + in_quotes = False + temp = "" + + for idx,char in enumerate(s): + if char == '"' and idx>155: + temp += char + if not in_quotes: + result.append(temp) + temp = "" + + in_quotes = not in_quotes + continue + if in_quotes: + if char.isspace(): + pass # have space token + + result.append("“" + char + "”") + else: + temp += char + + if temp: + result.append(temp) + + return result + + for idx, (txt, imgs) in enumerate(zip(text_list, ref_images)): + + messages = [{"role": "user", "content": []}] + + messages[0]["content"].append({"type": "text", "text": f"{self.prefix}"}) + + messages[0]["content"].append({"type": "image", "image": imgs}) + + # 再添加 text + messages[0]["content"].append({"type": "text", "text": f"{txt}"}) + + # Preparation for inference + text = self.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, add_vision_id=True + ) + + image_inputs = [imgs] + + inputs = self.processor( + text=[text], + images=image_inputs, + padding=True, + return_tensors="pt", + ) + + old_inputs_ids = inputs.input_ids + text_split_list = split_string(text) + + token_list = [] + for text_each in text_split_list: + txt_inputs = self.processor( + text=text_each, + images=None, + videos=None, + padding=True, + return_tensors="pt", + ) + token_each = txt_inputs.input_ids + if token_each[0][0] == 2073 and token_each[0][-1] == 854: + token_each = token_each[:, 1:-1] + token_list.append(token_each) + else: + token_list.append(token_each) + + new_txt_ids = torch.cat(token_list, dim=1).to("cuda") + + new_txt_ids = new_txt_ids.to(old_inputs_ids.device) + + idx1 = (old_inputs_ids == 151653).nonzero(as_tuple=True)[1][0] + idx2 = (new_txt_ids == 151653).nonzero(as_tuple=True)[1][0] + inputs.input_ids = ( + torch.cat([old_inputs_ids[0, :idx1], new_txt_ids[0, idx2:]], dim=0) + .unsqueeze(0) + .to("cuda") + ) + inputs.attention_mask = (inputs.input_ids > 0).long().to("cuda") + outputs = self.model_forward( + self.model, + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + pixel_values=inputs.pixel_values.to("cuda"), + image_grid_thw=inputs.image_grid_thw.to("cuda"), + output_hidden_states=True, + ) + + emb = outputs[-1] + + embs[idx, : min(self.max_length, emb.shape[1] - 217)] = emb[0, 217:][ + : self.max_length + ] + + masks[idx, : min(self.max_length, emb.shape[1] - 217)] = torch.ones( + (min(self.max_length, emb.shape[1] - 217)), + dtype=torch.long, + device=torch.cuda.current_device(), + ) + + return embs, masks diff --git a/diffsynth/models/wan_video_animate_adapter.py b/diffsynth/models/wan_video_animate_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..3ace70d8b3162e77844f0957cd40207a54e674a9 --- /dev/null +++ b/diffsynth/models/wan_video_animate_adapter.py @@ -0,0 +1,650 @@ +import torch +import torch.nn as nn +from torch.nn import functional as F +import math +from typing import Tuple, Optional, List +from einops import rearrange + + + +MEMORY_LAYOUT = { + "flash": ( + lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]), + lambda x: x, + ), + "torch": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), + "vanilla": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), +} + + +def attention( + q, + k, + v, + mode="torch", + drop_rate=0, + attn_mask=None, + causal=False, + max_seqlen_q=None, + batch_size=1, +): + pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode] + + if mode == "torch": + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) + + x = post_attn_layout(x) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +class CausalConv1d(nn.Module): + + def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs): + super().__init__() + + self.pad_mode = pad_mode + padding = (kernel_size - 1, 0) # T + self.time_causal_padding = padding + + self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + + +class FaceEncoder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.num_heads = num_heads + self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1) + self.norm1 = nn.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.act = nn.SiLU() + self.conv2 = CausalConv1d(1024, 1024, 3, stride=2) + self.conv3 = CausalConv1d(1024, 1024, 3, stride=2) + + self.out_proj = nn.Linear(1024, hidden_dim) + self.norm1 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.norm2 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.norm3 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim)) + + def forward(self, x): + + x = rearrange(x, "b t c -> b c t") + b, c, t = x.shape + + x = self.conv1_local(x) + x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads) + + x = self.norm1(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv2(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm2(x) + x = self.act(x) + x = rearrange(x, "b t c -> b c t") + x = self.conv3(x) + x = rearrange(x, "b c t -> b t c") + x = self.norm3(x) + x = self.act(x) + x = self.out_proj(x) + x = rearrange(x, "(b n) t c -> b t n c", b=b) + padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1) + x = torch.cat([x, padding], dim=-2) + x_local = x.clone() + + return x_local + + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") + + +class FaceAdapter(nn.Module): + def __init__( + self, + hidden_dim: int, + heads_num: int, + qk_norm: bool = True, + qk_norm_type: str = "rms", + num_adapter_layers: int = 1, + dtype=None, + device=None, + ): + + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.hidden_size = hidden_dim + self.heads_num = heads_num + self.fuser_blocks = nn.ModuleList( + [ + FaceBlock( + self.hidden_size, + self.heads_num, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) + for _ in range(num_adapter_layers) + ] + ) + + def forward( + self, + x: torch.Tensor, + motion_embed: torch.Tensor, + idx: int, + freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None, + freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None, + ) -> torch.Tensor: + + return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k) + + + +class FaceBlock(nn.Module): + def __init__( + self, + hidden_size: int, + heads_num: int, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + self.scale = qk_scale or head_dim**-0.5 + + self.linear1_kv = nn.Linear(hidden_size, hidden_size * 2, **factory_kwargs) + self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + + self.linear2 = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + + self.pre_norm_feat = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.pre_norm_motion = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + def forward( + self, + x: torch.Tensor, + motion_vec: torch.Tensor, + motion_mask: Optional[torch.Tensor] = None, + use_context_parallel=False, + ) -> torch.Tensor: + + B, T, N, C = motion_vec.shape + T_comp = T + + x_motion = self.pre_norm_motion(motion_vec) + x_feat = self.pre_norm_feat(x) + + kv = self.linear1_kv(x_motion) + q = self.linear1_q(x_feat) + + k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num) + q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + k = rearrange(k, "B L N H D -> (B L) H N D") + v = rearrange(v, "B L N H D -> (B L) H N D") + + q = rearrange(q, "B (L S) H D -> (B L) H S D", L=T_comp) + # Compute attention. + attn = F.scaled_dot_product_attention(q, k, v) + + attn = rearrange(attn, "(B L) H S D -> B (L S) (H D)", L=T_comp) + + output = self.linear2(attn) + + if motion_mask is not None: + output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1) + + return output + + + +def custom_qr(input_tensor): + original_dtype = input_tensor.dtype + if original_dtype == torch.bfloat16: + q, r = torch.linalg.qr(input_tensor.to(torch.float32)) + return q.to(original_dtype), r.to(original_dtype) + return torch.linalg.qr(input_tensor) + +def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): + return F.leaky_relu(input + bias, negative_slope) * scale + + +def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1): + _, minor, in_h, in_w = input.shape + kernel_h, kernel_w = kernel.shape + + out = input.view(-1, minor, in_h, 1, in_w, 1) + out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0]) + out = out.view(-1, minor, in_h * up_y, in_w * up_x) + + out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]) + out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0), + max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ] + + out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) + w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) + out = F.conv2d(out, w) + out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, + in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, ) + return out[:, :, ::down_y, ::down_x] + + +def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): + return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]) + + +def make_kernel(k): + k = torch.tensor(k, dtype=torch.float32) + if k.ndim == 1: + k = k[None, :] * k[:, None] + k /= k.sum() + return k + + +class FusedLeakyReLU(nn.Module): + def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): + super().__init__() + self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1)) + self.negative_slope = negative_slope + self.scale = scale + + def forward(self, input): + out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) + return out + + +class Blur(nn.Module): + def __init__(self, kernel, pad, upsample_factor=1): + super().__init__() + + kernel = make_kernel(kernel) + + if upsample_factor > 1: + kernel = kernel * (upsample_factor ** 2) + + self.kernel = torch.nn.Parameter(kernel) + + self.pad = pad + + def forward(self, input): + return upfirdn2d(input, self.kernel, pad=self.pad) + + +class ScaledLeakyReLU(nn.Module): + def __init__(self, negative_slope=0.2): + super().__init__() + + self.negative_slope = negative_slope + + def forward(self, input): + return F.leaky_relu(input, negative_slope=self.negative_slope) + + +class EqualConv2d(nn.Module): + def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size)) + self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) + + self.stride = stride + self.padding = padding + + if bias: + self.bias = nn.Parameter(torch.zeros(out_channel)) + else: + self.bias = None + + def forward(self, input): + + return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding) + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' + f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' + ) + + +class EqualLinear(nn.Module): + def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) + else: + self.bias = None + + self.activation = activation + + self.scale = (1 / math.sqrt(in_dim)) * lr_mul + self.lr_mul = lr_mul + + def forward(self, input): + + if self.activation: + out = F.linear(input, self.weight * self.scale) + out = fused_leaky_relu(out, self.bias * self.lr_mul) + else: + out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul) + + return out + + def __repr__(self): + return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})') + + +class ConvLayer(nn.Sequential): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + downsample=False, + blur_kernel=[1, 3, 3, 1], + bias=True, + activate=True, + ): + layers = [] + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + layers.append(Blur(blur_kernel, pad=(pad0, pad1))) + + stride = 2 + self.padding = 0 + + else: + stride = 1 + self.padding = kernel_size // 2 + + layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride, + bias=bias and not activate)) + + if activate: + if bias: + layers.append(FusedLeakyReLU(out_channel)) + else: + layers.append(ScaledLeakyReLU(0.2)) + + super().__init__(*layers) + + +class ResBlock(nn.Module): + def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + self.conv1 = ConvLayer(in_channel, in_channel, 3) + self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True) + + self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False) + + def forward(self, input): + out = self.conv1(input) + out = self.conv2(out) + + skip = self.skip(input) + out = (out + skip) / math.sqrt(2) + + return out + + +class EncoderApp(nn.Module): + def __init__(self, size, w_dim=512): + super(EncoderApp, self).__init__() + + channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256, + 128: 128, + 256: 64, + 512: 32, + 1024: 16 + } + + self.w_dim = w_dim + log_size = int(math.log(size, 2)) + + self.convs = nn.ModuleList() + self.convs.append(ConvLayer(3, channels[size], 1)) + + in_channel = channels[size] + for i in range(log_size, 2, -1): + out_channel = channels[2 ** (i - 1)] + self.convs.append(ResBlock(in_channel, out_channel)) + in_channel = out_channel + + self.convs.append(EqualConv2d(in_channel, self.w_dim, 4, padding=0, bias=False)) + + def forward(self, x): + + res = [] + h = x + for conv in self.convs: + h = conv(h) + res.append(h) + + return res[-1].squeeze(-1).squeeze(-1), res[::-1][2:] + + +class Encoder(nn.Module): + def __init__(self, size, dim=512, dim_motion=20): + super(Encoder, self).__init__() + + # appearance netmork + self.net_app = EncoderApp(size, dim) + + # motion network + fc = [EqualLinear(dim, dim)] + for i in range(3): + fc.append(EqualLinear(dim, dim)) + + fc.append(EqualLinear(dim, dim_motion)) + self.fc = nn.Sequential(*fc) + + def enc_app(self, x): + h_source = self.net_app(x) + return h_source + + def enc_motion(self, x): + h, _ = self.net_app(x) + h_motion = self.fc(h) + return h_motion + + +class Direction(nn.Module): + def __init__(self, motion_dim): + super(Direction, self).__init__() + self.weight = nn.Parameter(torch.randn(512, motion_dim)) + + def forward(self, input): + + weight = self.weight + 1e-8 + Q, R = custom_qr(weight) + if input is None: + return Q + else: + input_diag = torch.diag_embed(input) # alpha, diagonal matrix + out = torch.matmul(input_diag, Q.T) + out = torch.sum(out, dim=1) + return out + + +class Synthesis(nn.Module): + def __init__(self, motion_dim): + super(Synthesis, self).__init__() + self.direction = Direction(motion_dim) + + +class Generator(nn.Module): + def __init__(self, size, style_dim=512, motion_dim=20): + super().__init__() + + self.enc = Encoder(size, style_dim, motion_dim) + self.dec = Synthesis(motion_dim) + + def get_motion(self, img): + #motion_feat = self.enc.enc_motion(img) + motion_feat = torch.utils.checkpoint.checkpoint((self.enc.enc_motion), img, use_reentrant=True) + motion = self.dec.direction(motion_feat) + return motion + + +class WanAnimateAdapter(torch.nn.Module): + def __init__(self): + super().__init__() + self.pose_patch_embedding = torch.nn.Conv3d(16, 5120, kernel_size=(1, 2, 2), stride=(1, 2, 2)) + self.motion_encoder = Generator(size=512, style_dim=512, motion_dim=20) + self.face_adapter = FaceAdapter(heads_num=40, hidden_dim=5120, num_adapter_layers=40 // 5) + self.face_encoder = FaceEncoder(in_dim=512, hidden_dim=5120, num_heads=4) + + def after_patch_embedding(self, x: List[torch.Tensor], pose_latents, face_pixel_values): + pose_latents = self.pose_patch_embedding(pose_latents) + x[:, :, 1:] += pose_latents + + b,c,T,h,w = face_pixel_values.shape + face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w") + + encode_bs = 8 + face_pixel_values_tmp = [] + for i in range(math.ceil(face_pixel_values.shape[0]/encode_bs)): + face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i*encode_bs:(i+1)*encode_bs])) + + motion_vec = torch.cat(face_pixel_values_tmp) + + motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T) + motion_vec = self.face_encoder(motion_vec) + + B, L, H, C = motion_vec.shape + pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec) + motion_vec = torch.cat([pad_face, motion_vec], dim=1) + return x, motion_vec + + def after_transformer_block(self, block_idx, x, motion_vec, motion_masks=None): + if block_idx % 5 == 0: + adapter_args = [x, motion_vec, motion_masks, False] + residual_out = self.face_adapter.fuser_blocks[block_idx // 5](*adapter_args) + x = residual_out + x + return x diff --git a/diffsynth/models/wan_video_camera_controller.py b/diffsynth/models/wan_video_camera_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..45a44ee6bcd408d7ee9d18653f933151ce351a72 --- /dev/null +++ b/diffsynth/models/wan_video_camera_controller.py @@ -0,0 +1,206 @@ +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange +import os +from typing_extensions import Literal + +class SimpleAdapter(nn.Module): + def __init__(self, in_dim, out_dim, kernel_size, stride, num_residual_blocks=1): + super(SimpleAdapter, self).__init__() + + # Pixel Unshuffle: reduce spatial dimensions by a factor of 8 + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=8) + + # Convolution: reduce spatial dimensions by a factor + # of 2 (without overlap) + self.conv = nn.Conv2d(in_dim * 64, out_dim, kernel_size=kernel_size, stride=stride, padding=0) + + # Residual blocks for feature extraction + self.residual_blocks = nn.Sequential( + *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)] + ) + + def forward(self, x): + # Reshape to merge the frame dimension into batch + bs, c, f, h, w = x.size() + x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w) + + # Pixel Unshuffle operation + x_unshuffled = self.pixel_unshuffle(x) + + # Convolution operation + x_conv = self.conv(x_unshuffled) + + # Feature extraction with residual blocks + out = self.residual_blocks(x_conv) + + # Reshape to restore original bf dimension + out = out.view(bs, f, out.size(1), out.size(2), out.size(3)) + + # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames + out = out.permute(0, 2, 1, 3, 4) + + return out + + def process_camera_coordinates( + self, + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + height: int, + width: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + ): + if origin is None: + origin = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + coordinates = generate_camera_coordinates(direction, length, speed, origin) + plucker_embedding = process_pose_file(coordinates, width, height) + return plucker_embedding + + + +class ResidualBlock(nn.Module): + def __init__(self, dim): + super(ResidualBlock, self).__init__() + self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + + def forward(self, x): + residual = x + out = self.relu(self.conv1(x)) + out = self.conv2(out) + out += residual + return out + +class Camera(object): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + def __init__(self, entry): + fx, fy, cx, cy = entry[1:5] + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + w2c_mat = np.array(entry[7:]).reshape(3, 4) + w2c_mat_4x4 = np.eye(4) + w2c_mat_4x4[:3, :] = w2c_mat + self.w2c_mat = w2c_mat_4x4 + self.c2w_mat = np.linalg.inv(w2c_mat_4x4) + +def get_relative_pose(cam_params): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params] + abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params] + cam_to_origin = 0 + target_cam_c2w = np.array([ + [1, 0, 0, 0], + [0, 1, 0, -cam_to_origin], + [0, 0, 1, 0], + [0, 0, 0, 1] + ]) + abs2rel = target_cam_c2w @ abs_w2cs[0] + ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]] + ret_poses = np.array(ret_poses, dtype=np.float32) + return ret_poses + +def custom_meshgrid(*args): + # torch>=2.0.0 only + return torch.meshgrid(*args, indexing='ij') + + +def ray_condition(K, c2w, H, W, device): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + # c2w: B, V, 4, 4 + # K: B, V, 4 + + B = K.shape[0] + + j, i = custom_meshgrid( + torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype), + torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype), + ) + i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + + fx, fy, cx, cy = K.chunk(4, dim=-1) # B,V, 1 + + zs = torch.ones_like(i) # [B, HxW] + xs = (i - cx) / fx * zs + ys = (j - cy) / fy * zs + zs = zs.expand_as(ys) + + directions = torch.stack((xs, ys, zs), dim=-1) # B, V, HW, 3 + directions = directions / directions.norm(dim=-1, keepdim=True) # B, V, HW, 3 + + rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2) # B, V, 3, HW + rays_o = c2w[..., :3, 3] # B, V, 3 + rays_o = rays_o[:, :, None].expand_as(rays_d) # B, V, 3, HW + # c2w @ dirctions + rays_dxo = torch.linalg.cross(rays_o, rays_d) + plucker = torch.cat([rays_dxo, rays_d], dim=-1) + plucker = plucker.reshape(B, c2w.shape[1], H, W, 6) # B, V, H, W, 6 + # plucker = plucker.permute(0, 1, 4, 2, 3) + return plucker + + +def process_pose_file(cam_params, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False): + if return_poses: + return cam_params + else: + cam_params = [Camera(cam_param) for cam_param in cam_params] + + sample_wh_ratio = width / height + pose_wh_ratio = original_pose_width / original_pose_height # Assuming placeholder ratios, change as needed + + if pose_wh_ratio > sample_wh_ratio: + resized_ori_w = height * pose_wh_ratio + for cam_param in cam_params: + cam_param.fx = resized_ori_w * cam_param.fx / width + else: + resized_ori_h = width / pose_wh_ratio + for cam_param in cam_params: + cam_param.fy = resized_ori_h * cam_param.fy / height + + intrinsic = np.asarray([[cam_param.fx * width, + cam_param.fy * height, + cam_param.cx * width, + cam_param.cy * height] + for cam_param in cam_params], dtype=np.float32) + + K = torch.as_tensor(intrinsic)[None] # [1, 1, 4] + c2ws = get_relative_pose(cam_params) # Assuming this function is defined elsewhere + c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4] + plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W + plucker_embedding = plucker_embedding[None] + plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0] + return plucker_embedding + + + +def generate_camera_coordinates( + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown", "In", "Out"], + length: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) +): + coordinates = [list(origin)] + while len(coordinates) < length: + coor = coordinates[-1].copy() + if "Left" in direction: + coor[9] += speed + if "Right" in direction: + coor[9] -= speed + if "Up" in direction: + coor[13] += speed + if "Down" in direction: + coor[13] -= speed + if "In" in direction: + coor[18] -= speed + if "Out" in direction: + coor[18] += speed + coordinates.append(coor) + return coordinates diff --git a/diffsynth/models/wan_video_dit.py b/diffsynth/models/wan_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..daafa7a6876a34c645a36b7d13ac582455eb6603 --- /dev/null +++ b/diffsynth/models/wan_video_dit.py @@ -0,0 +1,406 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from typing import Tuple, Optional +from einops import rearrange +from .wan_video_camera_controller import SimpleAdapter +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + + +def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False): + if compatibility_mode: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_3_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x,tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_2_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn.flash_attn_func(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif SAGE_ATTN_AVAILABLE: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = sageattn(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + else: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return (x * (1 + scale) + shift) + + +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer(position.type(torch.float64), torch.pow( + 10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + + +def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + return f_freqs_cis, h_freqs_cis, w_freqs_cis + + +def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2) + [: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(end, device=freqs.device), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + +def rope_apply(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward(self, x): + dtype = x.dtype + return self.norm(x.float()).to(dtype) * self.weight + + +class AttentionModule(nn.Module): + def __init__(self, num_heads): + super().__init__() + self.num_heads = num_heads + + def forward(self, q, k, v): + x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads) + return x + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + x = self.attn(q, k, v) + return self.o(x) + + +class CrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + self.has_image_input = has_image_input + if has_image_input: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + if self.has_image_input: + img = y[:, :257] + ctx = y[:, 257:] + else: + ctx = y + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(ctx)) + v = self.v(ctx) + x = self.attn(q, k, v) + if self.has_image_input: + k_img = self.norm_k_img(self.k_img(img)) + v_img = self.v_img(img) + y = flash_attention(q, k_img, v_img, num_heads=self.num_heads) + x = x + y + return self.o(x) + + +class GateModule(nn.Module): + def __init__(self,): + super().__init__() + + def forward(self, x, gate, residual): + return x + gate * residual + +class DiTBlock(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention( + dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( + approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), + shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), + ) + input_x = modulate(self.norm1(x), shift_msa, scale_msa) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x), context) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class MLP(torch.nn.Module): + def __init__(self, in_dim, out_dim, has_pos_emb=False): + super().__init__() + self.proj = torch.nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU(), + nn.Linear(in_dim, out_dim), + nn.LayerNorm(out_dim) + ) + self.has_pos_emb = has_pos_emb + if has_pos_emb: + self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) + + def forward(self, x): + if self.has_pos_emb: + x = x + self.emb_pos.to(dtype=x.dtype, device=x.device) + return self.proj(x) + + +class Head(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + if len(t_mod.shape) == 3: + shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2) + x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))) + else: + shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + scale) + shift)) + return x + + +class WanModel(torch.nn.Module): + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: Tuple[int, int, int], + num_heads: int, + num_layers: int, + has_image_input: bool, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + seperated_timestep: bool = False, + require_vae_embedding: bool = True, + require_clip_embedding: bool = True, + fuse_vae_embedding_in_latents: bool = False, + ): + super().__init__() + self.dim = dim + self.in_dim = in_dim + self.freq_dim = freq_dim + self.has_image_input = has_image_input + self.patch_size = patch_size + self.seperated_timestep = seperated_timestep + self.require_vae_embedding = require_vae_embedding + self.require_clip_embedding = require_clip_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) + for _ in range(num_layers) + ]) + self.head = Head(dim, out_dim, patch_size, eps) + head_dim = dim // num_heads + self.freqs = precompute_freqs_cis_3d(head_dim) + + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) # clip_feature_dim = 1280 + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + if add_control_adapter: + self.control_adapter = SimpleAdapter(in_dim_control_adapter, dim, kernel_size=patch_size[1:], stride=patch_size[1:]) + else: + self.control_adapter = None + + def patchify(self, x: torch.Tensor, control_camera_latents_input: Optional[torch.Tensor] = None): + x = self.patch_embedding(x) + if self.control_adapter is not None and control_camera_latents_input is not None: + y_camera = self.control_adapter(control_camera_latents_input) + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + return x + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], h=grid_size[1], w=grid_size[2], + x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] + ) + + def forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep).to(x.dtype)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + + x, (f, h, w) = self.patchify(x) + + freqs = torch.cat([ + self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block in self.blocks: + if self.training and use_gradient_checkpointing: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + x = self.head(x, t) + x = self.unpatchify(x, (f, h, w)) + return x diff --git a/diffsynth/models/wan_video_dit_s2v.py b/diffsynth/models/wan_video_dit_s2v.py new file mode 100644 index 0000000000000000000000000000000000000000..8fbed8c05a2b5acca4c1ecaa2a68100433c74366 --- /dev/null +++ b/diffsynth/models/wan_video_dit_s2v.py @@ -0,0 +1,594 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Tuple +from .wan_video_dit import rearrange, precompute_freqs_cis_3d, DiTBlock, Head, CrossAttention, modulate, sinusoidal_embedding_1d + + +def torch_dfs(model: nn.Module, parent_name='root'): + module_names, modules = [], [] + current_name = parent_name if parent_name else 'root' + module_names.append(current_name) + modules.append(model) + + for name, child in model.named_children(): + if parent_name: + child_name = f'{parent_name}.{name}' + else: + child_name = name + child_modules, child_names = torch_dfs(child, child_name) + module_names += child_names + modules += child_modules + return modules, module_names + + +def rope_precompute(x, grid_sizes, freqs, start=None): + b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2 + + # split freqs + if type(freqs) is list: + trainable_freqs = freqs[1] + freqs = freqs[0] + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = torch.view_as_complex(x.detach().reshape(b, s, n, -1, 2).to(torch.float64)) + seq_bucket = [0] + if not type(grid_sizes) is list: + grid_sizes = [grid_sizes] + for g in grid_sizes: + if not type(g) is list: + g = [torch.zeros_like(g), g] + batch_size = g[0].shape[0] + for i in range(batch_size): + if start is None: + f_o, h_o, w_o = g[0][i] + else: + f_o, h_o, w_o = start[i] + + f, h, w = g[1][i] + t_f, t_h, t_w = g[2][i] + seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o + seq_len = int(seq_f * seq_h * seq_w) + if seq_len > 0: + if t_f > 0: + factor_f, factor_h, factor_w = (t_f / seq_f).item(), (t_h / seq_h).item(), (t_w / seq_w).item() + # Generate a list of seq_f integers starting from f_o and ending at math.ceil(factor_f * seq_f.item() + f_o.item()) + if f_o >= 0: + f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist() + else: + f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist() + h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist() + w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist() + + assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 + freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = freqs_0.view(seq_f, 1, 1, -1) + + freqs_i = torch.cat( + [ + freqs_0.expand(seq_f, seq_h, seq_w, -1), + freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1), + freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1), + ], + dim=-1 + ).reshape(seq_len, 1, -1) + elif t_f < 0: + freqs_i = trainable_freqs.unsqueeze(1) + # apply rotary embedding + output[i, seq_bucket[-1]:seq_bucket[-1] + seq_len] = freqs_i + seq_bucket.append(seq_bucket[-1] + seq_len) + return output + + +class CausalConv1d(nn.Module): + + def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode='replicate', **kwargs): + super().__init__() + + self.pad_mode = pad_mode + padding = (kernel_size - 1, 0) # T + self.time_causal_padding = padding + + self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + +class MotionEncoder_tc(nn.Module): + + def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, need_global=True, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.num_heads = num_heads + self.need_global = need_global + self.conv1_local = CausalConv1d(in_dim, hidden_dim // 4 * num_heads, 3, stride=1) + if need_global: + self.conv1_global = CausalConv1d(in_dim, hidden_dim // 4, 3, stride=1) + self.norm1 = nn.LayerNorm(hidden_dim // 4, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.act = nn.SiLU() + self.conv2 = CausalConv1d(hidden_dim // 4, hidden_dim // 2, 3, stride=2) + self.conv3 = CausalConv1d(hidden_dim // 2, hidden_dim, 3, stride=2) + + if need_global: + self.final_linear = nn.Linear(hidden_dim, hidden_dim, **factory_kwargs) + + self.norm1 = nn.LayerNorm(hidden_dim // 4, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.norm2 = nn.LayerNorm(hidden_dim // 2, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.norm3 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim)) + + def forward(self, x): + x = rearrange(x, 'b t c -> b c t') + x_ori = x.clone() + b, c, t = x.shape + x = self.conv1_local(x) + x = rearrange(x, 'b (n c) t -> (b n) t c', n=self.num_heads) + x = self.norm1(x) + x = self.act(x) + x = rearrange(x, 'b t c -> b c t') + x = self.conv2(x) + x = rearrange(x, 'b c t -> b t c') + x = self.norm2(x) + x = self.act(x) + x = rearrange(x, 'b t c -> b c t') + x = self.conv3(x) + x = rearrange(x, 'b c t -> b t c') + x = self.norm3(x) + x = self.act(x) + x = rearrange(x, '(b n) t c -> b t n c', b=b) + padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1).to(device=x.device, dtype=x.dtype) + x = torch.cat([x, padding], dim=-2) + x_local = x.clone() + + if not self.need_global: + return x_local + + x = self.conv1_global(x_ori) + x = rearrange(x, 'b c t -> b t c') + x = self.norm1(x) + x = self.act(x) + x = rearrange(x, 'b t c -> b c t') + x = self.conv2(x) + x = rearrange(x, 'b c t -> b t c') + x = self.norm2(x) + x = self.act(x) + x = rearrange(x, 'b t c -> b c t') + x = self.conv3(x) + x = rearrange(x, 'b c t -> b t c') + x = self.norm3(x) + x = self.act(x) + x = self.final_linear(x) + x = rearrange(x, '(b n) t c -> b t n c', b=b) + + return x, x_local + + +class FramePackMotioner(nn.Module): + + def __init__(self, inner_dim=1024, num_heads=16, zip_frame_buckets=[1, 2, 16], drop_mode="drop", *args, **kwargs): + super().__init__(*args, **kwargs) + self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) + self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4)) + self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8)) + self.zip_frame_buckets = torch.tensor(zip_frame_buckets, dtype=torch.long) + + self.inner_dim = inner_dim + self.num_heads = num_heads + self.freqs = torch.cat(precompute_freqs_cis_3d(inner_dim // num_heads), dim=1) + self.drop_mode = drop_mode + + def forward(self, motion_latents, add_last_motion=2): + motion_frames = motion_latents[0].shape[1] + mot = [] + mot_remb = [] + for m in motion_latents: + lat_height, lat_width = m.shape[2], m.shape[3] + padd_lat = torch.zeros(16, self.zip_frame_buckets.sum(), lat_height, lat_width).to(device=m.device, dtype=m.dtype) + overlap_frame = min(padd_lat.shape[1], m.shape[1]) + if overlap_frame > 0: + padd_lat[:, -overlap_frame:] = m[:, -overlap_frame:] + + if add_last_motion < 2 and self.drop_mode != "drop": + zero_end_frame = self.zip_frame_buckets[:self.zip_frame_buckets.__len__() - add_last_motion - 1].sum() + padd_lat[:, -zero_end_frame:] = 0 + + padd_lat = padd_lat.unsqueeze(0) + clean_latents_4x, clean_latents_2x, clean_latents_post = padd_lat[:, :, -self.zip_frame_buckets.sum():, :, :].split( + list(self.zip_frame_buckets)[::-1], dim=2 + ) # 16, 2 ,1 + + # patchfy + clean_latents_post = self.proj(clean_latents_post).flatten(2).transpose(1, 2) + clean_latents_2x = self.proj_2x(clean_latents_2x).flatten(2).transpose(1, 2) + clean_latents_4x = self.proj_4x(clean_latents_4x).flatten(2).transpose(1, 2) + + if add_last_motion < 2 and self.drop_mode == "drop": + clean_latents_post = clean_latents_post[:, :0] if add_last_motion < 2 else clean_latents_post + clean_latents_2x = clean_latents_2x[:, :0] if add_last_motion < 1 else clean_latents_2x + + motion_lat = torch.cat([clean_latents_post, clean_latents_2x, clean_latents_4x], dim=1) + + # rope + start_time_id = -(self.zip_frame_buckets[:1].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[0] + grid_sizes = [] if add_last_motion < 2 and self.drop_mode == "drop" else \ + [ + [torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[0], lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), ] + ] + + start_time_id = -(self.zip_frame_buckets[:2].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[1] // 2 + grid_sizes_2x = [] if add_last_motion < 1 and self.drop_mode == "drop" else \ + [ + [torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 4, lat_width // 4]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[1], lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), ] + ] + + start_time_id = -(self.zip_frame_buckets[:3].sum()) + end_time_id = start_time_id + self.zip_frame_buckets[2] // 4 + grid_sizes_4x = [ + [ + torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1), + torch.tensor([end_time_id, lat_height // 8, lat_width // 8]).unsqueeze(0).repeat(1, 1), + torch.tensor([self.zip_frame_buckets[2], lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), + ] + ] + + grid_sizes = grid_sizes + grid_sizes_2x + grid_sizes_4x + + motion_rope_emb = rope_precompute( + motion_lat.detach().view(1, motion_lat.shape[1], self.num_heads, self.inner_dim // self.num_heads), + grid_sizes, + self.freqs, + start=None + ) + + mot.append(motion_lat) + mot_remb.append(motion_rope_emb) + return mot, mot_remb + + +class AdaLayerNorm(nn.Module): + + def __init__( + self, + embedding_dim: int, + output_dim: int, + norm_eps: float = 1e-5, + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, output_dim) + self.norm = nn.LayerNorm(output_dim // 2, norm_eps, elementwise_affine=False) + + def forward(self, x, temb): + temb = self.linear(F.silu(temb)) + shift, scale = temb.chunk(2, dim=1) + shift = shift[:, None, :] + scale = scale[:, None, :] + x = self.norm(x) * (1 + scale) + shift + return x + + +class AudioInjector_WAN(nn.Module): + + def __init__( + self, + all_modules, + all_modules_names, + dim=2048, + num_heads=32, + inject_layer=[0, 27], + enable_adain=False, + adain_dim=2048, + ): + super().__init__() + self.injected_block_id = {} + audio_injector_id = 0 + for mod_name, mod in zip(all_modules_names, all_modules): + if isinstance(mod, DiTBlock): + for inject_id in inject_layer: + if f'transformer_blocks.{inject_id}' in mod_name: + self.injected_block_id[inject_id] = audio_injector_id + audio_injector_id += 1 + + self.injector = nn.ModuleList([CrossAttention( + dim=dim, + num_heads=num_heads, + ) for _ in range(audio_injector_id)]) + self.injector_pre_norm_feat = nn.ModuleList([nn.LayerNorm( + dim, + elementwise_affine=False, + eps=1e-6, + ) for _ in range(audio_injector_id)]) + self.injector_pre_norm_vec = nn.ModuleList([nn.LayerNorm( + dim, + elementwise_affine=False, + eps=1e-6, + ) for _ in range(audio_injector_id)]) + if enable_adain: + self.injector_adain_layers = nn.ModuleList([AdaLayerNorm(output_dim=dim * 2, embedding_dim=adain_dim) for _ in range(audio_injector_id)]) + + +class CausalAudioEncoder(nn.Module): + + def __init__(self, dim=5120, num_layers=25, out_dim=2048, num_token=4, need_global=False): + super().__init__() + self.encoder = MotionEncoder_tc(in_dim=dim, hidden_dim=out_dim, num_heads=num_token, need_global=need_global) + weight = torch.ones((1, num_layers, 1, 1)) * 0.01 + + self.weights = torch.nn.Parameter(weight) + self.act = torch.nn.SiLU() + + def forward(self, features): + # features B * num_layers * dim * video_length + weights = self.act(self.weights.to(device=features.device, dtype=features.dtype)) + weights_sum = weights.sum(dim=1, keepdims=True) + weighted_feat = ((features * weights) / weights_sum).sum(dim=1) # b dim f + weighted_feat = weighted_feat.permute(0, 2, 1) # b f dim + res = self.encoder(weighted_feat) # b f n dim + return res # b f n dim + + +class WanS2VDiTBlock(DiTBlock): + + def forward(self, x, context, t_mod, seq_len_x, freqs): + t_mod = (self.modulation.unsqueeze(2).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + # t_mod[:, :, 0] for x, t_mod[:, :, 1] for other like ref, motion, etc. + t_mod = [ + torch.cat([element[:, :, 0].expand(1, seq_len_x, x.shape[-1]), element[:, :, 1].expand(1, x.shape[1] - seq_len_x, x.shape[-1])], dim=1) + for element in t_mod + ] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = t_mod + input_x = modulate(self.norm1(x), shift_msa, scale_msa) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x), context) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class WanS2VModel(torch.nn.Module): + + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: Tuple[int, int, int], + num_heads: int, + num_layers: int, + cond_dim: int, + audio_dim: int, + num_audio_token: int, + enable_adain: bool = True, + audio_inject_layers: list = [0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39], + zero_timestep: bool = True, + add_last_motion: bool = True, + framepack_drop_mode: str = "padd", + fuse_vae_embedding_in_latents: bool = True, + require_vae_embedding: bool = False, + seperated_timestep: bool = False, + require_clip_embedding: bool = False, + ): + super().__init__() + self.dim = dim + self.in_dim = in_dim + self.freq_dim = freq_dim + self.patch_size = patch_size + self.num_heads = num_heads + self.enbale_adain = enable_adain + self.add_last_motion = add_last_motion + self.zero_timestep = zero_timestep + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + self.require_vae_embedding = require_vae_embedding + self.seperated_timestep = seperated_timestep + self.require_clip_embedding = require_clip_embedding + + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), nn.Linear(dim, dim)) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + self.blocks = nn.ModuleList([WanS2VDiTBlock(False, dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]) + self.head = Head(dim, out_dim, patch_size, eps) + self.freqs = torch.cat(precompute_freqs_cis_3d(dim // num_heads), dim=1) + + self.cond_encoder = nn.Conv3d(cond_dim, dim, kernel_size=patch_size, stride=patch_size) + self.casual_audio_encoder = CausalAudioEncoder(dim=audio_dim, out_dim=dim, num_token=num_audio_token, need_global=enable_adain) + all_modules, all_modules_names = torch_dfs(self.blocks, parent_name="root.transformer_blocks") + self.audio_injector = AudioInjector_WAN( + all_modules, + all_modules_names, + dim=dim, + num_heads=num_heads, + inject_layer=audio_inject_layers, + enable_adain=enable_adain, + adain_dim=dim, + ) + self.trainable_cond_mask = nn.Embedding(3, dim) + self.frame_packer = FramePackMotioner(inner_dim=dim, num_heads=num_heads, zip_frame_buckets=[1, 2, 16], drop_mode=framepack_drop_mode) + + def patchify(self, x: torch.Tensor): + grid_size = x.shape[2:] + x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous() + return x, grid_size # x, grid_size: (f, h, w) + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, + 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], + h=grid_size[1], + w=grid_size[2], + x=self.patch_size[0], + y=self.patch_size[1], + z=self.patch_size[2] + ) + + def process_motion_frame_pack(self, motion_latents, drop_motion_frames=False, add_last_motion=2): + flattern_mot, mot_remb = self.frame_packer(motion_latents, add_last_motion) + if drop_motion_frames: + return [m[:, :0] for m in flattern_mot], [m[:, :0] for m in mot_remb] + else: + return flattern_mot, mot_remb + + def inject_motion(self, x, rope_embs, mask_input, motion_latents, drop_motion_frames=True, add_last_motion=2): + # inject the motion frames token to the hidden states + mot, mot_remb = self.process_motion_frame_pack(motion_latents, drop_motion_frames=drop_motion_frames, add_last_motion=add_last_motion) + if len(mot) > 0: + x = torch.cat([x, mot[0]], dim=1) + rope_embs = torch.cat([rope_embs, mot_remb[0]], dim=1) + mask_input = torch.cat( + [mask_input, 2 * torch.ones([1, x.shape[1] - mask_input.shape[1]], device=mask_input.device, dtype=mask_input.dtype)], dim=1 + ) + return x, rope_embs, mask_input + + def after_transformer_block(self, block_idx, hidden_states, audio_emb_global, audio_emb, original_seq_len, use_unified_sequence_parallel=False): + if block_idx in self.audio_injector.injected_block_id.keys(): + audio_attn_id = self.audio_injector.injected_block_id[block_idx] + num_frames = audio_emb.shape[1] + if use_unified_sequence_parallel: + from xfuser.core.distributed import get_sp_group + hidden_states = get_sp_group().all_gather(hidden_states, dim=1) + + input_hidden_states = hidden_states[:, :original_seq_len].clone() # b (f h w) c + input_hidden_states = rearrange(input_hidden_states, "b (t n) c -> (b t) n c", t=num_frames) + + audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c") + adain_hidden_states = self.audio_injector.injector_adain_layers[audio_attn_id](input_hidden_states, temb=audio_emb_global[:, 0]) + attn_hidden_states = adain_hidden_states + + audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames) + attn_audio_emb = audio_emb + residual_out = self.audio_injector.injector[audio_attn_id](attn_hidden_states, attn_audio_emb) + residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) + hidden_states[:, :original_seq_len] = hidden_states[:, :original_seq_len] + residual_out + if use_unified_sequence_parallel: + from xfuser.core.distributed import get_sequence_parallel_world_size, get_sequence_parallel_rank + hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()] + return hidden_states + + def cal_audio_emb(self, audio_input, motion_frames=[73, 19]): + audio_input = torch.cat([audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], dim=-1) + audio_emb_global, audio_emb = self.casual_audio_encoder(audio_input) + audio_emb_global = audio_emb_global[:, motion_frames[1]:].clone() + merged_audio_emb = audio_emb[:, motion_frames[1]:, :] + return audio_emb_global, merged_audio_emb + + def get_grid_sizes(self, grid_size_x, grid_size_ref): + f, h, w = grid_size_x + rf, rh, rw = grid_size_ref + grid_sizes_x = torch.tensor([f, h, w], dtype=torch.long).unsqueeze(0) + grid_sizes_x = [[torch.zeros_like(grid_sizes_x), grid_sizes_x, grid_sizes_x]] + grid_sizes_ref = [[ + torch.tensor([30, 0, 0]).unsqueeze(0), + torch.tensor([31, rh, rw]).unsqueeze(0), + torch.tensor([1, rh, rw]).unsqueeze(0), + ]] + return grid_sizes_x + grid_sizes_ref + + def forward( + self, + latents, + timestep, + context, + audio_input, + motion_latents, + pose_cond, + use_gradient_checkpointing_offload=False, + use_gradient_checkpointing=False + ): + origin_ref_latents = latents[:, :, 0:1] + x = latents[:, :, 1:] + + # context embedding + context = self.text_embedding(context) + + # audio encode + audio_emb_global, merged_audio_emb = self.cal_audio_emb(audio_input) + + # x and pose_cond + pose_cond = torch.zeros_like(x) if pose_cond is None else pose_cond + x, (f, h, w) = self.patchify(self.patch_embedding(x) + self.cond_encoder(pose_cond)) # torch.Size([1, 29120, 5120]) + seq_len_x = x.shape[1] + + # reference image + ref_latents, (rf, rh, rw) = self.patchify(self.patch_embedding(origin_ref_latents)) # torch.Size([1, 1456, 5120]) + grid_sizes = self.get_grid_sizes((f, h, w), (rf, rh, rw)) + x = torch.cat([x, ref_latents], dim=1) + # mask + mask = torch.cat([torch.zeros([1, seq_len_x]), torch.ones([1, ref_latents.shape[1]])], dim=1).to(torch.long).to(x.device) + # freqs + pre_compute_freqs = rope_precompute( + x.detach().view(1, x.size(1), self.num_heads, self.dim // self.num_heads), grid_sizes, self.freqs, start=None + ) + # motion + x, pre_compute_freqs, mask = self.inject_motion(x, pre_compute_freqs, mask, motion_latents, add_last_motion=2) + + x = x + self.trainable_cond_mask(mask).to(x.dtype) + + # t_mod + timestep = torch.cat([timestep, torch.zeros([1], dtype=timestep.dtype, device=timestep.device)]) + t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)).unsqueeze(2).transpose(0, 2) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block_id, block in enumerate(self.blocks): + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + seq_len_x, + pre_compute_freqs[0], + use_reentrant=False, + ) + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)), + x, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + seq_len_x, + pre_compute_freqs[0], + use_reentrant=False, + ) + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)), + x, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, seq_len_x, pre_compute_freqs[0]) + x = self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x) + + x = x[:, :seq_len_x] + x = self.head(x, t[:-1]) + x = self.unpatchify(x, (f, h, w)) + # make compatible with wan video + x = torch.cat([origin_ref_latents, x], dim=2) + return x diff --git a/diffsynth/models/wan_video_image_encoder.py b/diffsynth/models/wan_video_image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..37d17d6a183f5d6290c3f4b1e417bf0310bfa353 --- /dev/null +++ b/diffsynth/models/wan_video_image_encoder.py @@ -0,0 +1,878 @@ +""" +Concise re-implementation of +``https://github.com/openai/CLIP'' and +``https://github.com/mlfoundations/open_clip''. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as T +from .wan_video_dit import flash_attention + + +class SelfAttention(nn.Module): + + def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, mask): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + + # compute attention + p = self.dropout.p if self.training else 0.0 + x = F.scaled_dot_product_attention(q, k, v, mask, p) + x = x.permute(0, 2, 1, 3).reshape(b, s, c) + + # output + x = self.o(x) + x = self.dropout(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.post_norm = post_norm + self.eps = eps + + # layers + self.attn = SelfAttention(dim, num_heads, dropout, eps) + self.norm1 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), + nn.Dropout(dropout)) + self.norm2 = nn.LayerNorm(dim, eps=eps) + + def forward(self, x, mask): + if self.post_norm: + x = self.norm1(x + self.attn(x, mask)) + x = self.norm2(x + self.ffn(x)) + else: + x = x + self.attn(self.norm1(x), mask) + x = x + self.ffn(self.norm2(x)) + return x + + +class XLMRoberta(nn.Module): + """ + XLMRobertaModel with no pooler and no LM head. + """ + + def __init__(self, + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5): + super().__init__() + self.vocab_size = vocab_size + self.max_seq_len = max_seq_len + self.type_size = type_size + self.pad_id = pad_id + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + self.post_norm = post_norm + self.eps = eps + + # embeddings + self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id) + self.type_embedding = nn.Embedding(type_size, dim) + self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id) + self.dropout = nn.Dropout(dropout) + + # blocks + self.blocks = nn.ModuleList([ + AttentionBlock(dim, num_heads, post_norm, dropout, eps) + for _ in range(num_layers) + ]) + + # norm layer + self.norm = nn.LayerNorm(dim, eps=eps) + + def forward(self, ids): + """ + ids: [B, L] of torch.LongTensor. + """ + b, s = ids.shape + mask = ids.ne(self.pad_id).long() + + # embeddings + x = self.token_embedding(ids) + \ + self.type_embedding(torch.zeros_like(ids)) + \ + self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask) + if self.post_norm: + x = self.norm(x) + x = self.dropout(x) + + # blocks + mask = torch.where( + mask.view(b, 1, 1, s).gt(0), 0.0, + torch.finfo(x.dtype).min) + for block in self.blocks: + x = block(x, mask) + + # output + if not self.post_norm: + x = self.norm(x) + return x + + +def xlm_roberta_large(pretrained=False, + return_tokenizer=False, + device='cpu', + **kwargs): + """ + XLMRobertaLarge adapted from Huggingface. + """ + # params + cfg = dict( + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5) + cfg.update(**kwargs) + + # init model + if pretrained: + from sora import DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = XLMRoberta(**cfg) + + # load checkpoint + model.load_state_dict( + torch.load( + DOWNLOAD_TO_CACHE('models/xlm_roberta/xlm_roberta_large.pth'), + map_location=device), + assign=True) + else: + # init a model on device + with torch.device(device): + model = XLMRoberta(**cfg) + + # init tokenizer + if return_tokenizer: + from sora.data import HuggingfaceTokenizer + tokenizer = HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.text_len, + clean='whitespace') + return model, tokenizer + else: + return model + + + +def pos_interpolate(pos, seq_len): + if pos.size(1) == seq_len: + return pos + else: + src_grid = int(math.sqrt(pos.size(1))) + tar_grid = int(math.sqrt(seq_len)) + n = pos.size(1) - src_grid * src_grid + return torch.cat([ + pos[:, :n], + F.interpolate( + pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute( + 0, 3, 1, 2), + size=(tar_grid, tar_grid), + mode='bicubic', + align_corners=False).flatten(2).transpose(1, 2) + ], + dim=1) + + +class QuickGELU(nn.Module): + + def forward(self, x): + return x * torch.sigmoid(1.702 * x) + + +class LayerNorm(nn.LayerNorm): + + def forward(self, x): + return super().forward(x).type_as(x) + + +class SelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + causal=False, + attn_dropout=0.0, + proj_dropout=0.0): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.causal = causal + self.attn_dropout = attn_dropout + self.proj_dropout = proj_dropout + + # layers + self.to_qkv = nn.Linear(dim, dim * 3) + self.proj = nn.Linear(dim, dim) + + def forward(self, x): + """ + x: [B, L, C]. + """ + # compute query, key, value + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + return x + + +class SwiGLU(nn.Module): + + def __init__(self, dim, mid_dim): + super().__init__() + self.dim = dim + self.mid_dim = mid_dim + + # layers + self.fc1 = nn.Linear(dim, mid_dim) + self.fc2 = nn.Linear(dim, mid_dim) + self.fc3 = nn.Linear(mid_dim, dim) + + def forward(self, x): + x = F.silu(self.fc1(x)) * self.fc2(x) + x = self.fc3(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + post_norm=False, + causal=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + norm_eps=1e-5): + assert activation in ['quick_gelu', 'gelu', 'swi_glu'] + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.post_norm = post_norm + self.causal = causal + self.norm_eps = norm_eps + + # layers + self.norm1 = LayerNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, + proj_dropout) + self.norm2 = LayerNorm(dim, eps=norm_eps) + if activation == 'swi_glu': + self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) + else: + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + if self.post_norm: + x = x + self.norm1(self.attn(x)) + x = x + self.norm2(self.mlp(x)) + else: + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class AttentionPool(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + activation='gelu', + proj_dropout=0.0, + norm_eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.proj_dropout = proj_dropout + self.norm_eps = norm_eps + + # layers + gain = 1.0 / math.sqrt(dim) + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.to_q = nn.Linear(dim, dim) + self.to_kv = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + self.norm = LayerNorm(dim, eps=norm_eps) + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.to_q(self.cls_embedding).view(1, 1, n*d).expand(b, -1, -1) + k, v = self.to_kv(x).chunk(2, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + x = x.reshape(b, 1, c) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + + # mlp + x = x + self.mlp(self.norm(x)) + return x[:, 0] + + +class VisionTransformer(nn.Module): + + def __init__(self, + image_size=224, + patch_size=16, + dim=768, + mlp_ratio=4, + out_dim=512, + num_heads=12, + num_layers=12, + pool_type='token', + pre_norm=True, + post_norm=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + if image_size % patch_size != 0: + print( + '[WARNING] image_size is not divisible by patch_size', + flush=True) + assert pool_type in ('token', 'token_fc', 'attn_pool') + out_dim = out_dim or dim + super().__init__() + self.image_size = image_size + self.patch_size = patch_size + self.num_patches = (image_size // patch_size)**2 + self.dim = dim + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.pool_type = pool_type + self.post_norm = post_norm + self.norm_eps = norm_eps + + # embeddings + gain = 1.0 / math.sqrt(dim) + self.patch_embedding = nn.Conv2d( + 3, + dim, + kernel_size=patch_size, + stride=patch_size, + bias=not pre_norm) + if pool_type in ('token', 'token_fc'): + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.pos_embedding = nn.Parameter(gain * torch.randn( + 1, self.num_patches + + (1 if pool_type in ('token', 'token_fc') else 0), dim)) + self.dropout = nn.Dropout(embedding_dropout) + + # transformer + self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None + self.transformer = nn.Sequential(*[ + AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, + activation, attn_dropout, proj_dropout, norm_eps) + for _ in range(num_layers) + ]) + self.post_norm = LayerNorm(dim, eps=norm_eps) + + # head + if pool_type == 'token': + self.head = nn.Parameter(gain * torch.randn(dim, out_dim)) + elif pool_type == 'token_fc': + self.head = nn.Linear(dim, out_dim) + elif pool_type == 'attn_pool': + self.head = AttentionPool(dim, mlp_ratio, num_heads, activation, + proj_dropout, norm_eps) + + def forward(self, x, interpolation=False, use_31_block=False): + b = x.size(0) + + # embeddings + x = self.patch_embedding(x).flatten(2).permute(0, 2, 1) + if self.pool_type in ('token', 'token_fc'): + x = torch.cat([self.cls_embedding.expand(b, -1, -1).to(dtype=x.dtype, device=x.device), x], dim=1) + if interpolation: + e = pos_interpolate(self.pos_embedding, x.size(1)) + else: + e = self.pos_embedding + e = e.to(dtype=x.dtype, device=x.device) + x = self.dropout(x + e) + if self.pre_norm is not None: + x = self.pre_norm(x) + + # transformer + if use_31_block: + x = self.transformer[:-1](x) + return x + else: + x = self.transformer(x) + return x + + +class CLIP(nn.Module): + + def __init__(self, + embed_dim=512, + image_size=224, + patch_size=16, + vision_dim=768, + vision_mlp_ratio=4, + vision_heads=12, + vision_layers=12, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + vocab_size=49408, + text_len=77, + text_dim=512, + text_mlp_ratio=4, + text_heads=8, + text_layers=12, + text_causal=True, + text_pool='argmax', + text_head_bias=False, + logit_bias=None, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pool = vision_pool + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.vocab_size = vocab_size + self.text_len = text_len + self.text_dim = text_dim + self.text_mlp_ratio = text_mlp_ratio + self.text_heads = text_heads + self.text_layers = text_layers + self.text_causal = text_causal + self.text_pool = text_pool + self.text_head_bias = text_head_bias + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = TextTransformer( + vocab_size=vocab_size, + text_len=text_len, + dim=text_dim, + mlp_ratio=text_mlp_ratio, + out_dim=embed_dim, + num_heads=text_heads, + num_layers=text_layers, + causal=text_causal, + pool_type=text_pool, + head_bias=text_head_bias, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + if logit_bias is not None: + self.logit_bias = nn.Parameter(logit_bias * torch.ones([])) + + # initialize weights + self.init_weights() + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def init_weights(self): + # embeddings + nn.init.normal_(self.textual.token_embedding.weight, std=0.02) + nn.init.normal_(self.visual.patch_embedding.weight, std=0.1) + + # attentions + for modality in ['visual', 'textual']: + dim = self.vision_dim if modality == 'visual' else self.text_dim + transformer = getattr(self, modality).transformer + proj_gain = (1.0 / math.sqrt(dim)) * ( + 1.0 / math.sqrt(2 * len(transformer))) + attn_gain = 1.0 / math.sqrt(dim) + mlp_gain = 1.0 / math.sqrt(2.0 * dim) + for block in transformer: + nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain) + nn.init.normal_(block.attn.proj.weight, std=proj_gain) + nn.init.normal_(block.mlp[0].weight, std=mlp_gain) + nn.init.normal_(block.mlp[2].weight, std=proj_gain) + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +class XLMRobertaWithHead(XLMRoberta): + + def __init__(self, **kwargs): + self.out_dim = kwargs.pop('out_dim') + super().__init__(**kwargs) + + # head + mid_dim = (self.dim + self.out_dim) // 2 + self.head = nn.Sequential( + nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(), + nn.Linear(mid_dim, self.out_dim, bias=False)) + + def forward(self, ids): + # xlm-roberta + x = super().forward(ids) + + # average pooling + mask = ids.ne(self.pad_id).unsqueeze(-1).to(x) + x = (x * mask).sum(dim=1) / mask.sum(dim=1) + + # head + x = self.head(x) + return x + + +class XLMRobertaCLIP(nn.Module): + + def __init__(self, + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.activation = activation + self.vocab_size = vocab_size + self.max_text_len = max_text_len + self.type_size = type_size + self.pad_id = pad_id + self.text_dim = text_dim + self.text_heads = text_heads + self.text_layers = text_layers + self.text_post_norm = text_post_norm + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = None + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. + Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +def _clip(pretrained=False, + pretrained_name=None, + model_cls=CLIP, + return_transforms=False, + return_tokenizer=False, + tokenizer_padding='eos', + dtype=torch.float32, + device='cpu', + **kwargs): + # init model + if pretrained and pretrained_name: + from sora import BUCKET, DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = model_cls(**kwargs) + + # checkpoint path + checkpoint = f'models/clip/{pretrained_name}' + if dtype in (torch.float16, torch.bfloat16): + suffix = '-' + { + torch.float16: 'fp16', + torch.bfloat16: 'bf16' + }[dtype] + if object_exists(BUCKET, f'{checkpoint}{suffix}.pth'): + checkpoint = f'{checkpoint}{suffix}' + checkpoint += '.pth' + + # load + model.load_state_dict( + torch.load(DOWNLOAD_TO_CACHE(checkpoint), map_location=device), + assign=True, + strict=False) + else: + # init a model on device + with torch.device(device): + model = model_cls(**kwargs) + + # set device + output = (model,) + + # init transforms + if return_transforms: + # mean and std + if 'siglip' in pretrained_name.lower(): + mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] + else: + mean = [0.48145466, 0.4578275, 0.40821073] + std = [0.26862954, 0.26130258, 0.27577711] + + # transforms + transforms = T.Compose([ + T.Resize((model.image_size, model.image_size), + interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor(), + T.Normalize(mean=mean, std=std) + ]) + output += (transforms,) + + # init tokenizer + if return_tokenizer: + from sora import data + if 'siglip' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name=f'timm/{pretrained_name}', + seq_len=model.text_len, + clean='canonicalize') + elif 'xlm' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.max_text_len - 2, + clean='whitespace') + elif 'mba' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='facebook/xlm-roberta-xl', + seq_len=model.max_text_len - 2, + clean='whitespace') + else: + tokenizer = data.CLIPTokenizer( + seq_len=model.text_len, padding=tokenizer_padding) + output += (tokenizer,) + return output[0] if len(output) == 1 else output + + +def clip_xlm_roberta_vit_h_14( + pretrained=False, + pretrained_name='open-clip-xlm-roberta-large-vit-huge-14', + **kwargs): + cfg = dict( + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0) + cfg.update(**kwargs) + return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg) + + +class WanImageEncoder(torch.nn.Module): + + def __init__(self): + super().__init__() + # init model + self.model, self.transforms = clip_xlm_roberta_vit_h_14( + pretrained=False, + return_transforms=True, + return_tokenizer=False, + dtype=torch.float32, + device="cpu") + + def encode_image(self, videos): + # preprocess + size = (self.model.image_size,) * 2 + videos = torch.cat([ + F.interpolate( + u, + size=size, + mode='bicubic', + align_corners=False) for u in videos + ]) + videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5)) + + # forward + out = self.model.visual(videos, use_31_block=True) + return out diff --git a/diffsynth/models/wan_video_mot.py b/diffsynth/models/wan_video_mot.py new file mode 100644 index 0000000000000000000000000000000000000000..4091c91777355dce91ccefac56679f8b936e7abb --- /dev/null +++ b/diffsynth/models/wan_video_mot.py @@ -0,0 +1,169 @@ +import torch +from .wan_video_dit import DiTBlock, SelfAttention, rope_apply, flash_attention, modulate, MLP +import einops +import torch.nn as nn + + +class MotSelfAttention(SelfAttention): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__(dim, num_heads, eps) + def forward(self, x, freqs, is_before_attn=False): + if is_before_attn: + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + return q, k, v + else: + return self.o(x) + + +class MotWanAttentionBlock(DiTBlock): + def __init__(self, has_image_input, dim, num_heads, ffn_dim, eps=1e-6, block_id=0): + super().__init__(has_image_input, dim, num_heads, ffn_dim, eps=eps) + self.block_id = block_id + + self.self_attn = MotSelfAttention(dim, num_heads, eps) + + + def forward(self, wan_block, x, context, t_mod, freqs, x_mot, context_mot, t_mod_mot, freqs_mot): + + # 1. prepare scale parameter + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + wan_block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + + scale_params_mot_ref = self.modulation + t_mod_mot.float() + scale_params_mot_ref = einops.rearrange(scale_params_mot_ref, '(b n) t c -> b n t c', n=1) + shift_msa_mot_ref, scale_msa_mot_ref, gate_msa_mot_ref, c_shift_msa_mot_ref, c_scale_msa_mot_ref, c_gate_msa_mot_ref = scale_params_mot_ref.chunk(6, dim=2) + + # 2. Self-attention + input_x = modulate(wan_block.norm1(x), shift_msa, scale_msa) + # original block self-attn + attn1 = wan_block.self_attn + q = attn1.norm_q(attn1.q(input_x)) + k = attn1.norm_k(attn1.k(input_x)) + v = attn1.v(input_x) + q = rope_apply(q, freqs, attn1.num_heads) + k = rope_apply(k, freqs, attn1.num_heads) + + # mot block self-attn + norm_x_mot = einops.rearrange(self.norm1(x_mot.float()), 'b (n t) c -> b n t c', n=1) + norm_x_mot = modulate(norm_x_mot, shift_msa_mot_ref, scale_msa_mot_ref).type_as(x_mot) + norm_x_mot = einops.rearrange(norm_x_mot, 'b n t c -> b (n t) c', n=1) + q_mot,k_mot,v_mot = self.self_attn(norm_x_mot, freqs_mot, is_before_attn=True) + + tmp_hidden_states = flash_attention( + torch.cat([q, q_mot], dim=-2), + torch.cat([k, k_mot], dim=-2), + torch.cat([v, v_mot], dim=-2), + num_heads=attn1.num_heads) + + attn_output, attn_output_mot = torch.split(tmp_hidden_states, [q.shape[-2], q_mot.shape[-2]], dim=-2) + + attn_output = attn1.o(attn_output) + x = wan_block.gate(x, gate_msa, attn_output) + + attn_output_mot = self.self_attn(x=attn_output_mot,freqs=freqs_mot, is_before_attn=False) + # gate + attn_output_mot = einops.rearrange(attn_output_mot, 'b (n t) c -> b n t c', n=1) + attn_output_mot = attn_output_mot * gate_msa_mot_ref + attn_output_mot = einops.rearrange(attn_output_mot, 'b n t c -> b (n t) c', n=1) + x_mot = (x_mot.float() + attn_output_mot).type_as(x_mot) + + # 3. cross-attention and feed-forward + x = x + wan_block.cross_attn(wan_block.norm3(x), context) + input_x = modulate(wan_block.norm2(x), shift_mlp, scale_mlp) + x = wan_block.gate(x, gate_mlp, wan_block.ffn(input_x)) + + x_mot = x_mot + self.cross_attn(self.norm3(x_mot),context_mot) + # modulate + norm_x_mot_ref = einops.rearrange(self.norm2(x_mot.float()), 'b (n t) c -> b n t c', n=1) + norm_x_mot_ref = (norm_x_mot_ref * (1 + c_scale_msa_mot_ref) + c_shift_msa_mot_ref).type_as(x_mot) + norm_x_mot_ref = einops.rearrange(norm_x_mot_ref, 'b n t c -> b (n t) c', n=1) + input_x_mot = self.ffn(norm_x_mot_ref) + # gate + input_x_mot = einops.rearrange(input_x_mot, 'b (n t) c -> b n t c', n=1) + input_x_mot = input_x_mot.float() * c_gate_msa_mot_ref + input_x_mot = einops.rearrange(input_x_mot, 'b n t c -> b (n t) c', n=1) + x_mot = (x_mot.float() + input_x_mot).type_as(x_mot) + + return x, x_mot + + +class MotWanModel(torch.nn.Module): + def __init__( + self, + mot_layers=(0, 4, 8, 12, 16, 20, 24, 28, 32, 36), + patch_size=(1, 2, 2), + has_image_input=True, + has_image_pos_emb=False, + dim=5120, + num_heads=40, + ffn_dim=13824, + freq_dim=256, + text_dim=4096, + in_dim=36, + eps=1e-6, + ): + super().__init__() + self.mot_layers = mot_layers + self.freq_dim = freq_dim + self.dim = dim + + self.mot_layers_mapping = {i: n for n, i in enumerate(self.mot_layers)} + self.head_dim = dim // num_heads + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) + + # mot blocks + self.blocks = torch.nn.ModuleList([ + MotWanAttentionBlock(has_image_input, dim, num_heads, ffn_dim, eps, block_id=i) + for i in self.mot_layers + ]) + + + def patchify(self, x: torch.Tensor): + x = self.patch_embedding(x) + return x + + def compute_freqs_mot(self, f, h, w, end: int = 1024, theta: float = 10000.0): + def precompute_freqs_cis(dim: int, start: int = 0, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2) + [: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(start, end, device=freqs.device), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + f_freqs_cis = precompute_freqs_cis(self.head_dim - 2 * (self.head_dim // 3), -f, end, theta) + h_freqs_cis = precompute_freqs_cis(self.head_dim // 3, 0, end, theta) + w_freqs_cis = precompute_freqs_cis(self.head_dim // 3, 0, end, theta) + + freqs = torch.cat([ + f_freqs_cis[:f].view(f, 1, 1, -1).expand(f, h, w, -1), + h_freqs_cis[:h].view(1, h, 1, -1).expand(f, h, w, -1), + w_freqs_cis[:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1) + return freqs + + def forward(self, wan_block, x, context, t_mod, freqs, x_mot, context_mot, t_mod_mot, freqs_mot, block_id): + block = self.blocks[self.mot_layers_mapping[block_id]] + x, x_mot = block(wan_block, x, context, t_mod, freqs, x_mot, context_mot, t_mod_mot, freqs_mot) + return x, x_mot diff --git a/diffsynth/models/wan_video_motion_controller.py b/diffsynth/models/wan_video_motion_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..34763a8d76e57bc8efff84f23863938cc2309029 --- /dev/null +++ b/diffsynth/models/wan_video_motion_controller.py @@ -0,0 +1,27 @@ +import torch +import torch.nn as nn +from .wan_video_dit import sinusoidal_embedding_1d + + + +class WanMotionControllerModel(torch.nn.Module): + def __init__(self, freq_dim=256, dim=1536): + super().__init__() + self.freq_dim = freq_dim + self.linear = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim), + nn.SiLU(), + nn.Linear(dim, dim * 6), + ) + + def forward(self, motion_bucket_id): + emb = sinusoidal_embedding_1d(self.freq_dim, motion_bucket_id * 10) + emb = self.linear(emb) + return emb + + def init(self): + state_dict = self.linear[-1].state_dict() + state_dict = {i: state_dict[i] * 0 for i in state_dict} + self.linear[-1].load_state_dict(state_dict) diff --git a/diffsynth/models/wan_video_text_encoder.py b/diffsynth/models/wan_video_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..64090db8c65138abfdb60a822b3ba2e74fefeb4c --- /dev/null +++ b/diffsynth/models/wan_video_text_encoder.py @@ -0,0 +1,330 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import AutoTokenizer +import ftfy +import html +import string +import regex as re + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ + # torch.arange(lq).unsqueeze(1).to(device) + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class WanTextEncoder(torch.nn.Module): + + def __init__(self, + vocab=256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1): + super(WanTextEncoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \ + else nn.Embedding(vocab, dim) + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) if shared_pos else None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), + x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +def canonicalize(text, keep_punctuation_exact_string=None): + text = text.replace('_', ' ') + if keep_punctuation_exact_string: + text = keep_punctuation_exact_string.join( + part.translate(str.maketrans('', '', string.punctuation)) + for part in text.split(keep_punctuation_exact_string)) + else: + text = text.translate(str.maketrans('', '', string.punctuation)) + text = text.lower() + text = re.sub(r'\s+', ' ', text) + return text.strip() + + +class HuggingfaceTokenizer: + + def __init__(self, name, seq_len=None, clean=None, **kwargs): + assert clean in (None, 'whitespace', 'lower', 'canonicalize') + self.name = name + self.seq_len = seq_len + self.clean = clean + + # init tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs) + self.vocab_size = self.tokenizer.vocab_size + + def __call__(self, sequence, **kwargs): + return_mask = kwargs.pop('return_mask', False) + + # arguments + _kwargs = {'return_tensors': 'pt'} + if self.seq_len is not None: + _kwargs.update({ + 'padding': 'max_length', + 'truncation': True, + 'max_length': self.seq_len + }) + _kwargs.update(**kwargs) + + # tokenization + if isinstance(sequence, str): + sequence = [sequence] + if self.clean: + sequence = [self._clean(u) for u in sequence] + ids = self.tokenizer(sequence, **_kwargs) + + # output + if return_mask: + return ids.input_ids, ids.attention_mask + else: + return ids.input_ids + + def _clean(self, text): + if self.clean == 'whitespace': + text = whitespace_clean(basic_clean(text)) + elif self.clean == 'lower': + text = whitespace_clean(basic_clean(text)).lower() + elif self.clean == 'canonicalize': + text = canonicalize(basic_clean(text)) + return text \ No newline at end of file diff --git a/diffsynth/models/wan_video_vace.py b/diffsynth/models/wan_video_vace.py new file mode 100644 index 0000000000000000000000000000000000000000..f3367f788891cb22b8a5bf6eaa50cabbc202ab4a --- /dev/null +++ b/diffsynth/models/wan_video_vace.py @@ -0,0 +1,87 @@ +import torch +from .wan_video_dit import DiTBlock + + +class VaceWanAttentionBlock(DiTBlock): + def __init__(self, has_image_input, dim, num_heads, ffn_dim, eps=1e-6, block_id=0): + super().__init__(has_image_input, dim, num_heads, ffn_dim, eps=eps) + self.block_id = block_id + if block_id == 0: + self.before_proj = torch.nn.Linear(self.dim, self.dim) + self.after_proj = torch.nn.Linear(self.dim, self.dim) + + def forward(self, c, x, context, t_mod, freqs): + if self.block_id == 0: + c = self.before_proj(c) + x + all_c = [] + else: + all_c = list(torch.unbind(c)) + c = all_c.pop(-1) + c = super().forward(c, context, t_mod, freqs) + c_skip = self.after_proj(c) + all_c += [c_skip, c] + c = torch.stack(all_c) + return c + + +class VaceWanModel(torch.nn.Module): + def __init__( + self, + vace_layers=(0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28), + vace_in_dim=96, + patch_size=(1, 2, 2), + has_image_input=False, + dim=1536, + num_heads=12, + ffn_dim=8960, + eps=1e-6, + ): + super().__init__() + self.vace_layers = vace_layers + self.vace_in_dim = vace_in_dim + self.vace_layers_mapping = {i: n for n, i in enumerate(self.vace_layers)} + + # vace blocks + self.vace_blocks = torch.nn.ModuleList([ + VaceWanAttentionBlock(has_image_input, dim, num_heads, ffn_dim, eps, block_id=i) + for i in self.vace_layers + ]) + + # vace patch embeddings + self.vace_patch_embedding = torch.nn.Conv3d(vace_in_dim, dim, kernel_size=patch_size, stride=patch_size) + + def forward( + self, x, vace_context, context, t_mod, freqs, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + ): + c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context] + c = [u.flatten(2).transpose(1, 2) for u in c] + c = torch.cat([ + torch.cat([u, u.new_zeros(1, x.shape[1] - u.size(1), u.size(2))], + dim=1) for u in c + ]) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block in self.vace_blocks: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + c = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + c, x, context, t_mod, freqs, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + c = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + c, x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + c = block(c, x, context, t_mod, freqs) + hints = torch.unbind(c)[:-1] + return hints diff --git a/diffsynth/models/wan_video_vae.py b/diffsynth/models/wan_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..d24e29d9398f95a59cbd1466542c6e7059f7c7af --- /dev/null +++ b/diffsynth/models/wan_video_vae.py @@ -0,0 +1,1382 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _, device = *x.size(), x.device + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + else: + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) # may add tensor offload + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + self.z_dim = z_dim + + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if not left_bound: + x[:border_width] = (torch.arange(border_width) + 1) / border_width + if not right_bound: + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, device, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = T * 4 - 3 + weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype, device=data_device) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.clamp_(-1, 1) + return values + + + def tiled_encode(self, video, device, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = (T + 3) // 4 + weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype, device=data_device) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + return values + + + def single_encode(self, video, device): + video = video.to(device) + x = self.model.encode(video, self.scale) + return x + + + def single_decode(self, hidden_state, device): + hidden_state = hidden_state.to(device) + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + + def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + videos = [video.to("cpu") for video in videos] + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, device, tile_size, tile_stride) + else: + hidden_state = self.single_encode(video, device) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + + def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states] + videos = [] + for hidden_state in hidden_states: + hidden_state = hidden_state.unsqueeze(0) + if tiled: + video = self.tiled_decode(hidden_state, device, tile_size, tile_stride) + else: + video = self.single_decode(hidden_state, device) + video = video.squeeze(0) + videos.append(video) + videos = torch.stack(videos) + return videos + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True) + else: + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) + self.upsampling_factor = 16 + self.z_dim = z_dim diff --git a/diffsynth/models/wav2vec.py b/diffsynth/models/wav2vec.py new file mode 100644 index 0000000000000000000000000000000000000000..8807302d815a917123b794a597fc5fe84d3394fc --- /dev/null +++ b/diffsynth/models/wav2vec.py @@ -0,0 +1,191 @@ +import math +import numpy as np +import torch +import torch.nn.functional as F + + +def get_sample_indices(original_fps, total_frames, target_fps, num_sample, fixed_start=None): + required_duration = num_sample / target_fps + required_origin_frames = int(np.ceil(required_duration * original_fps)) + if required_duration > total_frames / original_fps: + raise ValueError("required_duration must be less than video length") + + if not fixed_start is None and fixed_start >= 0: + start_frame = fixed_start + else: + max_start = total_frames - required_origin_frames + if max_start < 0: + raise ValueError("video length is too short") + start_frame = np.random.randint(0, max_start + 1) + start_time = start_frame / original_fps + + end_time = start_time + required_duration + time_points = np.linspace(start_time, end_time, num_sample, endpoint=False) + + frame_indices = np.round(np.array(time_points) * original_fps).astype(int) + frame_indices = np.clip(frame_indices, 0, total_frames - 1) + return frame_indices + + +def linear_interpolation(features, input_fps, output_fps, output_len=None): + """ + features: shape=[1, T, 512] + input_fps: fps for audio, f_a + output_fps: fps for video, f_m + output_len: video length + """ + features = features.transpose(1, 2) + seq_len = features.shape[2] / float(input_fps) + if output_len is None: + output_len = int(seq_len * output_fps) + output_features = F.interpolate(features, size=output_len, align_corners=True, mode='linear') # [1, 512, output_len] + return output_features.transpose(1, 2) + + +class WanS2VAudioEncoder(torch.nn.Module): + + def __init__(self): + super().__init__() + from transformers import Wav2Vec2ForCTC, Wav2Vec2Config + config = { + "_name_or_path": "facebook/wav2vec2-large-xlsr-53", + "activation_dropout": 0.05, + "apply_spec_augment": True, + "architectures": ["Wav2Vec2ForCTC"], + "attention_dropout": 0.1, + "bos_token_id": 1, + "conv_bias": True, + "conv_dim": [512, 512, 512, 512, 512, 512, 512], + "conv_kernel": [10, 3, 3, 3, 3, 2, 2], + "conv_stride": [5, 2, 2, 2, 2, 2, 2], + "ctc_loss_reduction": "mean", + "ctc_zero_infinity": True, + "do_stable_layer_norm": True, + "eos_token_id": 2, + "feat_extract_activation": "gelu", + "feat_extract_dropout": 0.0, + "feat_extract_norm": "layer", + "feat_proj_dropout": 0.05, + "final_dropout": 0.0, + "hidden_act": "gelu", + "hidden_dropout": 0.05, + "hidden_size": 1024, + "initializer_range": 0.02, + "intermediate_size": 4096, + "layer_norm_eps": 1e-05, + "layerdrop": 0.05, + "mask_channel_length": 10, + "mask_channel_min_space": 1, + "mask_channel_other": 0.0, + "mask_channel_prob": 0.0, + "mask_channel_selection": "static", + "mask_feature_length": 10, + "mask_feature_prob": 0.0, + "mask_time_length": 10, + "mask_time_min_space": 1, + "mask_time_other": 0.0, + "mask_time_prob": 0.05, + "mask_time_selection": "static", + "model_type": "wav2vec2", + "num_attention_heads": 16, + "num_conv_pos_embedding_groups": 16, + "num_conv_pos_embeddings": 128, + "num_feat_extract_layers": 7, + "num_hidden_layers": 24, + "pad_token_id": 0, + "transformers_version": "4.7.0.dev0", + "vocab_size": 33 + } + self.model = Wav2Vec2ForCTC(Wav2Vec2Config(**config)) + self.video_rate = 30 + + def extract_audio_feat(self, input_audio, sample_rate, processor, return_all_layers=False, dtype=torch.float32, device='cpu'): + input_values = processor(input_audio, sampling_rate=sample_rate, return_tensors="pt").input_values.to(dtype=dtype, device=device) + + # retrieve logits & take argmax + res = self.model(input_values, output_hidden_states=True) + if return_all_layers: + feat = torch.cat(res.hidden_states) + else: + feat = res.hidden_states[-1] + feat = linear_interpolation(feat, input_fps=50, output_fps=self.video_rate) + return feat + + def get_audio_embed_bucket(self, audio_embed, stride=2, batch_frames=12, m=2): + num_layers, audio_frame_num, audio_dim = audio_embed.shape + + if num_layers > 1: + return_all_layers = True + else: + return_all_layers = False + + min_batch_num = int(audio_frame_num / (batch_frames * stride)) + 1 + + bucket_num = min_batch_num * batch_frames + batch_idx = [stride * i for i in range(bucket_num)] + batch_audio_eb = [] + for bi in batch_idx: + if bi < audio_frame_num: + audio_sample_stride = 2 + chosen_idx = list(range(bi - m * audio_sample_stride, bi + (m + 1) * audio_sample_stride, audio_sample_stride)) + chosen_idx = [0 if c < 0 else c for c in chosen_idx] + chosen_idx = [audio_frame_num - 1 if c >= audio_frame_num else c for c in chosen_idx] + + if return_all_layers: + frame_audio_embed = audio_embed[:, chosen_idx].flatten(start_dim=-2, end_dim=-1) + else: + frame_audio_embed = audio_embed[0][chosen_idx].flatten() + else: + frame_audio_embed = \ + torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \ + else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device) + batch_audio_eb.append(frame_audio_embed) + batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb], dim=0) + + return batch_audio_eb, min_batch_num + + def get_audio_embed_bucket_fps(self, audio_embed, fps=16, batch_frames=81, m=0): + num_layers, audio_frame_num, audio_dim = audio_embed.shape + + if num_layers > 1: + return_all_layers = True + else: + return_all_layers = False + + scale = self.video_rate / fps + + min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1 + + bucket_num = min_batch_num * batch_frames + padd_audio_num = math.ceil(min_batch_num * batch_frames / fps * self.video_rate) - audio_frame_num + batch_idx = get_sample_indices( + original_fps=self.video_rate, total_frames=audio_frame_num + padd_audio_num, target_fps=fps, num_sample=bucket_num, fixed_start=0 + ) + batch_audio_eb = [] + audio_sample_stride = int(self.video_rate / fps) + for bi in batch_idx: + if bi < audio_frame_num: + + chosen_idx = list(range(bi - m * audio_sample_stride, bi + (m + 1) * audio_sample_stride, audio_sample_stride)) + chosen_idx = [0 if c < 0 else c for c in chosen_idx] + chosen_idx = [audio_frame_num - 1 if c >= audio_frame_num else c for c in chosen_idx] + + if return_all_layers: + frame_audio_embed = audio_embed[:, chosen_idx].flatten(start_dim=-2, end_dim=-1) + else: + frame_audio_embed = audio_embed[0][chosen_idx].flatten() + else: + frame_audio_embed = \ + torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \ + else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device) + batch_audio_eb.append(frame_audio_embed) + batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb], dim=0) + + return batch_audio_eb, min_batch_num + + def get_audio_feats_per_inference(self, input_audio, sample_rate, processor, fps=16, batch_frames=80, m=0, dtype=torch.float32, device='cpu'): + audio_feat = self.extract_audio_feat(input_audio, sample_rate, processor, return_all_layers=True, dtype=dtype, device=device) + audio_embed_bucket, min_batch_num = self.get_audio_embed_bucket_fps(audio_feat, fps=fps, batch_frames=batch_frames, m=m) + audio_embed_bucket = audio_embed_bucket.unsqueeze(0).permute(0, 2, 3, 1).to(device, dtype) + audio_embeds = [audio_embed_bucket[..., i * batch_frames:(i + 1) * batch_frames] for i in range(min_batch_num)] + return audio_embeds diff --git a/diffsynth/models/z_image_dit.py b/diffsynth/models/z_image_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..7664fc5a37a2bf071677f09888cfa5e263ce7143 --- /dev/null +++ b/diffsynth/models/z_image_dit.py @@ -0,0 +1,621 @@ +import math +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from torch.nn import RMSNorm +from ..core.attention import attention_forward +from ..core.gradient import gradient_checkpoint_forward + + +ADALN_EMBED_DIM = 256 +SEQ_MULTI_OF = 32 + + +class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): + super().__init__() + if mid_size is None: + mid_size = out_size + self.mlp = nn.Sequential( + nn.Linear( + frequency_embedding_size, + mid_size, + bias=True, + ), + nn.SiLU(), + nn.Linear( + mid_size, + out_size, + bias=True, + ), + ) + + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq.to(torch.bfloat16)) + return t_emb + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def _forward_silu_gating(self, x1, x3): + return F.silu(x1) * x3 + + def forward(self, x): + return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) + + +class Attention(torch.nn.Module): + + def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): + super().__init__() + dim_inner = head_dim * num_heads + kv_dim = kv_dim if kv_dim is not None else q_dim + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) + self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) + self.to_out = torch.nn.ModuleList([torch.nn.Linear(dim_inner, q_dim, bias=bias_out)]) + + self.norm_q = RMSNorm(head_dim, eps=1e-5) + self.norm_k = RMSNorm(head_dim, eps=1e-5) + + def forward(self, hidden_states, freqs_cis): + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + query = query.unflatten(-1, (self.num_heads, -1)) + key = key.unflatten(-1, (self.num_heads, -1)) + value = value.unflatten(-1, (self.num_heads, -1)) + + # Apply Norms + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + # Apply RoPE + def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) # todo + + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) + + # Cast to correct dtype + dtype = query.dtype + query, key = query.to(dtype), key.to(dtype) + + # Compute joint attention + hidden_states = attention_forward( + query, + key, + value, + q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d", + ) + + # Reshape back + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(dtype) + + output = self.to_out[0](hidden_states) + if len(self.to_out) > 1: # dropout + output = self.to_out[1](output) + + return output + + +class ZImageTransformerBlock(nn.Module): + def __init__( + self, + layer_id: int, + dim: int, + n_heads: int, + n_kv_heads: int, + norm_eps: float, + qk_norm: bool, + modulation=True, + ): + super().__init__() + self.dim = dim + self.head_dim = dim // n_heads + + # Refactored to use diffusers Attention with custom processor + # Original Z-Image params: dim, n_heads, n_kv_heads, qk_norm + self.attention = Attention( + q_dim=dim, + num_heads=n_heads, + head_dim=dim // n_heads, + ) + + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + self.layer_id = layer_id + + self.attention_norm1 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + + self.attention_norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + self.modulation = modulation + if modulation: + self.adaLN_modulation = nn.Sequential( + nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True), + ) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, + adaln_input: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert adaln_input is not None + scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # Attention block + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, + freqs_cis=freqs_cis, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + + # FFN block + x = x + gate_mlp * self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x) * scale_mlp, + ) + ) + else: + # Attention block + attn_out = self.attention( + self.attention_norm1(x), + freqs_cis=freqs_cis, + ) + x = x + self.attention_norm2(attn_out) + + # FFN block + x = x + self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x), + ) + ) + + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), + ) + + def forward(self, x, c): + scale = 1.0 + self.adaLN_modulation(c) + x = self.norm_final(x) * scale.unsqueeze(1) + x = self.linear(x) + return x + + +class RopeEmbedder: + def __init__( + self, + theta: float = 256.0, + axes_dims: List[int] = (16, 56, 56), + axes_lens: List[int] = (64, 128, 128), + ): + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length" + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + with torch.device("cpu"): + freqs_cis = [] + for i, (d, e) in enumerate(zip(dim, end)): + freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64) # complex64 + freqs_cis.append(freqs_cis_i) + + return freqs_cis + + def __call__(self, ids: torch.Tensor): + assert ids.ndim == 2 + assert ids.shape[-1] == len(self.axes_dims) + device = ids.device + + if self.freqs_cis is None: + self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i] + result.append(self.freqs_cis[i][index]) + return torch.cat(result, dim=-1) + + +class ZImageDiT(nn.Module): + _supports_gradient_checkpointing = True + _no_split_modules = ["ZImageTransformerBlock"] + + def __init__( + self, + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=3840, + n_layers=30, + n_refiner_layers=2, + n_heads=30, + n_kv_heads=30, + norm_eps=1e-5, + qk_norm=True, + cap_feat_dim=2560, + rope_theta=256.0, + t_scale=1000.0, + axes_dims=[32, 48, 48], + axes_lens=[1024, 512, 512], + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.all_patch_size = all_patch_size + self.all_f_patch_size = all_f_patch_size + self.dim = dim + self.n_heads = n_heads + + self.rope_theta = rope_theta + self.t_scale = t_scale + self.gradient_checkpointing = False + + assert len(all_patch_size) == len(all_f_patch_size) + + all_x_embedder = {} + all_final_layer = {} + for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): + x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True) + all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder + + final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels) + all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer + + self.all_x_embedder = nn.ModuleDict(all_x_embedder) + self.all_final_layer = nn.ModuleDict(all_final_layer) + self.noise_refiner = nn.ModuleList( + [ + ZImageTransformerBlock( + 1000 + layer_id, + dim, + n_heads, + n_kv_heads, + norm_eps, + qk_norm, + modulation=True, + ) + for layer_id in range(n_refiner_layers) + ] + ) + self.context_refiner = nn.ModuleList( + [ + ZImageTransformerBlock( + layer_id, + dim, + n_heads, + n_kv_heads, + norm_eps, + qk_norm, + modulation=False, + ) + for layer_id in range(n_refiner_layers) + ] + ) + self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024) + self.cap_embedder = nn.Sequential( + RMSNorm(cap_feat_dim, eps=norm_eps), + nn.Linear(cap_feat_dim, dim, bias=True), + ) + + self.x_pad_token = nn.Parameter(torch.empty((1, dim))) + self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) + + self.layers = nn.ModuleList( + [ + ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm) + for layer_id in range(n_layers) + ] + ) + head_dim = dim // n_heads + assert head_dim == sum(axes_dims) + self.axes_dims = axes_dims + self.axes_lens = axes_lens + + self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + + def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]: + pH = pW = patch_size + pF = f_patch_size + bsz = len(x) + assert len(size) == bsz + for i in range(bsz): + F, H, W = size[i] + ori_len = (F // pF) * (H // pH) * (W // pW) + # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)" + x[i] = ( + x[i][:ori_len] + .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) + .permute(6, 0, 3, 1, 4, 2, 5) + .reshape(self.out_channels, F, H, W) + ) + return x + + @staticmethod + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + + axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, dim=-1) + + def patchify_and_embed( + self, + all_image: List[torch.Tensor], + all_cap_feats: List[torch.Tensor], + patch_size: int, + f_patch_size: int, + ): + pH = pW = patch_size + pF = f_patch_size + device = all_image[0].device + + all_image_out = [] + all_image_size = [] + all_image_pos_ids = [] + all_image_pad_mask = [] + all_cap_pos_ids = [] + all_cap_pad_mask = [] + all_cap_feats_out = [] + + for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): + ### Process Caption + cap_ori_len = len(cap_feat) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + # padded position ids + cap_padded_pos_ids = self.create_coordinate_grid( + size=(cap_ori_len + cap_padding_len, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2) + all_cap_pos_ids.append(cap_padded_pos_ids) + # pad mask + all_cap_pad_mask.append( + torch.cat( + [ + torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), + torch.ones((cap_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + ) + # padded feature + cap_padded_feat = torch.cat( + [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], + dim=0, + ) + all_cap_feats_out.append(cap_padded_feat) + + ### Process Image + C, F, H, W = image.size() + all_image_size.append((F, H, W)) + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + + image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + # "c f pf h ph w pw -> (f h w) (pf ph pw c)" + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C) + + image_ori_len = len(image) + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + + image_ori_pos_ids = self.create_coordinate_grid( + size=(F_tokens, H_tokens, W_tokens), + start=(cap_ori_len + cap_padding_len + 1, 0, 0), + device=device, + ).flatten(0, 2) + image_padding_pos_ids = ( + self.create_coordinate_grid( + size=(1, 1, 1), + start=(0, 0, 0), + device=device, + ) + .flatten(0, 2) + .repeat(image_padding_len, 1) + ) + image_padded_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0) + all_image_pos_ids.append(image_padded_pos_ids) + # pad mask + all_image_pad_mask.append( + torch.cat( + [ + torch.zeros((image_ori_len,), dtype=torch.bool, device=device), + torch.ones((image_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + ) + # padded feature + image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0) + all_image_out.append(image_padded_feat) + + return ( + all_image_out, + all_cap_feats_out, + all_image_size, + all_image_pos_ids, + all_cap_pos_ids, + all_image_pad_mask, + all_cap_pad_mask, + ) + + def forward( + self, + x: List[torch.Tensor], + t, + cap_feats: List[torch.Tensor], + patch_size=2, + f_patch_size=1, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + ): + assert patch_size in self.all_patch_size + assert f_patch_size in self.all_f_patch_size + + bsz = len(x) + device = x[0].device + t = t * self.t_scale + t = self.t_embedder(t) + + adaln_input = t + + ( + x, + cap_feats, + x_size, + x_pos_ids, + cap_pos_ids, + x_inner_pad_mask, + cap_inner_pad_mask, + ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + + # x embed & refine + x_item_seqlens = [len(_) for _ in x] + assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) + x_max_item_seqlen = max(x_item_seqlens) + + x = torch.cat(x, dim=0) + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + x[torch.cat(x_inner_pad_mask)] = self.x_pad_token.to(dtype=x.dtype, device=x.device) + x = list(x.split(x_item_seqlens, dim=0)) + x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) + + x = pad_sequence(x, batch_first=True, padding_value=0.0) + x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) + x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(x_item_seqlens): + x_attn_mask[i, :seq_len] = 1 + + for layer in self.noise_refiner: + x = gradient_checkpoint_forward( + layer, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + x=x, + attn_mask=x_attn_mask, + freqs_cis=x_freqs_cis, + adaln_input=adaln_input, + ) + + # cap embed & refine + cap_item_seqlens = [len(_) for _ in cap_feats] + assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) + cap_max_item_seqlen = max(cap_item_seqlens) + + cap_feats = torch.cat(cap_feats, dim=0) + cap_feats = self.cap_embedder(cap_feats) + cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token.to(dtype=x.dtype, device=x.device) + cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) + cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) + + cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) + cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) + cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(cap_item_seqlens): + cap_attn_mask[i, :seq_len] = 1 + + for layer in self.context_refiner: + cap_feats = gradient_checkpoint_forward( + layer, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + x=cap_feats, + attn_mask=cap_attn_mask, + freqs_cis=cap_freqs_cis, + ) + + # unified + unified = [] + unified_freqs_cis = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) + unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + assert unified_item_seqlens == [len(_) for _ in unified] + unified_max_item_seqlen = max(unified_item_seqlens) + + unified = pad_sequence(unified, batch_first=True, padding_value=0.0) + unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) + unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + for layer in self.layers: + unified = gradient_checkpoint_forward( + layer, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + x=unified, + attn_mask=unified_attn_mask, + freqs_cis=unified_freqs_cis, + adaln_input=adaln_input, + ) + + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) + unified = list(unified.unbind(dim=0)) + x = self.unpatchify(unified, x_size, patch_size, f_patch_size) + + return x, {} diff --git a/diffsynth/models/z_image_text_encoder.py b/diffsynth/models/z_image_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4eba636058b6299e41bec476ac348cbacc39cc37 --- /dev/null +++ b/diffsynth/models/z_image_text_encoder.py @@ -0,0 +1,41 @@ +from transformers import Qwen3Model, Qwen3Config +import torch + + +class ZImageTextEncoder(torch.nn.Module): + def __init__(self): + super().__init__() + config = Qwen3Config(**{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "max_position_embeddings": 40960, + "max_window_layers": 36, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "rope_theta": 1000000, + "sliding_window": None, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936 + }) + self.model = Qwen3Model(config) + + def forward(self, *args, **kwargs): + return self.model(*args, **kwargs) diff --git a/diffsynth/pipelines/flux2_image.py b/diffsynth/pipelines/flux2_image.py new file mode 100644 index 0000000000000000000000000000000000000000..8b0046949b8766a087d4746f473406ac54f686bb --- /dev/null +++ b/diffsynth/pipelines/flux2_image.py @@ -0,0 +1,370 @@ +import torch, math +from PIL import Image +from typing import Union +from tqdm import tqdm +from einops import rearrange +import numpy as np +from typing import Union, List, Optional, Tuple + +from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig, gradient_checkpoint_forward +from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput + +from transformers import AutoProcessor +from ..models.flux2_text_encoder import Flux2TextEncoder +from ..models.flux2_dit import Flux2DiT +from ..models.flux2_vae import Flux2VAE + + +class Flux2ImagePipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, + ) + self.scheduler = FlowMatchScheduler("FLUX.2") + self.text_encoder: Flux2TextEncoder = None + self.dit: Flux2DiT = None + self.vae: Flux2VAE = None + self.tokenizer: AutoProcessor = None + self.in_iteration_models = ("dit",) + self.units = [ + Flux2Unit_ShapeChecker(), + Flux2Unit_PromptEmbedder(), + Flux2Unit_NoiseInitializer(), + Flux2Unit_InputImageEmbedder(), + Flux2Unit_ImageIDs(), + ] + self.model_fn = model_fn_flux2 + + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_config: ModelConfig = ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"), + vram_limit: float = None, + ): + # Initialize pipeline + pipe = Flux2ImagePipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder = model_pool.fetch_model("flux2_text_encoder") + pipe.dit = model_pool.fetch_model("flux2_dit") + pipe.vae = model_pool.fetch_model("flux2_vae") + if tokenizer_config is not None: + tokenizer_config.download_if_necessary() + pipe.tokenizer = AutoProcessor.from_pretrained(tokenizer_config.path) + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: str = "", + cfg_scale: float = 1.0, + embedded_guidance: float = 4.0, + # Image + input_image: Image.Image = None, + denoising_strength: float = 1.0, + # Shape + height: int = 1024, + width: int = 1024, + # Randomness + seed: int = None, + rand_device: str = "cpu", + # Steps + num_inference_steps: int = 30, + # Progress bar + progress_bar_cmd = tqdm, + ): + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, dynamic_shift_len=height//16*width//16) + + # Parameters + inputs_posi = { + "prompt": prompt, + } + inputs_nega = { + "negative_prompt": negative_prompt, + } + inputs_shared = { + "cfg_scale": cfg_scale, "embedded_guidance": embedded_guidance, + "input_image": input_image, "denoising_strength": denoising_strength, + "height": height, "width": width, + "seed": seed, "rand_device": rand_device, + "num_inference_steps": num_inference_steps, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + noise_pred = self.cfg_guided_model_fn( + self.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) + + # Decode + self.load_models_to_device(['vae']) + latents = rearrange(inputs_shared["latents"], "B (H W) C -> B C H W", H=inputs_shared["height"]//16, W=inputs_shared["width"]//16) + image = self.vae.decode(latents) + image = self.vae_output_to_image(image) + self.load_models_to_device([]) + + return image + + +class Flux2Unit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width"), + output_params=("height", "width"), + ) + + def process(self, pipe: Flux2ImagePipeline, height, width): + height, width = pipe.check_resize_height_width(height, width) + return {"height": height, "width": width} + + +class Flux2Unit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt"}, + input_params_nega={"prompt": "negative_prompt"}, + output_params=("prompt_emb", "prompt_emb_mask"), + onload_model_names=("text_encoder",) + ) + self.system_message = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation." + + def format_text_input(self, prompts: List[str], system_message: str = None): + # Remove [IMG] tokens from prompts to avoid Pixtral validation issues + # when truncation is enabled. The processor counts [IMG] tokens and fails + # if the count changes after truncation. + cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts] + + return [ + [ + { + "role": "system", + "content": [{"type": "text", "text": system_message}], + }, + {"role": "user", "content": [{"type": "text", "text": prompt}]}, + ] + for prompt in cleaned_txt + ] + + def get_mistral_3_small_prompt_embeds( + self, + text_encoder, + tokenizer, + prompt: Union[str, List[str]], + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + max_sequence_length: int = 512, + # fmt: off + system_message: str = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation.", + # fmt: on + hidden_states_layers: List[int] = (10, 20, 30), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + # Format input messages + messages_batch = self.format_text_input(prompts=prompt, system_message=system_message) + + # Process all messages at once + inputs = tokenizer.apply_chat_template( + messages_batch, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + # Move to device + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + def prepare_text_ids( + self, + x: torch.Tensor, # (B, L, D) or (L, D) + t_coord: Optional[torch.Tensor] = None, + ): + B, L, _ = x.shape + out_ids = [] + + for i in range(B): + t = torch.arange(1) if t_coord is None else t_coord[i] + h = torch.arange(1) + w = torch.arange(1) + l = torch.arange(L) + + coords = torch.cartesian_prod(t, h, w, l) + out_ids.append(coords) + + return torch.stack(out_ids) + + def encode_prompt( + self, + text_encoder, + tokenizer, + prompt: Union[str, List[str]], + dtype = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + max_sequence_length: int = 512, + text_encoder_out_layers: Tuple[int] = (10, 20, 30), + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds = self.get_mistral_3_small_prompt_embeds( + text_encoder=text_encoder, + tokenizer=tokenizer, + prompt=prompt, + dtype=dtype, + device=device, + max_sequence_length=max_sequence_length, + system_message=self.system_message, + hidden_states_layers=text_encoder_out_layers, + ) + + batch_size, seq_len, _ = prompt_embeds.shape + 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) + + text_ids = self.prepare_text_ids(prompt_embeds) + text_ids = text_ids.to(device) + return prompt_embeds, text_ids + + def process(self, pipe: Flux2ImagePipeline, prompt): + pipe.load_models_to_device(self.onload_model_names) + prompt_embeds, text_ids = self.encode_prompt( + pipe.text_encoder, pipe.tokenizer, prompt, + dtype=pipe.torch_dtype, device=pipe.device, + ) + return {"prompt_embeds": prompt_embeds, "text_ids": text_ids} + + +class Flux2Unit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "seed", "rand_device"), + output_params=("noise",), + ) + + def process(self, pipe: Flux2ImagePipeline, height, width, seed, rand_device): + noise = pipe.generate_noise((1, 128, height//16, width//16), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype) + noise = noise.reshape(1, 128, height//16 * width//16).permute(0, 2, 1) + return {"noise": noise} + + +class Flux2Unit_InputImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "noise"), + output_params=("latents", "input_latents"), + onload_model_names=("vae",) + ) + + def process(self, pipe: Flux2ImagePipeline, input_image, noise): + if input_image is None: + return {"latents": noise, "input_latents": None} + pipe.load_models_to_device(['vae']) + image = pipe.preprocess_image(input_image) + input_latents = pipe.vae.encode(image) + input_latents = rearrange(input_latents, "B C H W -> B (H W) C") + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents, "input_latents": input_latents} + + +class Flux2Unit_ImageIDs(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width"), + output_params=("image_ids",), + ) + + def prepare_latent_ids(self, height, width): + t = torch.arange(1) # [0] - time dimension + h = torch.arange(height) + w = torch.arange(width) + l = torch.arange(1) # [0] - layer dimension + + # Create position IDs: (H*W, 4) + latent_ids = torch.cartesian_prod(t, h, w, l) + + # Expand to batch: (B, H*W, 4) + latent_ids = latent_ids.unsqueeze(0).expand(1, -1, -1) + + return latent_ids + + def process(self, pipe: Flux2ImagePipeline, height, width): + image_ids = self.prepare_latent_ids(height // 16, width // 16).to(pipe.device) + return {"image_ids": image_ids} + + +def model_fn_flux2( + dit: Flux2DiT, + latents=None, + timestep=None, + embedded_guidance=None, + prompt_embeds=None, + text_ids=None, + image_ids=None, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + **kwargs, +): + embedded_guidance = torch.tensor([embedded_guidance], device=latents.device) + model_output = dit( + hidden_states=latents, + timestep=timestep / 1000, + guidance=embedded_guidance, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=image_ids, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + ) + return model_output diff --git a/diffsynth/pipelines/flux_image.py b/diffsynth/pipelines/flux_image.py new file mode 100644 index 0000000000000000000000000000000000000000..1ee5635ee16eac74ae4a3fd322d4aa4bd323893d --- /dev/null +++ b/diffsynth/pipelines/flux_image.py @@ -0,0 +1,1205 @@ +import torch, math +from PIL import Image +from typing import Union +from tqdm import tqdm +from einops import rearrange, repeat +import numpy as np +from transformers import CLIPTokenizer, T5TokenizerFast + +from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig, gradient_checkpoint_forward, load_state_dict +from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput +from ..utils.lora.flux import FluxLoRALoader + +from ..models.flux_dit import FluxDiT +from ..models.flux_text_encoder_clip import FluxTextEncoderClip +from ..models.flux_text_encoder_t5 import FluxTextEncoderT5 +from ..models.flux_vae import FluxVAEEncoder, FluxVAEDecoder +from ..models.flux_value_control import MultiValueEncoder +from ..models.step1x_text_encoder import Step1xEditEmbedder +from ..core.vram.layers import AutoWrappedLinear + +class MultiControlNet(torch.nn.Module): + def __init__(self, models: list[torch.nn.Module]): + super().__init__() + if not isinstance(models, list): + models = [models] + self.models = torch.nn.ModuleList(models) + + def process_single_controlnet(self, controlnet_input: ControlNetInput, conditioning: torch.Tensor, **kwargs): + model = self.models[controlnet_input.controlnet_id] + res_stack, single_res_stack = model( + controlnet_conditioning=conditioning, + processor_id=controlnet_input.processor_id, + **kwargs + ) + res_stack = [res * controlnet_input.scale for res in res_stack] + single_res_stack = [res * controlnet_input.scale for res in single_res_stack] + return res_stack, single_res_stack + + def forward(self, conditionings: list[torch.Tensor], controlnet_inputs: list[ControlNetInput], progress_id, num_inference_steps, **kwargs): + res_stack, single_res_stack = None, None + for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): + progress = (num_inference_steps - 1 - progress_id) / max(num_inference_steps - 1, 1) + if progress > controlnet_input.start or progress < controlnet_input.end: + continue + res_stack_, single_res_stack_ = self.process_single_controlnet(controlnet_input, conditioning, **kwargs) + if res_stack is None: + res_stack = res_stack_ + single_res_stack = single_res_stack_ + else: + res_stack = [i + j for i, j in zip(res_stack, res_stack_)] + single_res_stack = [i + j for i, j in zip(single_res_stack, single_res_stack_)] + return res_stack, single_res_stack + + +class FluxImagePipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, + ) + self.scheduler = FlowMatchScheduler("FLUX.1") + self.tokenizer_1: CLIPTokenizer = None + self.tokenizer_2: T5TokenizerFast = None + self.text_encoder_1: FluxTextEncoderClip = None + self.text_encoder_2: FluxTextEncoderT5 = None + self.dit: FluxDiT = None + self.vae_decoder: FluxVAEDecoder = None + self.vae_encoder: FluxVAEEncoder = None + self.controlnet = None + self.ipadapter = None + self.ipadapter_image_encoder = None + self.qwenvl = None + self.step1x_connector = None + self.nexus_gen = None + self.nexus_gen_generation_adapter = None + self.nexus_gen_editing_adapter = None + self.value_controller = None + self.infinityou_processor = None + self.image_proj_model = None + self.lora_patcher = None + self.lora_encoder = None + self.in_iteration_models = ("dit", "step1x_connector", "controlnet", "lora_patcher") + self.units = [ + FluxImageUnit_ShapeChecker(), + FluxImageUnit_NoiseInitializer(), + FluxImageUnit_PromptEmbedder(), + FluxImageUnit_InputImageEmbedder(), + FluxImageUnit_ImageIDs(), + FluxImageUnit_EmbeddedGuidanceEmbedder(), + FluxImageUnit_Kontext(), + FluxImageUnit_InfiniteYou(), + FluxImageUnit_ControlNet(), + FluxImageUnit_IPAdapter(), + FluxImageUnit_EntityControl(), + FluxImageUnit_NexusGen(), + FluxImageUnit_TeaCache(), + FluxImageUnit_Flex(), + FluxImageUnit_Step1x(), + FluxImageUnit_ValueControl(), + FluxImageUnit_LoRAEncode(), + ] + self.model_fn = model_fn_flux_image + self.lora_loader = FluxLoRALoader + + def enable_lora_merger(self): + if not (hasattr(self.dit, "vram_management_enabled") and getattr(self.dit, "vram_management_enabled")): + raise ValueError("DiT VRAM management is not enabled.") + if self.lora_patcher is not None: + for name, module in self.dit.named_modules(): + if isinstance(module, AutoWrappedLinear): + merger_name = name.replace(".", "___") + if merger_name in self.lora_patcher.model_dict: + module.lora_merger = self.lora_patcher.model_dict[merger_name] + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_1_config: ModelConfig = ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="tokenizer/"), + tokenizer_2_config: ModelConfig = ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="tokenizer_2/"), + nexus_gen_processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="processor/"), + step1x_processor_config: ModelConfig = ModelConfig(model_id="Qwen/Qwen2.5-VL-7B-Instruct", origin_file_pattern=""), + vram_limit: float = None, + ): + # Initialize pipeline + pipe = FluxImagePipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder_1 = model_pool.fetch_model("flux_text_encoder_clip") + pipe.text_encoder_2 = model_pool.fetch_model("flux_text_encoder_t5") + pipe.dit = model_pool.fetch_model("flux_dit") + pipe.vae_encoder = model_pool.fetch_model("flux_vae_encoder") + pipe.vae_decoder = model_pool.fetch_model("flux_vae_decoder") + if tokenizer_1_config is not None: + tokenizer_1_config.download_if_necessary() + pipe.tokenizer_1 = CLIPTokenizer.from_pretrained(tokenizer_1_config.path) + if tokenizer_2_config is not None: + tokenizer_2_config.download_if_necessary() + pipe.tokenizer_2 = T5TokenizerFast.from_pretrained(tokenizer_2_config.path) + + value_controllers = model_pool.fetch_model("flux_value_controller") + if value_controllers is not None: + pipe.value_controller = MultiValueEncoder(value_controllers) + if hasattr(pipe.value_controller.encoders[0], "vram_management_enabled"): + pipe.value_controller.vram_management_enabled = pipe.value_controller.encoders[0].vram_management_enabled + controlnets = model_pool.fetch_model("flux_controlnet") + if controlnets is not None: pipe.controlnet = MultiControlNet(controlnets) + pipe.ipadapter = model_pool.fetch_model("flux_ipadapter") + pipe.ipadapter_image_encoder = model_pool.fetch_model("siglip_vision_model") + qwenvl = model_pool.fetch_model("qwen_image_text_encoder") + if qwenvl is not None: + from transformers import AutoProcessor + step1x_processor_config.download_if_necessary() + processor = AutoProcessor.from_pretrained(step1x_processor_config.path, min_pixels=256 * 28 * 28, max_pixels=324 * 28 * 28) + pipe.qwenvl = Step1xEditEmbedder(qwenvl, processor) + pipe.step1x_connector = model_pool.fetch_model("step1x_connector") + pipe.image_proj_model = model_pool.fetch_model("infiniteyou_image_projector") + if pipe.image_proj_model is not None: + pipe.infinityou_processor = InfinitYou(device=device) + pipe.lora_patcher = model_pool.fetch_model("flux_lora_patcher") + pipe.lora_encoder = model_pool.fetch_model("flux_lora_encoder") + pipe.nexus_gen = model_pool.fetch_model("nexus_gen_llm") + pipe.nexus_gen_generation_adapter = model_pool.fetch_model("nexus_gen_generation_adapter") + pipe.nexus_gen_editing_adapter = model_pool.fetch_model("nexus_gen_editing_adapter") + if pipe.nexus_gen is not None: + nexus_gen_processor_config.download_if_necessary() + pipe.nexus_gen.load_processor(nexus_gen_processor_config.path) + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: str = "", + cfg_scale: float = 1.0, + embedded_guidance: float = 3.5, + t5_sequence_length: int = 512, + # Image + input_image: Image.Image = None, + denoising_strength: float = 1.0, + # Shape + height: int = 1024, + width: int = 1024, + # Randomness + seed: int = None, + rand_device: str = "cpu", + # Scheduler + sigma_shift: float = None, + # Steps + num_inference_steps: int = 30, + # local prompts + multidiffusion_prompts=(), + multidiffusion_masks=(), + multidiffusion_scales=(), + # Kontext + kontext_images: Union[list[Image.Image], Image.Image] = None, + # ControlNet + controlnet_inputs: list[ControlNetInput] = None, + # IP-Adapter + ipadapter_images: Union[list[Image.Image], Image.Image] = None, + ipadapter_scale: float = 1.0, + # EliGen + eligen_entity_prompts: list[str] = None, + eligen_entity_masks: list[Image.Image] = None, + eligen_enable_on_negative: bool = False, + eligen_enable_inpaint: bool = False, + # InfiniteYou + infinityou_id_image: Image.Image = None, + infinityou_guidance: float = 1.0, + # Flex + flex_inpaint_image: Image.Image = None, + flex_inpaint_mask: Image.Image = None, + flex_control_image: Image.Image = None, + flex_control_strength: float = 0.5, + flex_control_stop: float = 0.5, + # Value Controller + value_controller_inputs: Union[list[float], float] = None, + # Step1x + step1x_reference_image: Image.Image = None, + # NexusGen + nexus_gen_reference_image: Image.Image = None, + # LoRA Encoder + lora_encoder_inputs: Union[list[ModelConfig], ModelConfig, str] = None, + lora_encoder_scale: float = 1.0, + # TeaCache + tea_cache_l1_thresh: float = None, + # Tile + tiled: bool = False, + tile_size: int = 128, + tile_stride: int = 64, + # Progress bar + progress_bar_cmd = tqdm, + ): + # Scheduler + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, shift=sigma_shift) + + inputs_posi = { + "prompt": prompt, + } + inputs_nega = { + "negative_prompt": negative_prompt, + } + inputs_shared = { + "cfg_scale": cfg_scale, "embedded_guidance": embedded_guidance, "t5_sequence_length": t5_sequence_length, + "input_image": input_image, "denoising_strength": denoising_strength, + "height": height, "width": width, + "seed": seed, "rand_device": rand_device, + "sigma_shift": sigma_shift, "num_inference_steps": num_inference_steps, + "multidiffusion_prompts": multidiffusion_prompts, "multidiffusion_masks": multidiffusion_masks, "multidiffusion_scales": multidiffusion_scales, + "kontext_images": kontext_images, + "controlnet_inputs": controlnet_inputs, + "ipadapter_images": ipadapter_images, "ipadapter_scale": ipadapter_scale, + "eligen_entity_prompts": eligen_entity_prompts, "eligen_entity_masks": eligen_entity_masks, "eligen_enable_on_negative": eligen_enable_on_negative, "eligen_enable_inpaint": eligen_enable_inpaint, + "infinityou_id_image": infinityou_id_image, "infinityou_guidance": infinityou_guidance, + "flex_inpaint_image": flex_inpaint_image, "flex_inpaint_mask": flex_inpaint_mask, "flex_control_image": flex_control_image, "flex_control_strength": flex_control_strength, "flex_control_stop": flex_control_stop, + "value_controller_inputs": value_controller_inputs, + "step1x_reference_image": step1x_reference_image, + "nexus_gen_reference_image": nexus_gen_reference_image, + "lora_encoder_inputs": lora_encoder_inputs, "lora_encoder_scale": lora_encoder_scale, + "tea_cache_l1_thresh": tea_cache_l1_thresh, + "tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride, + "progress_bar_cmd": progress_bar_cmd, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + noise_pred = self.cfg_guided_model_fn( + self.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) + + # Decode + self.load_models_to_device(['vae_decoder']) + image = self.vae_decoder(inputs_shared["latents"], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + image = self.vae_output_to_image(image) + self.load_models_to_device([]) + + return image + + +class FluxImageUnit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__(input_params=("height", "width"), output_params=("height", "width")) + + def process(self, pipe: FluxImagePipeline, height, width): + height, width = pipe.check_resize_height_width(height, width) + return {"height": height, "width": width} + + + +class FluxImageUnit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__(input_params=("height", "width", "seed", "rand_device"), output_params=("noise",)) + + def process(self, pipe: FluxImagePipeline, height, width, seed, rand_device): + noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device) + return {"noise": noise} + + + +class FluxImageUnit_InputImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "noise", "tiled", "tile_size", "tile_stride"), + output_params=("latents", "input_latents"), + onload_model_names=("vae_encoder",) + ) + + def process(self, pipe: FluxImagePipeline, input_image, noise, tiled, tile_size, tile_stride): + if input_image is None: + return {"latents": noise, "input_latents": None} + pipe.load_models_to_device(['vae_encoder']) + image = pipe.preprocess_image(input_image).to(device=pipe.device, dtype=pipe.torch_dtype) + input_latents = pipe.vae_encoder(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents, "input_latents": None} + + + +class FluxImageUnit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt", "positive": "positive"}, + input_params_nega={"prompt": "negative_prompt", "positive": "positive"}, + input_params=("t5_sequence_length",), + output_params=("prompt_emb", "pooled_prompt_emb", "text_ids"), + onload_model_names=("text_encoder_1", "text_encoder_2") + ) + + def encode_prompt_using_clip(self, prompt, text_encoder, tokenizer, max_length, device): + input_ids = tokenizer( + prompt, + return_tensors="pt", + padding="max_length", + max_length=max_length, + truncation=True + ).input_ids.to(device) + pooled_prompt_emb, _ = text_encoder(input_ids) + return pooled_prompt_emb + + def encode_prompt_using_t5(self, prompt, text_encoder, tokenizer, max_length, device): + input_ids = tokenizer( + prompt, + return_tensors="pt", + padding="max_length", + max_length=max_length, + truncation=True, + ).input_ids.to(device) + prompt_emb = text_encoder(input_ids) + return prompt_emb + + def encode_prompt( + self, + tokenizer_1, + tokenizer_2, + text_encoder_1, + text_encoder_2, + prompt, + positive=True, + device="cuda", + t5_sequence_length=512, + ): + pooled_prompt_emb = self.encode_prompt_using_clip(prompt, text_encoder_1, tokenizer_1, 77, device) + prompt_emb = self.encode_prompt_using_t5(prompt, text_encoder_2, tokenizer_2, t5_sequence_length, device) + text_ids = torch.zeros(prompt_emb.shape[0], prompt_emb.shape[1], 3).to(device=device, dtype=prompt_emb.dtype) + return prompt_emb, pooled_prompt_emb, text_ids + + def process(self, pipe: FluxImagePipeline, prompt, t5_sequence_length, positive) -> dict: + if pipe.text_encoder_1 is not None and pipe.text_encoder_2 is not None: + prompt_emb, pooled_prompt_emb, text_ids = self.encode_prompt( + tokenizer_1=pipe.tokenizer_1, tokenizer_2=pipe.tokenizer_2, + text_encoder_1=pipe.text_encoder_1, text_encoder_2=pipe.text_encoder_2, + prompt=prompt, device=pipe.device, positive=positive, t5_sequence_length=t5_sequence_length, + ) + return {"prompt_emb": prompt_emb, "pooled_prompt_emb": pooled_prompt_emb, "text_ids": text_ids} + else: + return {} + + +class FluxImageUnit_ImageIDs(PipelineUnit): + def __init__(self): + super().__init__(input_params=("latents",), output_params=("image_ids",)) + + def process(self, pipe: FluxImagePipeline, latents): + latent_image_ids = pipe.dit.prepare_image_ids(latents) + return {"image_ids": latent_image_ids} + + + +class FluxImageUnit_EmbeddedGuidanceEmbedder(PipelineUnit): + def __init__(self): + super().__init__(input_params=("embedded_guidance", "latents"), output_params=("guidance",)) + + def process(self, pipe: FluxImagePipeline, embedded_guidance, latents): + guidance = torch.Tensor([embedded_guidance] * latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + return {"guidance": guidance} + + + +class FluxImageUnit_Kontext(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("kontext_images", "tiled", "tile_size", "tile_stride"), + output_params=("kontext_latents", "kontext_image_ids"), + onload_model_names=("vae_encoder",) + ) + + def process(self, pipe: FluxImagePipeline, kontext_images, tiled, tile_size, tile_stride): + if kontext_images is None: + return {} + if not isinstance(kontext_images, list): + kontext_images = [kontext_images] + + kontext_latents = [] + kontext_image_ids = [] + for kontext_image in kontext_images: + kontext_image = pipe.preprocess_image(kontext_image) + kontext_latent = pipe.vae_encoder(kontext_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + image_ids = pipe.dit.prepare_image_ids(kontext_latent) + image_ids[..., 0] = 1 + kontext_image_ids.append(image_ids) + kontext_latent = pipe.dit.patchify(kontext_latent) + kontext_latents.append(kontext_latent) + kontext_latents = torch.concat(kontext_latents, dim=1) + kontext_image_ids = torch.concat(kontext_image_ids, dim=-2) + return {"kontext_latents": kontext_latents, "kontext_image_ids": kontext_image_ids} + + + +class FluxImageUnit_ControlNet(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("controlnet_inputs", "tiled", "tile_size", "tile_stride"), + output_params=("controlnet_conditionings",), + onload_model_names=("vae_encoder",) + ) + + def apply_controlnet_mask_on_latents(self, pipe, latents, mask): + mask = (pipe.preprocess_image(mask) + 1) / 2 + mask = mask.mean(dim=1, keepdim=True) + mask = 1 - torch.nn.functional.interpolate(mask, size=latents.shape[-2:]) + latents = torch.concat([latents, mask], dim=1) + return latents + + def apply_controlnet_mask_on_image(self, pipe, image, mask): + mask = mask.resize(image.size) + mask = pipe.preprocess_image(mask).mean(dim=[0, 1]).cpu() + image = np.array(image) + image[mask > 0] = 0 + image = Image.fromarray(image) + return image + + def process(self, pipe: FluxImagePipeline, controlnet_inputs: list[ControlNetInput], tiled, tile_size, tile_stride): + if controlnet_inputs is None: + return {} + pipe.load_models_to_device(['vae_encoder']) + conditionings = [] + for controlnet_input in controlnet_inputs: + image = controlnet_input.image + if controlnet_input.inpaint_mask is not None: + image = self.apply_controlnet_mask_on_image(pipe, image, controlnet_input.inpaint_mask) + + image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) + image = pipe.vae_encoder(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + + if controlnet_input.inpaint_mask is not None: + image = self.apply_controlnet_mask_on_latents(pipe, image, controlnet_input.inpaint_mask) + conditionings.append(image) + return {"controlnet_conditionings": conditionings} + + + +class FluxImageUnit_IPAdapter(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("ipadapter_images", "ipadapter_scale"), + output_params=("ipadapter_kwargs_list",), + onload_model_names=("ipadapter_image_encoder", "ipadapter") + ) + + def process(self, pipe: FluxImagePipeline, inputs_shared, inputs_posi, inputs_nega): + ipadapter_images, ipadapter_scale = inputs_shared.get("ipadapter_images", None), inputs_shared.get("ipadapter_scale", 1.0) + if ipadapter_images is None: + return inputs_shared, inputs_posi, inputs_nega + if not isinstance(ipadapter_images, list): + ipadapter_images = [ipadapter_images] + + pipe.load_models_to_device(self.onload_model_names) + images = [image.convert("RGB").resize((384, 384), resample=3) for image in ipadapter_images] + images = [pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) for image in images] + ipadapter_images = torch.cat(images, dim=0) + ipadapter_image_encoding = pipe.ipadapter_image_encoder(ipadapter_images).pooler_output + + inputs_posi.update({"ipadapter_kwargs_list": pipe.ipadapter(ipadapter_image_encoding, scale=ipadapter_scale)}) + if inputs_shared.get("cfg_scale", 1.0) != 1.0: + inputs_nega.update({"ipadapter_kwargs_list": pipe.ipadapter(torch.zeros_like(ipadapter_image_encoding))}) + return inputs_shared, inputs_posi, inputs_nega + + + +class FluxImageUnit_EntityControl(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("eligen_entity_prompts", "eligen_entity_masks", "eligen_enable_on_negative", "width", "height", "t5_sequence_length", "cfg_scale"), + output_params=("entity_prompt_emb", "entity_masks"), + onload_model_names=("text_encoder_1", "text_encoder_2") + ) + + def encode_prompt_using_clip(self, prompt, text_encoder, tokenizer, max_length, device): + input_ids = tokenizer( + prompt, + return_tensors="pt", + padding="max_length", + max_length=max_length, + truncation=True + ).input_ids.to(device) + pooled_prompt_emb, _ = text_encoder(input_ids) + return pooled_prompt_emb + + def encode_prompt_using_t5(self, prompt, text_encoder, tokenizer, max_length, device): + input_ids = tokenizer( + prompt, + return_tensors="pt", + padding="max_length", + max_length=max_length, + truncation=True, + ).input_ids.to(device) + prompt_emb = text_encoder(input_ids) + return prompt_emb + + def encode_prompt( + self, + tokenizer_1, + tokenizer_2, + text_encoder_1, + text_encoder_2, + prompt, + positive=True, + device="cuda", + t5_sequence_length=512, + ): + pooled_prompt_emb = self.encode_prompt_using_clip(prompt, text_encoder_1, tokenizer_1, 77, device) + prompt_emb = self.encode_prompt_using_t5(prompt, text_encoder_2, tokenizer_2, t5_sequence_length, device) + text_ids = torch.zeros(prompt_emb.shape[0], prompt_emb.shape[1], 3).to(device=device, dtype=prompt_emb.dtype) + return prompt_emb, pooled_prompt_emb, text_ids + + def preprocess_masks(self, pipe, masks, height, width, dim): + out_masks = [] + for mask in masks: + mask = pipe.preprocess_image(mask.resize((width, height), resample=Image.NEAREST)).mean(dim=1, keepdim=True) > 0 + mask = mask.repeat(1, dim, 1, 1).to(device=pipe.device, dtype=pipe.torch_dtype) + out_masks.append(mask) + return out_masks + + def prepare_entity_inputs(self, pipe, entity_prompts, entity_masks, width, height, t5_sequence_length=512): + entity_masks = self.preprocess_masks(pipe, entity_masks, height//8, width//8, 1) + entity_masks = torch.cat(entity_masks, dim=0).unsqueeze(0) # b, n_mask, c, h, w + + prompt_emb, _, _ = self.encode_prompt( + tokenizer_1=pipe.tokenizer_1, tokenizer_2=pipe.tokenizer_2, + text_encoder_1=pipe.text_encoder_1, text_encoder_2=pipe.text_encoder_2, + prompt=entity_prompts, device=pipe.device, t5_sequence_length=t5_sequence_length, + ) + return prompt_emb.unsqueeze(0), entity_masks + + def prepare_eligen(self, pipe, prompt_emb_nega, eligen_entity_prompts, eligen_entity_masks, width, height, t5_sequence_length, enable_eligen_on_negative, cfg_scale): + entity_prompt_emb_posi, entity_masks_posi = self.prepare_entity_inputs(pipe, eligen_entity_prompts, eligen_entity_masks, width, height, t5_sequence_length) + if enable_eligen_on_negative and cfg_scale != 1.0: + entity_prompt_emb_nega = prompt_emb_nega['prompt_emb'].unsqueeze(1).repeat(1, entity_masks_posi.shape[1], 1, 1) + entity_masks_nega = entity_masks_posi + else: + entity_prompt_emb_nega, entity_masks_nega = None, None + eligen_kwargs_posi = {"entity_prompt_emb": entity_prompt_emb_posi, "entity_masks": entity_masks_posi} + eligen_kwargs_nega = {"entity_prompt_emb": entity_prompt_emb_nega, "entity_masks": entity_masks_nega} + return eligen_kwargs_posi, eligen_kwargs_nega + + def process(self, pipe: FluxImagePipeline, inputs_shared, inputs_posi, inputs_nega): + eligen_entity_prompts, eligen_entity_masks = inputs_shared.get("eligen_entity_prompts", None), inputs_shared.get("eligen_entity_masks", None) + if eligen_entity_prompts is None or eligen_entity_masks is None: + return inputs_shared, inputs_posi, inputs_nega + pipe.load_models_to_device(self.onload_model_names) + eligen_enable_on_negative = inputs_shared.get("eligen_enable_on_negative", False) + eligen_kwargs_posi, eligen_kwargs_nega = self.prepare_eligen(pipe, inputs_nega, + eligen_entity_prompts, eligen_entity_masks, inputs_shared["width"], inputs_shared["height"], + inputs_shared["t5_sequence_length"], eligen_enable_on_negative, inputs_shared["cfg_scale"]) + inputs_posi.update(eligen_kwargs_posi) + if inputs_shared.get("cfg_scale", 1.0) != 1.0: + inputs_nega.update(eligen_kwargs_nega) + return inputs_shared, inputs_posi, inputs_nega + + +class FluxImageUnit_NexusGen(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("nexus_gen_reference_image", "prompt", "latents"), + output_params=("prompt_emb", "text_ids"), + onload_model_names=("nexus_gen", "nexus_gen_generation_adapter", "nexus_gen_editing_adapter"), + ) + + def process(self, pipe: FluxImagePipeline, inputs_shared, inputs_posi, inputs_nega): + if pipe.nexus_gen is None: + return inputs_shared, inputs_posi, inputs_nega + pipe.load_models_to_device(self.onload_model_names) + if inputs_shared.get("nexus_gen_reference_image", None) is None: + assert pipe.nexus_gen_generation_adapter is not None, "NexusGen requires a generation adapter to be set." + embed = pipe.nexus_gen(inputs_posi["prompt"])[0].unsqueeze(0) + inputs_posi["prompt_emb"] = pipe.nexus_gen_generation_adapter(embed) + inputs_posi['text_ids'] = torch.zeros(embed.shape[0], embed.shape[1], 3).to(device=pipe.device, dtype=pipe.torch_dtype) + else: + assert pipe.nexus_gen_editing_adapter is not None, "NexusGen requires an editing adapter to be set." + embed, ref_embed, grids = pipe.nexus_gen(inputs_posi["prompt"], inputs_shared["nexus_gen_reference_image"]) + embeds_grid = grids[1:2].to(device=pipe.device, dtype=torch.long) + ref_embeds_grid = grids[0:1].to(device=pipe.device, dtype=torch.long) + + inputs_posi["prompt_emb"] = pipe.nexus_gen_editing_adapter(embed.unsqueeze(0), embeds_grid, ref_embed.unsqueeze(0), ref_embeds_grid) + inputs_posi["text_ids"] = self.get_editing_text_ids( + inputs_shared["latents"], + embeds_grid[0][1].item(), embeds_grid[0][2].item(), + ref_embeds_grid[0][1].item(), ref_embeds_grid[0][2].item(), + ) + return inputs_shared, inputs_posi, inputs_nega + + + def get_editing_text_ids(self, latents, target_embed_height, target_embed_width, ref_embed_height, ref_embed_width): + # prepare text ids for target and reference embeddings + batch_size, height, width = latents.shape[0], target_embed_height, target_embed_width + embed_ids = torch.zeros(height // 2, width // 2, 3) + scale_factor_height, scale_factor_width = latents.shape[-2] / height, latents.shape[-1] / width + embed_ids[..., 1] = embed_ids[..., 1] + torch.arange(height // 2)[:, None] * scale_factor_height + embed_ids[..., 2] = embed_ids[..., 2] + torch.arange(width // 2)[None, :] * scale_factor_width + embed_ids = embed_ids[None, :].repeat(batch_size, 1, 1, 1).reshape(batch_size, height // 2 * width // 2, 3) + embed_text_ids = embed_ids.to(device=latents.device, dtype=latents.dtype) + + batch_size, height, width = latents.shape[0], ref_embed_height, ref_embed_width + ref_embed_ids = torch.zeros(height // 2, width // 2, 3) + scale_factor_height, scale_factor_width = latents.shape[-2] / height, latents.shape[-1] / width + ref_embed_ids[..., 0] = ref_embed_ids[..., 0] + 1.0 + ref_embed_ids[..., 1] = ref_embed_ids[..., 1] + torch.arange(height // 2)[:, None] * scale_factor_height + ref_embed_ids[..., 2] = ref_embed_ids[..., 2] + torch.arange(width // 2)[None, :] * scale_factor_width + ref_embed_ids = ref_embed_ids[None, :].repeat(batch_size, 1, 1, 1).reshape(batch_size, height // 2 * width // 2, 3) + ref_embed_text_ids = ref_embed_ids.to(device=latents.device, dtype=latents.dtype) + + text_ids = torch.cat([embed_text_ids, ref_embed_text_ids], dim=1) + return text_ids + + +class FluxImageUnit_Step1x(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("step1x_reference_image", "prompt", "negative_prompt"), + output_params=("step1x_llm_embedding", "step1x_mask", "step1x_reference_latents"), + onload_model_names=("qwenvl","vae_encoder") + ) + + def process(self, pipe: FluxImagePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict): + image = inputs_shared.get("step1x_reference_image",None) + if image is None: + return inputs_shared, inputs_posi, inputs_nega + else: + pipe.load_models_to_device(self.onload_model_names) + prompt = inputs_posi["prompt"] + nega_prompt = inputs_nega["negative_prompt"] + captions = [prompt, nega_prompt] + ref_images = [image, image] + embs, masks = pipe.qwenvl(captions, ref_images) + image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) + image = pipe.vae_encoder(image) + inputs_posi.update({"step1x_llm_embedding": embs[0:1], "step1x_mask": masks[0:1], "step1x_reference_latents": image}) + if inputs_shared.get("cfg_scale", 1) != 1: + inputs_nega.update({"step1x_llm_embedding": embs[1:2], "step1x_mask": masks[1:2], "step1x_reference_latents": image}) + return inputs_shared, inputs_posi, inputs_nega + + +class FluxImageUnit_TeaCache(PipelineUnit): + def __init__(self): + super().__init__(input_params=("num_inference_steps","tea_cache_l1_thresh"), output_params=("tea_cache",)) + + def process(self, pipe: FluxImagePipeline, num_inference_steps, tea_cache_l1_thresh): + if tea_cache_l1_thresh is None: + return {} + else: + return {"tea_cache": TeaCache(num_inference_steps=num_inference_steps, rel_l1_thresh=tea_cache_l1_thresh)} + +class FluxImageUnit_Flex(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("latents", "flex_inpaint_image", "flex_inpaint_mask", "flex_control_image", "flex_control_strength", "flex_control_stop", "tiled", "tile_size", "tile_stride"), + output_params=("flex_condition", "flex_uncondition", "flex_control_stop_timestep"), + onload_model_names=("vae_encoder",) + ) + + def process(self, pipe: FluxImagePipeline, latents, flex_inpaint_image, flex_inpaint_mask, flex_control_image, flex_control_strength, flex_control_stop, tiled, tile_size, tile_stride): + if pipe.dit.input_dim == 196: + if flex_control_stop is None: + flex_control_stop = 1 + pipe.load_models_to_device(self.onload_model_names) + if flex_inpaint_image is None: + flex_inpaint_image = torch.zeros_like(latents) + else: + flex_inpaint_image = pipe.preprocess_image(flex_inpaint_image).to(device=pipe.device, dtype=pipe.torch_dtype) + flex_inpaint_image = pipe.vae_encoder(flex_inpaint_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + if flex_inpaint_mask is None: + flex_inpaint_mask = torch.ones_like(latents)[:, 0:1, :, :] + else: + flex_inpaint_mask = flex_inpaint_mask.resize((latents.shape[3], latents.shape[2])) + flex_inpaint_mask = pipe.preprocess_image(flex_inpaint_mask).to(device=pipe.device, dtype=pipe.torch_dtype) + flex_inpaint_mask = (flex_inpaint_mask[:, 0:1, :, :] + 1) / 2 + flex_inpaint_image = flex_inpaint_image * (1 - flex_inpaint_mask) + if flex_control_image is None: + flex_control_image = torch.zeros_like(latents) + else: + flex_control_image = pipe.preprocess_image(flex_control_image).to(device=pipe.device, dtype=pipe.torch_dtype) + flex_control_image = pipe.vae_encoder(flex_control_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) * flex_control_strength + flex_condition = torch.concat([flex_inpaint_image, flex_inpaint_mask, flex_control_image], dim=1) + flex_uncondition = torch.concat([flex_inpaint_image, flex_inpaint_mask, torch.zeros_like(flex_control_image)], dim=1) + flex_control_stop_timestep = pipe.scheduler.timesteps[int(flex_control_stop * (len(pipe.scheduler.timesteps) - 1))] + return {"flex_condition": flex_condition, "flex_uncondition": flex_uncondition, "flex_control_stop_timestep": flex_control_stop_timestep} + else: + return {} + + + +class FluxImageUnit_InfiniteYou(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("infinityou_id_image", "infinityou_guidance"), + output_params=("id_emb", "infinityou_guidance"), + onload_model_names=("infinityou_processor",) + ) + + def process(self, pipe: FluxImagePipeline, infinityou_id_image, infinityou_guidance): + pipe.load_models_to_device("infinityou_processor") + if infinityou_id_image is not None: + return pipe.infinityou_processor.prepare_infinite_you(pipe.image_proj_model, infinityou_id_image, infinityou_guidance, pipe.device) + else: + return {} + + + +class FluxImageUnit_ValueControl(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt_emb": "prompt_emb", "text_ids": "text_ids"}, + input_params_nega={"prompt_emb": "prompt_emb", "text_ids": "text_ids"}, + input_params=("value_controller_inputs",), + output_params=("prompt_emb", "text_ids"), + onload_model_names=("value_controller",) + ) + + def add_to_text_embedding(self, prompt_emb, text_ids, value_emb): + prompt_emb = torch.concat([prompt_emb, value_emb], dim=1) + extra_text_ids = torch.zeros((value_emb.shape[0], value_emb.shape[1], 3), device=value_emb.device, dtype=value_emb.dtype) + text_ids = torch.concat([text_ids, extra_text_ids], dim=1) + return prompt_emb, text_ids + + def process(self, pipe: FluxImagePipeline, prompt_emb, text_ids, value_controller_inputs): + if value_controller_inputs is None: + return {} + if not isinstance(value_controller_inputs, list): + value_controller_inputs = [value_controller_inputs] + value_controller_inputs = torch.tensor(value_controller_inputs).to(dtype=pipe.torch_dtype, device=pipe.device) + pipe.load_models_to_device(["value_controller"]) + value_emb = pipe.value_controller(value_controller_inputs, pipe.torch_dtype) + value_emb = value_emb.unsqueeze(0) + prompt_emb, text_ids = self.add_to_text_embedding(prompt_emb, text_ids, value_emb) + return {"prompt_emb": prompt_emb, "text_ids": text_ids} + + + +class InfinitYou(torch.nn.Module): + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__() + from facexlib.recognition import init_recognition_model + from insightface.app import FaceAnalysis + self.device = device + self.torch_dtype = torch_dtype + insightface_root_path = 'models/ByteDance/InfiniteYou/supports/insightface' + self.app_640 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) + self.app_640.prepare(ctx_id=0, det_size=(640, 640)) + self.app_320 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) + self.app_320.prepare(ctx_id=0, det_size=(320, 320)) + self.app_160 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) + self.app_160.prepare(ctx_id=0, det_size=(160, 160)) + self.arcface_model = init_recognition_model('arcface', device=self.device).to(torch_dtype) + + def _detect_face(self, id_image_cv2): + face_info = self.app_640.get(id_image_cv2) + if len(face_info) > 0: + return face_info + face_info = self.app_320.get(id_image_cv2) + if len(face_info) > 0: + return face_info + face_info = self.app_160.get(id_image_cv2) + return face_info + + def extract_arcface_bgr_embedding(self, in_image, landmark, device): + from insightface.utils import face_align + arc_face_image = face_align.norm_crop(in_image, landmark=np.array(landmark), image_size=112) + arc_face_image = torch.from_numpy(arc_face_image).unsqueeze(0).permute(0, 3, 1, 2) / 255. + arc_face_image = 2 * arc_face_image - 1 + arc_face_image = arc_face_image.contiguous().to(device=device, dtype=self.torch_dtype) + face_emb = self.arcface_model(arc_face_image)[0] # [512], normalized + return face_emb + + def prepare_infinite_you(self, model, id_image, infinityou_guidance, device): + import cv2 + if id_image is None: + return {'id_emb': None} + id_image_cv2 = cv2.cvtColor(np.array(id_image), cv2.COLOR_RGB2BGR) + face_info = self._detect_face(id_image_cv2) + if len(face_info) == 0: + raise ValueError('No face detected in the input ID image') + landmark = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1]['kps'] # only use the maximum face + id_emb = self.extract_arcface_bgr_embedding(id_image_cv2, landmark, device) + id_emb = model(id_emb.unsqueeze(0).reshape([1, -1, 512]).to(dtype=self.torch_dtype)) + infinityou_guidance = torch.Tensor([infinityou_guidance]).to(device=device, dtype=self.torch_dtype) + return {'id_emb': id_emb, 'infinityou_guidance': infinityou_guidance} + + + +class FluxImageUnit_LoRAEncode(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("lora_encoder_inputs", "lora_encoder_scale"), + output_params=("prompt_emb", "text_ids"), + onload_model_names=("lora_encoder",) + ) + + def parse_lora_encoder_inputs(self, lora_encoder_inputs): + if not isinstance(lora_encoder_inputs, list): + lora_encoder_inputs = [lora_encoder_inputs] + lora_configs = [] + for lora_encoder_input in lora_encoder_inputs: + if isinstance(lora_encoder_input, str): + lora_encoder_input = ModelConfig(path=lora_encoder_input) + lora_encoder_input.download_if_necessary() + lora_configs.append(lora_encoder_input) + return lora_configs + + def load_lora(self, lora_config, dtype, device): + loader = FluxLoRALoader(torch_dtype=dtype, device=device) + lora = load_state_dict(lora_config.path, torch_dtype=dtype, device=device) + lora = loader.convert_state_dict(lora) + return lora + + def lora_embedding(self, pipe, lora_encoder_inputs): + lora_emb = [] + for lora_config in self.parse_lora_encoder_inputs(lora_encoder_inputs): + lora = self.load_lora(lora_config, pipe.torch_dtype, pipe.device) + lora_emb.append(pipe.lora_encoder(lora)) + lora_emb = torch.concat(lora_emb, dim=1) + return lora_emb + + def add_to_text_embedding(self, prompt_emb, text_ids, lora_emb): + prompt_emb = torch.concat([prompt_emb, lora_emb], dim=1) + extra_text_ids = torch.zeros((lora_emb.shape[0], lora_emb.shape[1], 3), device=lora_emb.device, dtype=lora_emb.dtype) + text_ids = torch.concat([text_ids, extra_text_ids], dim=1) + return prompt_emb, text_ids + + def process(self, pipe: FluxImagePipeline, inputs_shared, inputs_posi, inputs_nega): + if inputs_shared.get("lora_encoder_inputs", None) is None: + return inputs_shared, inputs_posi, inputs_nega + + # Encode + pipe.load_models_to_device(["lora_encoder"]) + lora_encoder_inputs = inputs_shared["lora_encoder_inputs"] + lora_emb = self.lora_embedding(pipe, lora_encoder_inputs) + + # Scale + lora_encoder_scale = inputs_shared.get("lora_encoder_scale", None) + if lora_encoder_scale is not None: + lora_emb = lora_emb * lora_encoder_scale + + # Add to prompt embedding + inputs_posi["prompt_emb"], inputs_posi["text_ids"] = self.add_to_text_embedding( + inputs_posi["prompt_emb"], inputs_posi["text_ids"], lora_emb) + return inputs_shared, inputs_posi, inputs_nega + + + +class TeaCache: + def __init__(self, num_inference_steps, rel_l1_thresh): + self.num_inference_steps = num_inference_steps + self.step = 0 + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.rel_l1_thresh = rel_l1_thresh + self.previous_residual = None + self.previous_hidden_states = None + + def check(self, dit: FluxDiT, hidden_states, conditioning): + inp = hidden_states.clone() + temb_ = conditioning.clone() + modulated_inp, _, _, _, _ = dit.blocks[0].norm1_a(inp, emb=temb_) + if self.step == 0 or self.step == self.num_inference_steps - 1: + should_calc = True + self.accumulated_rel_l1_distance = 0 + else: + coefficients = [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01] + rescale_func = np.poly1d(coefficients) + self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()) + if self.accumulated_rel_l1_distance < self.rel_l1_thresh: + should_calc = False + else: + should_calc = True + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = modulated_inp + self.step += 1 + if self.step == self.num_inference_steps: + self.step = 0 + if should_calc: + self.previous_hidden_states = hidden_states.clone() + return not should_calc + + def store(self, hidden_states): + self.previous_residual = hidden_states - self.previous_hidden_states + self.previous_hidden_states = None + + def update(self, hidden_states): + hidden_states = hidden_states + self.previous_residual + return hidden_states + + +class FastTileWorker: + def __init__(self): + pass + + + def build_mask(self, data, is_bound): + _, _, H, W = data.shape + h = repeat(torch.arange(H), "H -> H W", H=H, W=W) + w = repeat(torch.arange(W), "W -> H W", H=H, W=W) + border_width = (H + W) // 4 + pad = torch.ones_like(h) * border_width + mask = torch.stack([ + pad if is_bound[0] else h + 1, + pad if is_bound[1] else H - h, + pad if is_bound[2] else w + 1, + pad if is_bound[3] else W - w + ]).min(dim=0).values + mask = mask.clip(1, border_width) + mask = (mask / border_width).to(dtype=data.dtype, device=data.device) + mask = rearrange(mask, "H W -> 1 H W") + return mask + + + def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_device="cpu", tile_dtype=torch.float32, border_width=None): + # Prepare + B, C, H, W = model_input.shape + border_width = int(tile_stride*0.5) if border_width is None else border_width + weight = torch.zeros((1, 1, H, W), dtype=tile_dtype, device=tile_device) + values = torch.zeros((B, C, H, W), dtype=tile_dtype, device=tile_device) + + # Split tasks + tasks = [] + for h in range(0, H, tile_stride): + for w in range(0, W, tile_stride): + if (h-tile_stride >= 0 and h-tile_stride+tile_size >= H) or (w-tile_stride >= 0 and w-tile_stride+tile_size >= W): + continue + h_, w_ = h + tile_size, w + tile_size + if h_ > H: h, h_ = H - tile_size, H + if w_ > W: w, w_ = W - tile_size, W + tasks.append((h, h_, w, w_)) + + # Run + for hl, hr, wl, wr in tasks: + # Forward + hidden_states_batch = forward_fn(hl, hr, wl, wr).to(dtype=tile_dtype, device=tile_device) + + mask = self.build_mask(hidden_states_batch, is_bound=(hl==0, hr>=H, wl==0, wr>=W)) + values[:, :, hl:hr, wl:wr] += hidden_states_batch * mask + weight[:, :, hl:hr, wl:wr] += mask + values /= weight + return values + + +def model_fn_flux_image( + dit: FluxDiT, + controlnet=None, + step1x_connector=None, + latents=None, + timestep=None, + prompt_emb=None, + pooled_prompt_emb=None, + guidance=None, + text_ids=None, + image_ids=None, + kontext_latents=None, + kontext_image_ids=None, + controlnet_inputs=None, + controlnet_conditionings=None, + tiled=False, + tile_size=128, + tile_stride=64, + entity_prompt_emb=None, + entity_masks=None, + ipadapter_kwargs_list={}, + id_emb=None, + infinityou_guidance=None, + flex_condition=None, + flex_uncondition=None, + flex_control_stop_timestep=None, + step1x_llm_embedding=None, + step1x_mask=None, + step1x_reference_latents=None, + tea_cache: TeaCache = None, + progress_id=0, + num_inference_steps=1, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + **kwargs +): + if tiled: + def flux_forward_fn(hl, hr, wl, wr): + tiled_controlnet_conditionings = [f[:, :, hl: hr, wl: wr] for f in controlnet_conditionings] if controlnet_conditionings is not None else None + return model_fn_flux_image( + dit=dit, + controlnet=controlnet, + latents=latents[:, :, hl: hr, wl: wr], + timestep=timestep, + prompt_emb=prompt_emb, + pooled_prompt_emb=pooled_prompt_emb, + guidance=guidance, + text_ids=text_ids, + image_ids=None, + controlnet_inputs=controlnet_inputs, + controlnet_conditionings=tiled_controlnet_conditionings, + tiled=False, + **kwargs + ) + return FastTileWorker().tiled_forward( + flux_forward_fn, + latents, + tile_size=tile_size, + tile_stride=tile_stride, + tile_device=latents.device, + tile_dtype=latents.dtype + ) + + hidden_states = latents + + # ControlNet + if controlnet is not None and controlnet_conditionings is not None: + controlnet_extra_kwargs = { + "hidden_states": hidden_states, + "timestep": timestep, + "prompt_emb": prompt_emb, + "pooled_prompt_emb": pooled_prompt_emb, + "guidance": guidance, + "text_ids": text_ids, + "image_ids": image_ids, + "controlnet_inputs": controlnet_inputs, + "tiled": tiled, + "tile_size": tile_size, + "tile_stride": tile_stride, + "progress_id": progress_id, + "num_inference_steps": num_inference_steps, + } + if id_emb is not None: + controlnet_text_ids = torch.zeros(id_emb.shape[0], id_emb.shape[1], 3).to(device=hidden_states.device, dtype=hidden_states.dtype) + controlnet_extra_kwargs.update({"prompt_emb": id_emb, 'text_ids': controlnet_text_ids, 'guidance': infinityou_guidance}) + controlnet_res_stack, controlnet_single_res_stack = controlnet( + controlnet_conditionings, **controlnet_extra_kwargs + ) + + # Flex + if flex_condition is not None: + if timestep.tolist()[0] >= flex_control_stop_timestep: + hidden_states = torch.concat([hidden_states, flex_condition], dim=1) + else: + hidden_states = torch.concat([hidden_states, flex_uncondition], dim=1) + + # Step1x + if step1x_llm_embedding is not None: + prompt_emb, pooled_prompt_emb = step1x_connector(step1x_llm_embedding, timestep / 1000, step1x_mask) + text_ids = torch.zeros((1, prompt_emb.shape[1], 3), dtype=prompt_emb.dtype, device=prompt_emb.device) + + if image_ids is None: + image_ids = dit.prepare_image_ids(hidden_states) + + conditioning = dit.time_embedder(timestep, hidden_states.dtype) + dit.pooled_text_embedder(pooled_prompt_emb) + if dit.guidance_embedder is not None: + guidance = guidance * 1000 + conditioning = conditioning + dit.guidance_embedder(guidance, hidden_states.dtype) + + height, width = hidden_states.shape[-2:] + hidden_states = dit.patchify(hidden_states) + + # Kontext + if kontext_latents is not None: + image_ids = torch.concat([image_ids, kontext_image_ids], dim=-2) + hidden_states = torch.concat([hidden_states, kontext_latents], dim=1) + + # Step1x + if step1x_reference_latents is not None: + step1x_reference_image_ids = dit.prepare_image_ids(step1x_reference_latents) + step1x_reference_latents = dit.patchify(step1x_reference_latents) + image_ids = torch.concat([image_ids, step1x_reference_image_ids], dim=-2) + hidden_states = torch.concat([hidden_states, step1x_reference_latents], dim=1) + + hidden_states = dit.x_embedder(hidden_states) + + # EliGen + if entity_prompt_emb is not None and entity_masks is not None: + prompt_emb, image_rotary_emb, attention_mask = dit.process_entity_masks(hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids, latents.shape[1]) + else: + prompt_emb = dit.context_embedder(prompt_emb) + image_rotary_emb = dit.pos_embedder(torch.cat((text_ids, image_ids), dim=1)) + attention_mask = None + + # TeaCache + if tea_cache is not None: + tea_cache_update = tea_cache.check(dit, hidden_states, conditioning) + else: + tea_cache_update = False + + if tea_cache_update: + hidden_states = tea_cache.update(hidden_states) + else: + # Joint Blocks + for block_id, block in enumerate(dit.blocks): + hidden_states, prompt_emb = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + hidden_states, + prompt_emb, + conditioning, + image_rotary_emb, + attention_mask, + ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id, None), + ) + # ControlNet + if controlnet is not None and controlnet_conditionings is not None and controlnet_res_stack is not None: + if kontext_latents is None: + hidden_states = hidden_states + controlnet_res_stack[block_id] + else: + hidden_states[:, :-kontext_latents.shape[1]] = hidden_states[:, :-kontext_latents.shape[1]] + controlnet_res_stack[block_id] + + # Single Blocks + hidden_states = torch.cat([prompt_emb, hidden_states], dim=1) + num_joint_blocks = len(dit.blocks) + for block_id, block in enumerate(dit.single_blocks): + hidden_states, prompt_emb = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + hidden_states, + prompt_emb, + conditioning, + image_rotary_emb, + attention_mask, + ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id + num_joint_blocks, None), + ) + # ControlNet + if controlnet is not None and controlnet_conditionings is not None and controlnet_single_res_stack is not None: + if kontext_latents is None: + hidden_states[:, prompt_emb.shape[1]:] = hidden_states[:, prompt_emb.shape[1]:] + controlnet_single_res_stack[block_id] + else: + hidden_states[:, prompt_emb.shape[1]:-kontext_latents.shape[1]] = hidden_states[:, prompt_emb.shape[1]:-kontext_latents.shape[1]] + controlnet_single_res_stack[block_id] + hidden_states = hidden_states[:, prompt_emb.shape[1]:] + + if tea_cache is not None: + tea_cache.store(hidden_states) + + hidden_states = dit.final_norm_out(hidden_states, conditioning) + hidden_states = dit.final_proj_out(hidden_states) + + # Step1x + if step1x_reference_latents is not None: + hidden_states = hidden_states[:, :hidden_states.shape[1] // 2] + + # Kontext + if kontext_latents is not None: + hidden_states = hidden_states[:, :-kontext_latents.shape[1]] + + hidden_states = dit.unpatchify(hidden_states, height, width) + + return hidden_states diff --git a/diffsynth/pipelines/qwen_image.py b/diffsynth/pipelines/qwen_image.py new file mode 100644 index 0000000000000000000000000000000000000000..fc6581e5e27d17c3fea831b00a51778719e3b1ce --- /dev/null +++ b/diffsynth/pipelines/qwen_image.py @@ -0,0 +1,622 @@ +import torch, math +from PIL import Image +from typing import Union +from tqdm import tqdm +from einops import rearrange +import numpy as np + +from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig, gradient_checkpoint_forward +from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput + +from ..models.qwen_image_dit import QwenImageDiT +from ..models.qwen_image_text_encoder import QwenImageTextEncoder +from ..models.qwen_image_vae import QwenImageVAE +from ..models.qwen_image_controlnet import QwenImageBlockWiseControlNet + + +class QwenImagePipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, + ) + from transformers import Qwen2Tokenizer, Qwen2VLProcessor + + self.scheduler = FlowMatchScheduler("Qwen-Image") + self.text_encoder: QwenImageTextEncoder = None + self.dit: QwenImageDiT = None + self.vae: QwenImageVAE = None + self.blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None + self.tokenizer: Qwen2Tokenizer = None + self.processor: Qwen2VLProcessor = None + self.in_iteration_models = ("dit", "blockwise_controlnet") + self.units = [ + QwenImageUnit_ShapeChecker(), + QwenImageUnit_NoiseInitializer(), + QwenImageUnit_InputImageEmbedder(), + QwenImageUnit_Inpaint(), + QwenImageUnit_EditImageEmbedder(), + QwenImageUnit_ContextImageEmbedder(), + QwenImageUnit_PromptEmbedder(), + QwenImageUnit_EntityControl(), + QwenImageUnit_BlockwiseControlNet(), + ] + self.model_fn = model_fn_qwen_image + + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_config: ModelConfig = ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), + processor_config: ModelConfig = None, + vram_limit: float = None, + ): + # Initialize pipeline + pipe = QwenImagePipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder = model_pool.fetch_model("qwen_image_text_encoder") + pipe.dit = model_pool.fetch_model("qwen_image_dit") + pipe.vae = model_pool.fetch_model("qwen_image_vae") + pipe.blockwise_controlnet = QwenImageBlockwiseMultiControlNet(model_pool.fetch_model("qwen_image_blockwise_controlnet", index="all")) + if tokenizer_config is not None: + tokenizer_config.download_if_necessary() + from transformers import Qwen2Tokenizer + pipe.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_config.path) + if processor_config is not None: + processor_config.download_if_necessary() + from transformers import Qwen2VLProcessor + pipe.processor = Qwen2VLProcessor.from_pretrained(processor_config.path) + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: str = "", + cfg_scale: float = 4.0, + # Image + input_image: Image.Image = None, + denoising_strength: float = 1.0, + # Inpaint + inpaint_mask: Image.Image = None, + inpaint_blur_size: int = None, + inpaint_blur_sigma: float = None, + # Shape + height: int = 1328, + width: int = 1328, + # Randomness + seed: int = None, + rand_device: str = "cpu", + # Steps + num_inference_steps: int = 30, + exponential_shift_mu: float = None, + # Blockwise ControlNet + blockwise_controlnet_inputs: list[ControlNetInput] = None, + # EliGen + eligen_entity_prompts: list[str] = None, + eligen_entity_masks: list[Image.Image] = None, + eligen_enable_on_negative: bool = False, + # Qwen-Image-Edit + edit_image: Image.Image = None, + edit_image_auto_resize: bool = True, + edit_rope_interpolation: bool = False, + # In-context control + context_image: Image.Image = None, + # Tile + tiled: bool = False, + tile_size: int = 128, + tile_stride: int = 64, + # Progress bar + progress_bar_cmd = tqdm, + ): + # Scheduler + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, dynamic_shift_len=(height // 16) * (width // 16), exponential_shift_mu=exponential_shift_mu) + + # Parameters + inputs_posi = { + "prompt": prompt, + } + inputs_nega = { + "negative_prompt": negative_prompt, + } + inputs_shared = { + "cfg_scale": cfg_scale, + "input_image": input_image, "denoising_strength": denoising_strength, + "inpaint_mask": inpaint_mask, "inpaint_blur_size": inpaint_blur_size, "inpaint_blur_sigma": inpaint_blur_sigma, + "height": height, "width": width, + "seed": seed, "rand_device": rand_device, + "num_inference_steps": num_inference_steps, + "blockwise_controlnet_inputs": blockwise_controlnet_inputs, + "tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride, + "eligen_entity_prompts": eligen_entity_prompts, "eligen_entity_masks": eligen_entity_masks, "eligen_enable_on_negative": eligen_enable_on_negative, + "edit_image": edit_image, "edit_image_auto_resize": edit_image_auto_resize, "edit_rope_interpolation": edit_rope_interpolation, + "context_image": context_image, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + noise_pred = self.cfg_guided_model_fn( + self.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) + + # Decode + self.load_models_to_device(['vae']) + image = self.vae.decode(inputs_shared["latents"], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + image = self.vae_output_to_image(image) + self.load_models_to_device([]) + + return image + + +class QwenImageBlockwiseMultiControlNet(torch.nn.Module): + def __init__(self, models: list[QwenImageBlockWiseControlNet]): + super().__init__() + if not isinstance(models, list): + models = [models] + self.models = torch.nn.ModuleList(models) + for model in models: + if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"): + self.vram_management_enabled = True + + def preprocess(self, controlnet_inputs: list[ControlNetInput], conditionings: list[torch.Tensor], **kwargs): + processed_conditionings = [] + for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): + conditioning = rearrange(conditioning, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) + model_output = self.models[controlnet_input.controlnet_id].process_controlnet_conditioning(conditioning) + processed_conditionings.append(model_output) + return processed_conditionings + + def blockwise_forward(self, image, conditionings: list[torch.Tensor], controlnet_inputs: list[ControlNetInput], progress_id, num_inference_steps, block_id, **kwargs): + res = 0 + for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): + progress = (num_inference_steps - 1 - progress_id) / max(num_inference_steps - 1, 1) + if progress > controlnet_input.start + (1e-4) or progress < controlnet_input.end - (1e-4): + continue + model_output = self.models[controlnet_input.controlnet_id].blockwise_forward(image, conditioning, block_id) + res = res + model_output * controlnet_input.scale + return res + + +class QwenImageUnit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width"), + output_params=("height", "width"), + ) + + def process(self, pipe: QwenImagePipeline, height, width): + height, width = pipe.check_resize_height_width(height, width) + return {"height": height, "width": width} + + + +class QwenImageUnit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "seed", "rand_device"), + output_params=("noise",), + ) + + def process(self, pipe: QwenImagePipeline, height, width, seed, rand_device): + noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype) + return {"noise": noise} + + + +class QwenImageUnit_InputImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "noise", "tiled", "tile_size", "tile_stride"), + output_params=("latents", "input_latents"), + onload_model_names=("vae",) + ) + + def process(self, pipe: QwenImagePipeline, input_image, noise, tiled, tile_size, tile_stride): + if input_image is None: + return {"latents": noise, "input_latents": None} + pipe.load_models_to_device(['vae']) + image = pipe.preprocess_image(input_image).to(device=pipe.device, dtype=pipe.torch_dtype) + input_latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents, "input_latents": input_latents} + + + +class QwenImageUnit_Inpaint(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("inpaint_mask", "height", "width", "inpaint_blur_size", "inpaint_blur_sigma"), + output_params=("inpaint_mask",), + ) + + def process(self, pipe: QwenImagePipeline, inpaint_mask, height, width, inpaint_blur_size, inpaint_blur_sigma): + if inpaint_mask is None: + return {} + inpaint_mask = pipe.preprocess_image(inpaint_mask.convert("RGB").resize((width // 8, height // 8)), min_value=0, max_value=1) + inpaint_mask = inpaint_mask.mean(dim=1, keepdim=True) + if inpaint_blur_size is not None and inpaint_blur_sigma is not None: + from torchvision.transforms import GaussianBlur + blur = GaussianBlur(kernel_size=inpaint_blur_size * 2 + 1, sigma=inpaint_blur_sigma) + inpaint_mask = blur(inpaint_mask) + return {"inpaint_mask": inpaint_mask} + + +class QwenImageUnit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt"}, + input_params_nega={"prompt": "negative_prompt"}, + input_params=("edit_image",), + output_params=("prompt_emb", "prompt_emb_mask"), + onload_model_names=("text_encoder",) + ) + + def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + return split_result + + def calculate_dimensions(self, target_area, ratio): + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + def resize_image(self, image, target_area=384*384): + width, height = self.calculate_dimensions(target_area, image.size[0] / image.size[1]) + return image.resize((width, height)) + + def encode_prompt(self, pipe: QwenImagePipeline, prompt): + template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 34 + txt = [template.format(e) for e in prompt] + model_inputs = pipe.tokenizer(txt, max_length=4096+drop_idx, padding=True, truncation=True, return_tensors="pt").to(pipe.device) + if model_inputs.input_ids.shape[1] >= 1024: + print(f"Warning!!! QwenImage model was trained on prompts up to 512 tokens. Current prompt requires {model_inputs['input_ids'].shape[1] - drop_idx} tokens, which may lead to unpredictable behavior.") + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def encode_prompt_edit(self, pipe: QwenImagePipeline, prompt, edit_image): + template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 64 + txt = [template.format(e) for e in prompt] + model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def encode_prompt_edit_multi(self, pipe: QwenImagePipeline, prompt, edit_image): + template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 64 + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + base_img_prompt = "".join([img_prompt_template.format(i + 1) for i in range(len(edit_image))]) + txt = [template.format(base_img_prompt + e) for e in prompt] + edit_image = [self.resize_image(image) for image in edit_image] + model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def process(self, pipe: QwenImagePipeline, prompt, edit_image=None) -> dict: + pipe.load_models_to_device(self.onload_model_names) + if pipe.text_encoder is not None: + prompt = [prompt] + if edit_image is None: + split_hidden_states = self.encode_prompt(pipe, prompt) + elif isinstance(edit_image, Image.Image): + split_hidden_states = self.encode_prompt_edit(pipe, prompt, edit_image) + else: + split_hidden_states = self.encode_prompt_edit_multi(pipe, prompt, edit_image) + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) + encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) + prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask} + else: + return {} + + +class QwenImageUnit_EntityControl(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("eligen_entity_prompts", "width", "height", "eligen_enable_on_negative", "cfg_scale"), + output_params=("entity_prompt_emb", "entity_masks", "entity_prompt_emb_mask"), + onload_model_names=("text_encoder",) + ) + + def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + return split_result + + def get_prompt_emb(self, pipe: QwenImagePipeline, prompt) -> dict: + if pipe.text_encoder is not None: + prompt = [prompt] + template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 34 + txt = [template.format(e) for e in prompt] + txt_tokens = pipe.tokenizer(txt, max_length=1024+drop_idx, padding=True, truncation=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, output_hidden_states=True,)[-1] + + split_hidden_states = self.extract_masked_hidden(hidden_states, txt_tokens.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) + encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) + prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask} + else: + return {} + + def preprocess_masks(self, pipe, masks, height, width, dim): + out_masks = [] + for mask in masks: + mask = pipe.preprocess_image(mask.resize((width, height), resample=Image.NEAREST)).mean(dim=1, keepdim=True) > 0 + mask = mask.repeat(1, dim, 1, 1).to(device=pipe.device, dtype=pipe.torch_dtype) + out_masks.append(mask) + return out_masks + + def prepare_entity_inputs(self, pipe, entity_prompts, entity_masks, width, height): + entity_masks = self.preprocess_masks(pipe, entity_masks, height//8, width//8, 1) + entity_masks = torch.cat(entity_masks, dim=0).unsqueeze(0) # b, n_mask, c, h, w + prompt_embs, prompt_emb_masks = [], [] + for entity_prompt in entity_prompts: + prompt_emb_dict = self.get_prompt_emb(pipe, entity_prompt) + prompt_embs.append(prompt_emb_dict['prompt_emb']) + prompt_emb_masks.append(prompt_emb_dict['prompt_emb_mask']) + return prompt_embs, prompt_emb_masks, entity_masks + + def prepare_eligen(self, pipe, prompt_emb_nega, eligen_entity_prompts, eligen_entity_masks, width, height, enable_eligen_on_negative, cfg_scale): + entity_prompt_emb_posi, entity_prompt_emb_posi_mask, entity_masks_posi = self.prepare_entity_inputs(pipe, eligen_entity_prompts, eligen_entity_masks, width, height) + if enable_eligen_on_negative and cfg_scale != 1.0: + entity_prompt_emb_nega = [prompt_emb_nega['prompt_emb']] * len(entity_prompt_emb_posi) + entity_prompt_emb_nega_mask = [prompt_emb_nega['prompt_emb_mask']] * len(entity_prompt_emb_posi) + entity_masks_nega = entity_masks_posi + else: + entity_prompt_emb_nega, entity_prompt_emb_nega_mask, entity_masks_nega = None, None, None + eligen_kwargs_posi = {"entity_prompt_emb": entity_prompt_emb_posi, "entity_masks": entity_masks_posi, "entity_prompt_emb_mask": entity_prompt_emb_posi_mask} + eligen_kwargs_nega = {"entity_prompt_emb": entity_prompt_emb_nega, "entity_masks": entity_masks_nega, "entity_prompt_emb_mask": entity_prompt_emb_nega_mask} + return eligen_kwargs_posi, eligen_kwargs_nega + + def process(self, pipe: QwenImagePipeline, inputs_shared, inputs_posi, inputs_nega): + eligen_entity_prompts, eligen_entity_masks = inputs_shared.get("eligen_entity_prompts", None), inputs_shared.get("eligen_entity_masks", None) + if eligen_entity_prompts is None or eligen_entity_masks is None or len(eligen_entity_prompts) == 0 or len(eligen_entity_masks) == 0: + return inputs_shared, inputs_posi, inputs_nega + pipe.load_models_to_device(self.onload_model_names) + eligen_enable_on_negative = inputs_shared.get("eligen_enable_on_negative", False) + eligen_kwargs_posi, eligen_kwargs_nega = self.prepare_eligen(pipe, inputs_nega, + eligen_entity_prompts, eligen_entity_masks, inputs_shared["width"], inputs_shared["height"], + eligen_enable_on_negative, inputs_shared["cfg_scale"]) + inputs_posi.update(eligen_kwargs_posi) + if inputs_shared.get("cfg_scale", 1.0) != 1.0: + inputs_nega.update(eligen_kwargs_nega) + return inputs_shared, inputs_posi, inputs_nega + + + +class QwenImageUnit_BlockwiseControlNet(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("blockwise_controlnet_inputs", "tiled", "tile_size", "tile_stride"), + output_params=("blockwise_controlnet_conditioning",), + onload_model_names=("vae",) + ) + + def apply_controlnet_mask_on_latents(self, pipe, latents, mask): + mask = (pipe.preprocess_image(mask) + 1) / 2 + mask = mask.mean(dim=1, keepdim=True) + mask = 1 - torch.nn.functional.interpolate(mask, size=latents.shape[-2:]) + latents = torch.concat([latents, mask], dim=1) + return latents + + def apply_controlnet_mask_on_image(self, pipe, image, mask): + mask = mask.resize(image.size) + mask = pipe.preprocess_image(mask).mean(dim=[0, 1]).cpu() + image = np.array(image) + image[mask > 0] = 0 + image = Image.fromarray(image) + return image + + def process(self, pipe: QwenImagePipeline, blockwise_controlnet_inputs: list[ControlNetInput], tiled, tile_size, tile_stride): + if blockwise_controlnet_inputs is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + conditionings = [] + for controlnet_input in blockwise_controlnet_inputs: + image = controlnet_input.image + if controlnet_input.inpaint_mask is not None: + image = self.apply_controlnet_mask_on_image(pipe, image, controlnet_input.inpaint_mask) + + image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) + image = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + + if controlnet_input.inpaint_mask is not None: + image = self.apply_controlnet_mask_on_latents(pipe, image, controlnet_input.inpaint_mask) + conditionings.append(image) + + return {"blockwise_controlnet_conditioning": conditionings} + + +class QwenImageUnit_EditImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("edit_image", "tiled", "tile_size", "tile_stride", "edit_image_auto_resize"), + output_params=("edit_latents", "edit_image"), + onload_model_names=("vae",) + ) + + + def calculate_dimensions(self, target_area, ratio): + import math + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + + def edit_image_auto_resize(self, edit_image): + calculated_width, calculated_height = self.calculate_dimensions(1024 * 1024, edit_image.size[0] / edit_image.size[1]) + return edit_image.resize((calculated_width, calculated_height)) + + + def process(self, pipe: QwenImagePipeline, edit_image, tiled, tile_size, tile_stride, edit_image_auto_resize=False): + if edit_image is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + if isinstance(edit_image, Image.Image): + resized_edit_image = self.edit_image_auto_resize(edit_image) if edit_image_auto_resize else edit_image + edit_image = pipe.preprocess_image(resized_edit_image).to(device=pipe.device, dtype=pipe.torch_dtype) + edit_latents = pipe.vae.encode(edit_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + else: + resized_edit_image, edit_latents = [], [] + for image in edit_image: + if edit_image_auto_resize: + image = self.edit_image_auto_resize(image) + resized_edit_image.append(image) + image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) + latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + edit_latents.append(latents) + return {"edit_latents": edit_latents, "edit_image": resized_edit_image} + + +class QwenImageUnit_ContextImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("context_image", "height", "width", "tiled", "tile_size", "tile_stride"), + output_params=("context_latents",), + onload_model_names=("vae",) + ) + + def process(self, pipe: QwenImagePipeline, context_image, height, width, tiled, tile_size, tile_stride): + if context_image is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + context_image = pipe.preprocess_image(context_image.resize((width, height))).to(device=pipe.device, dtype=pipe.torch_dtype) + context_latents = pipe.vae.encode(context_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + return {"context_latents": context_latents} + + +def model_fn_qwen_image( + dit: QwenImageDiT = None, + blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None, + latents=None, + timestep=None, + prompt_emb=None, + prompt_emb_mask=None, + height=None, + width=None, + blockwise_controlnet_conditioning=None, + blockwise_controlnet_inputs=None, + progress_id=0, + num_inference_steps=1, + entity_prompt_emb=None, + entity_prompt_emb_mask=None, + entity_masks=None, + edit_latents=None, + context_latents=None, + enable_fp8_attention=False, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + edit_rope_interpolation=False, + **kwargs +): + img_shapes = [(latents.shape[0], latents.shape[2]//2, latents.shape[3]//2)] + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + timestep = timestep / 1000 + + image = rearrange(latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + image_seq_len = image.shape[1] + + if context_latents is not None: + img_shapes += [(context_latents.shape[0], context_latents.shape[2]//2, context_latents.shape[3]//2)] + context_image = rearrange(context_latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=context_latents.shape[2]//2, W=context_latents.shape[3]//2, P=2, Q=2) + image = torch.cat([image, context_image], dim=1) + if edit_latents is not None: + edit_latents_list = edit_latents if isinstance(edit_latents, list) else [edit_latents] + img_shapes += [(e.shape[0], e.shape[2]//2, e.shape[3]//2) for e in edit_latents_list] + edit_image = [rearrange(e, "B C (H P) (W Q) -> B (H W) (C P Q)", H=e.shape[2]//2, W=e.shape[3]//2, P=2, Q=2) for e in edit_latents_list] + image = torch.cat([image] + edit_image, dim=1) + + image = dit.img_in(image) + conditioning = dit.time_text_embed(timestep, image.dtype) + + if entity_prompt_emb is not None: + text, image_rotary_emb, attention_mask = dit.process_entity_masks( + latents, prompt_emb, prompt_emb_mask, entity_prompt_emb, entity_prompt_emb_mask, + entity_masks, height, width, image, img_shapes, + ) + else: + text = dit.txt_in(dit.txt_norm(prompt_emb)) + if edit_rope_interpolation: + image_rotary_emb = dit.pos_embed.forward_sampling(img_shapes, txt_seq_lens, device=latents.device) + else: + image_rotary_emb = dit.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + attention_mask = None + + if blockwise_controlnet_conditioning is not None: + blockwise_controlnet_conditioning = blockwise_controlnet.preprocess( + blockwise_controlnet_inputs, blockwise_controlnet_conditioning) + + for block_id, block in enumerate(dit.transformer_blocks): + text, image = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + image=image, + text=text, + temb=conditioning, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + ) + if blockwise_controlnet_conditioning is not None: + image_slice = image[:, :image_seq_len].clone() + controlnet_output = blockwise_controlnet.blockwise_forward( + image=image_slice, conditionings=blockwise_controlnet_conditioning, + controlnet_inputs=blockwise_controlnet_inputs, block_id=block_id, + progress_id=progress_id, num_inference_steps=num_inference_steps, + ) + image[:, :image_seq_len] = image_slice + controlnet_output + + image = dit.norm_out(image, conditioning) + image = dit.proj_out(image) + image = image[:, :image_seq_len] + + latents = rearrange(image, "B (H W) (C P Q) -> B C (H P) (W Q)", H=height//16, W=width//16, P=2, Q=2) + return latents diff --git a/diffsynth/pipelines/wan_video.py b/diffsynth/pipelines/wan_video.py new file mode 100644 index 0000000000000000000000000000000000000000..fa43db1dd893ac5c9271c80eb9589c6d16e9b66e --- /dev/null +++ b/diffsynth/pipelines/wan_video.py @@ -0,0 +1,1517 @@ +import torch, types +import numpy as np +from PIL import Image +from einops import repeat +from typing import Optional, Union +from einops import rearrange +import numpy as np +from PIL import Image +from tqdm import tqdm +from typing import Optional +from typing_extensions import Literal +from transformers import Wav2Vec2Processor + +from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig, gradient_checkpoint_forward +from ..diffusion.base_pipeline import BasePipeline, PipelineUnit + +from ..models.wan_video_dit import WanModel, sinusoidal_embedding_1d +from ..models.wan_video_dit_s2v import rope_precompute +from ..models.wan_video_text_encoder import WanTextEncoder, HuggingfaceTokenizer +from ..models.wan_video_vae import WanVideoVAE +from ..models.wan_video_image_encoder import WanImageEncoder +from ..models.wan_video_vace import VaceWanModel +from ..models.wan_video_motion_controller import WanMotionControllerModel +from ..models.wan_video_animate_adapter import WanAnimateAdapter +from ..models.wan_video_mot import MotWanModel +from ..models.wav2vec import WanS2VAudioEncoder +from ..models.longcat_video_dit import LongCatVideoTransformer3DModel + + +class WanVideoPipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, time_division_factor=4, time_division_remainder=1 + ) + self.scheduler = FlowMatchScheduler("Wan") + self.tokenizer: HuggingfaceTokenizer = None + self.audio_processor: Wav2Vec2Processor = None + self.text_encoder: WanTextEncoder = None + self.image_encoder: WanImageEncoder = None + self.dit: WanModel = None + self.dit2: WanModel = None + self.vae: WanVideoVAE = None + self.motion_controller: WanMotionControllerModel = None + self.vace: VaceWanModel = None + self.vace2: VaceWanModel = None + self.vap: MotWanModel = None + self.animate_adapter: WanAnimateAdapter = None + self.audio_encoder: WanS2VAudioEncoder = None + self.in_iteration_models = ("dit", "motion_controller", "vace", "animate_adapter", "vap") + self.in_iteration_models_2 = ("dit2", "motion_controller", "vace2", "animate_adapter", "vap") + self.units = [ + WanVideoUnit_ShapeChecker(), + WanVideoUnit_NoiseInitializer(), + WanVideoUnit_PromptEmbedder(), + WanVideoUnit_S2V(), + WanVideoUnit_InputVideoEmbedder(), + WanVideoUnit_ImageEmbedderVAE(), + WanVideoUnit_ImageEmbedderCLIP(), + WanVideoUnit_ImageEmbedderFused(), + WanVideoUnit_FunControl(), + WanVideoUnit_FunReference(), + WanVideoUnit_FunCameraControl(), + WanVideoUnit_SpeedControl(), + WanVideoUnit_VACE(), + WanVideoUnit_AnimateVideoSplit(), + WanVideoUnit_AnimatePoseLatents(), + WanVideoUnit_AnimateFacePixelValues(), + WanVideoUnit_AnimateInpaint(), + WanVideoUnit_VAP(), + WanVideoUnit_UnifiedSequenceParallel(), + WanVideoUnit_TeaCache(), + WanVideoUnit_CfgMerger(), + WanVideoUnit_LongCatVideo(), + ] + self.post_units = [ + WanVideoPostUnit_S2V(), + ] + self.model_fn = model_fn_wan_video + + + def enable_usp(self): + from ..utils.xfuser import get_sequence_parallel_world_size, usp_attn_forward, usp_dit_forward + + for block in self.dit.blocks: + block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn) + self.dit.forward = types.MethodType(usp_dit_forward, self.dit) + if self.dit2 is not None: + for block in self.dit2.blocks: + block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn) + self.dit2.forward = types.MethodType(usp_dit_forward, self.dit2) + self.sp_size = get_sequence_parallel_world_size() + self.use_unified_sequence_parallel = True + + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_config: ModelConfig = ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/"), + audio_processor_config: ModelConfig = None, + redirect_common_files: bool = True, + use_usp: bool = False, + vram_limit: float = None, + ): + # Redirect model path + if redirect_common_files: + redirect_dict = { + "models_t5_umt5-xxl-enc-bf16.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "models_t5_umt5-xxl-enc-bf16.safetensors"), + "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14.safetensors"), + "Wan2.1_VAE.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "Wan2.1_VAE.safetensors"), + "Wan2.2_VAE.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "Wan2.2_VAE.safetensors"), + } + for model_config in model_configs: + if model_config.origin_file_pattern is None or model_config.model_id is None: + continue + if model_config.origin_file_pattern in redirect_dict and model_config.model_id != redirect_dict[model_config.origin_file_pattern][0]: + print(f"To avoid repeatedly downloading model files, ({model_config.model_id}, {model_config.origin_file_pattern}) is redirected to {redirect_dict[model_config.origin_file_pattern]}. You can use `redirect_common_files=False` to disable file redirection.") + model_config.model_id = redirect_dict[model_config.origin_file_pattern][0] + model_config.origin_file_pattern = redirect_dict[model_config.origin_file_pattern][1] + + # Initialize pipeline + pipe = WanVideoPipeline(device=device, torch_dtype=torch_dtype) + if use_usp: + from ..utils.xfuser import initialize_usp + initialize_usp() + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder = model_pool.fetch_model("wan_video_text_encoder") + dit = model_pool.fetch_model("wan_video_dit", index=2) + if isinstance(dit, list): + pipe.dit, pipe.dit2 = dit + else: + pipe.dit = dit + pipe.vae = model_pool.fetch_model("wan_video_vae") + pipe.image_encoder = model_pool.fetch_model("wan_video_image_encoder") + pipe.motion_controller = model_pool.fetch_model("wan_video_motion_controller") + vace = model_pool.fetch_model("wan_video_vace", index=2) + if isinstance(vace, list): + pipe.vace, pipe.vace2 = vace + else: + pipe.vace = vace + pipe.vap = model_pool.fetch_model("wan_video_vap") + pipe.audio_encoder = model_pool.fetch_model("wans2v_audio_encoder") + pipe.animate_adapter = model_pool.fetch_model("wan_video_animate_adapter") + + # Size division factor + if pipe.vae is not None: + pipe.height_division_factor = pipe.vae.upsampling_factor * 2 + pipe.width_division_factor = pipe.vae.upsampling_factor * 2 + + # Initialize tokenizer and processor + if tokenizer_config is not None: + tokenizer_config.download_if_necessary() + pipe.tokenizer = HuggingfaceTokenizer(name=tokenizer_config.path, seq_len=512, clean='whitespace') + if audio_processor_config is not None: + audio_processor_config.download_if_necessary() + pipe.audio_processor = Wav2Vec2Processor.from_pretrained(audio_processor_config.path) + + # Unified Sequence Parallel + if use_usp: pipe.enable_usp() + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: Optional[str] = "", + # Image-to-video + input_image: Optional[Image.Image] = None, + # First-last-frame-to-video + end_image: Optional[Image.Image] = None, + # Video-to-video + input_video: Optional[list[Image.Image]] = None, + denoising_strength: Optional[float] = 1.0, + # Speech-to-video + input_audio: Optional[np.array] = None, + audio_embeds: Optional[torch.Tensor] = None, + audio_sample_rate: Optional[int] = 16000, + s2v_pose_video: Optional[list[Image.Image]] = None, + s2v_pose_latents: Optional[torch.Tensor] = None, + motion_video: Optional[list[Image.Image]] = None, + # ControlNet + control_video: Optional[list[Image.Image]] = None, + reference_image: Optional[Image.Image] = None, + # Camera control + camera_control_direction: Optional[Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"]] = None, + camera_control_speed: Optional[float] = 1/54, + camera_control_origin: Optional[tuple] = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0), + # VACE + vace_video: Optional[list[Image.Image]] = None, + vace_video_mask: Optional[Image.Image] = None, + vace_reference_image: Optional[Image.Image] = None, + vace_scale: Optional[float] = 1.0, + # Animate + animate_pose_video: Optional[list[Image.Image]] = None, + animate_face_video: Optional[list[Image.Image]] = None, + animate_inpaint_video: Optional[list[Image.Image]] = None, + animate_mask_video: Optional[list[Image.Image]] = None, + # VAP + vap_video: Optional[list[Image.Image]] = None, + vap_prompt: Optional[str] = " ", + negative_vap_prompt: Optional[str] = " ", + # Randomness + seed: Optional[int] = None, + rand_device: Optional[str] = "cpu", + # Shape + height: Optional[int] = 480, + width: Optional[int] = 832, + num_frames=81, + # Classifier-free guidance + cfg_scale: Optional[float] = 5.0, + cfg_merge: Optional[bool] = False, + # Boundary + switch_DiT_boundary: Optional[float] = 0.875, + # Scheduler + num_inference_steps: Optional[int] = 50, + sigma_shift: Optional[float] = 5.0, + # Speed control + motion_bucket_id: Optional[int] = None, + # LongCat-Video + longcat_video: Optional[list[Image.Image]] = None, + # VAE tiling + tiled: Optional[bool] = True, + tile_size: Optional[tuple[int, int]] = (30, 52), + tile_stride: Optional[tuple[int, int]] = (15, 26), + # Sliding window + sliding_window_size: Optional[int] = None, + sliding_window_stride: Optional[int] = None, + # Teacache + tea_cache_l1_thresh: Optional[float] = None, + tea_cache_model_id: Optional[str] = "", + # progress_bar + progress_bar_cmd=tqdm, + ): + # Scheduler + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, shift=sigma_shift) + + # Inputs + inputs_posi = { + "prompt": prompt, + "vap_prompt": vap_prompt, + "tea_cache_l1_thresh": tea_cache_l1_thresh, "tea_cache_model_id": tea_cache_model_id, "num_inference_steps": num_inference_steps, + } + inputs_nega = { + "negative_prompt": negative_prompt, + "negative_vap_prompt": negative_vap_prompt, + "tea_cache_l1_thresh": tea_cache_l1_thresh, "tea_cache_model_id": tea_cache_model_id, "num_inference_steps": num_inference_steps, + } + inputs_shared = { + "input_image": input_image, + "end_image": end_image, + "input_video": input_video, "denoising_strength": denoising_strength, + "control_video": control_video, "reference_image": reference_image, + "camera_control_direction": camera_control_direction, "camera_control_speed": camera_control_speed, "camera_control_origin": camera_control_origin, + "vace_video": vace_video, "vace_video_mask": vace_video_mask, "vace_reference_image": vace_reference_image, "vace_scale": vace_scale, + "seed": seed, "rand_device": rand_device, + "height": height, "width": width, "num_frames": num_frames, + "cfg_scale": cfg_scale, "cfg_merge": cfg_merge, + "sigma_shift": sigma_shift, + "motion_bucket_id": motion_bucket_id, + "longcat_video": longcat_video, + "tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride, + "sliding_window_size": sliding_window_size, "sliding_window_stride": sliding_window_stride, + "input_audio": input_audio, "audio_sample_rate": audio_sample_rate, "s2v_pose_video": s2v_pose_video, "audio_embeds": audio_embeds, "s2v_pose_latents": s2v_pose_latents, "motion_video": motion_video, + "animate_pose_video": animate_pose_video, "animate_face_video": animate_face_video, "animate_inpaint_video": animate_inpaint_video, "animate_mask_video": animate_mask_video, + "vap_video": vap_video, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + # Switch DiT if necessary + if timestep.item() < switch_DiT_boundary * 1000 and self.dit2 is not None and not models["dit"] is self.dit2: + self.load_models_to_device(self.in_iteration_models_2) + models["dit"] = self.dit2 + models["vace"] = self.vace2 + + # Timestep + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + + # Inference + noise_pred_posi = self.model_fn(**models, **inputs_shared, **inputs_posi, timestep=timestep) + if cfg_scale != 1.0: + if cfg_merge: + noise_pred_posi, noise_pred_nega = noise_pred_posi.chunk(2, dim=0) + else: + noise_pred_nega = self.model_fn(**models, **inputs_shared, **inputs_nega, timestep=timestep) + noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega) + else: + noise_pred = noise_pred_posi + + # Scheduler + inputs_shared["latents"] = self.scheduler.step(noise_pred, self.scheduler.timesteps[progress_id], inputs_shared["latents"]) + if "first_frame_latents" in inputs_shared: + inputs_shared["latents"][:, :, 0:1] = inputs_shared["first_frame_latents"] + + # VACE (TODO: remove it) + if vace_reference_image is not None or (animate_pose_video is not None and animate_face_video is not None): + if vace_reference_image is not None and isinstance(vace_reference_image, list): + f = len(vace_reference_image) + else: + f = 1 + inputs_shared["latents"] = inputs_shared["latents"][:, :, f:] + # post-denoising, pre-decoding processing logic + for unit in self.post_units: + inputs_shared, _, _ = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + # Decode + self.load_models_to_device(['vae']) + video = self.vae.decode(inputs_shared["latents"], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + video = self.vae_output_to_video(video) + self.load_models_to_device([]) + + return video + + + +class WanVideoUnit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "num_frames"), + output_params=("height", "width", "num_frames"), + ) + + def process(self, pipe: WanVideoPipeline, height, width, num_frames): + height, width, num_frames = pipe.check_resize_height_width(height, width, num_frames) + return {"height": height, "width": width, "num_frames": num_frames} + + + +class WanVideoUnit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "num_frames", "seed", "rand_device", "vace_reference_image"), + output_params=("noise",) + ) + + def process(self, pipe: WanVideoPipeline, height, width, num_frames, seed, rand_device, vace_reference_image): + length = (num_frames - 1) // 4 + 1 + if vace_reference_image is not None: + f = len(vace_reference_image) if isinstance(vace_reference_image, list) else 1 + length += f + shape = (1, pipe.vae.model.z_dim, length, height // pipe.vae.upsampling_factor, width // pipe.vae.upsampling_factor) + noise = pipe.generate_noise(shape, seed=seed, rand_device=rand_device) + if vace_reference_image is not None: + noise = torch.concat((noise[:, :, -f:], noise[:, :, :-f]), dim=2) + return {"noise": noise} + + + +class WanVideoUnit_InputVideoEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_video", "noise", "tiled", "tile_size", "tile_stride", "vace_reference_image"), + output_params=("latents", "input_latents"), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, input_video, noise, tiled, tile_size, tile_stride, vace_reference_image): + if input_video is None: + return {"latents": noise} + pipe.load_models_to_device(self.onload_model_names) + input_video = pipe.preprocess_video(input_video) + input_latents = pipe.vae.encode(input_video, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + if vace_reference_image is not None: + if not isinstance(vace_reference_image, list): + vace_reference_image = [vace_reference_image] + vace_reference_image = pipe.preprocess_video(vace_reference_image) + vace_reference_latents = pipe.vae.encode(vace_reference_image, device=pipe.device).to(dtype=pipe.torch_dtype, device=pipe.device) + input_latents = torch.concat([vace_reference_latents, input_latents], dim=2) + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents} + + + +class WanVideoUnit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt", "positive": "positive"}, + input_params_nega={"prompt": "negative_prompt", "positive": "positive"}, + output_params=("context",), + onload_model_names=("text_encoder",) + ) + + def encode_prompt(self, pipe: WanVideoPipeline, prompt): + ids, mask = pipe.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids = ids.to(pipe.device) + mask = mask.to(pipe.device) + seq_lens = mask.gt(0).sum(dim=1).long() + prompt_emb = pipe.text_encoder(ids, mask) + for i, v in enumerate(seq_lens): + prompt_emb[:, v:] = 0 + return prompt_emb + + def process(self, pipe: WanVideoPipeline, prompt, positive) -> dict: + pipe.load_models_to_device(self.onload_model_names) + prompt_emb = self.encode_prompt(pipe, prompt) + return {"context": prompt_emb} + + + +class WanVideoUnit_ImageEmbedderCLIP(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "end_image", "height", "width"), + output_params=("clip_feature",), + onload_model_names=("image_encoder",) + ) + + def process(self, pipe: WanVideoPipeline, input_image, end_image, height, width): + if input_image is None or pipe.image_encoder is None or not pipe.dit.require_clip_embedding: + return {} + pipe.load_models_to_device(self.onload_model_names) + image = pipe.preprocess_image(input_image.resize((width, height))).to(pipe.device) + clip_context = pipe.image_encoder.encode_image([image]) + if end_image is not None: + end_image = pipe.preprocess_image(end_image.resize((width, height))).to(pipe.device) + if pipe.dit.has_image_pos_emb: + clip_context = torch.concat([clip_context, pipe.image_encoder.encode_image([end_image])], dim=1) + clip_context = clip_context.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"clip_feature": clip_context} + + + +class WanVideoUnit_ImageEmbedderVAE(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "end_image", "num_frames", "height", "width", "tiled", "tile_size", "tile_stride"), + output_params=("y",), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, input_image, end_image, num_frames, height, width, tiled, tile_size, tile_stride): + if input_image is None or not pipe.dit.require_vae_embedding: + return {} + pipe.load_models_to_device(self.onload_model_names) + image = pipe.preprocess_image(input_image.resize((width, height))).to(pipe.device) + msk = torch.ones(1, num_frames, height//8, width//8, device=pipe.device) + msk[:, 1:] = 0 + if end_image is not None: + end_image = pipe.preprocess_image(end_image.resize((width, height))).to(pipe.device) + vae_input = torch.concat([image.transpose(0,1), torch.zeros(3, num_frames-2, height, width).to(image.device), end_image.transpose(0,1)],dim=1) + msk[:, -1:] = 1 + else: + vae_input = torch.concat([image.transpose(0, 1), torch.zeros(3, num_frames-1, height, width).to(image.device)], dim=1) + + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, height//8, width//8) + msk = msk.transpose(1, 2)[0] + + y = pipe.vae.encode([vae_input.to(dtype=pipe.torch_dtype, device=pipe.device)], device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)[0] + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + y = torch.concat([msk, y]) + y = y.unsqueeze(0) + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"y": y} + + + +class WanVideoUnit_ImageEmbedderFused(PipelineUnit): + """ + Encode input image to latents using VAE. This unit is for Wan-AI/Wan2.2-TI2V-5B. + """ + def __init__(self): + super().__init__( + input_params=("input_image", "latents", "height", "width", "tiled", "tile_size", "tile_stride"), + output_params=("latents", "fuse_vae_embedding_in_latents", "first_frame_latents"), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, input_image, latents, height, width, tiled, tile_size, tile_stride): + if input_image is None or not pipe.dit.fuse_vae_embedding_in_latents: + return {} + pipe.load_models_to_device(self.onload_model_names) + image = pipe.preprocess_image(input_image.resize((width, height))).transpose(0, 1) + z = pipe.vae.encode([image], device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + latents[:, :, 0: 1] = z + return {"latents": latents, "fuse_vae_embedding_in_latents": True, "first_frame_latents": z} + + + +class WanVideoUnit_FunControl(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("control_video", "num_frames", "height", "width", "tiled", "tile_size", "tile_stride", "clip_feature", "y", "latents"), + output_params=("clip_feature", "y"), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, control_video, num_frames, height, width, tiled, tile_size, tile_stride, clip_feature, y, latents): + if control_video is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + control_video = pipe.preprocess_video(control_video) + control_latents = pipe.vae.encode(control_video, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + control_latents = control_latents.to(dtype=pipe.torch_dtype, device=pipe.device) + y_dim = pipe.dit.in_dim-control_latents.shape[1]-latents.shape[1] + if clip_feature is None or y is None: + clip_feature = torch.zeros((1, 257, 1280), dtype=pipe.torch_dtype, device=pipe.device) + y = torch.zeros((1, y_dim, (num_frames - 1) // 4 + 1, height//8, width//8), dtype=pipe.torch_dtype, device=pipe.device) + else: + y = y[:, -y_dim:] + y = torch.concat([control_latents, y], dim=1) + return {"clip_feature": clip_feature, "y": y} + + + +class WanVideoUnit_FunReference(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("reference_image", "height", "width", "reference_image"), + output_params=("reference_latents", "clip_feature"), + onload_model_names=("vae", "image_encoder") + ) + + def process(self, pipe: WanVideoPipeline, reference_image, height, width): + if reference_image is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + reference_image = reference_image.resize((width, height)) + reference_latents = pipe.preprocess_video([reference_image]) + reference_latents = pipe.vae.encode(reference_latents, device=pipe.device) + if pipe.image_encoder is None: + return {"reference_latents": reference_latents} + clip_feature = pipe.preprocess_image(reference_image) + clip_feature = pipe.image_encoder.encode_image([clip_feature]) + return {"reference_latents": reference_latents, "clip_feature": clip_feature} + + + +class WanVideoUnit_FunCameraControl(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "num_frames", "camera_control_direction", "camera_control_speed", "camera_control_origin", "latents", "input_image", "tiled", "tile_size", "tile_stride"), + output_params=("control_camera_latents_input", "y"), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, height, width, num_frames, camera_control_direction, camera_control_speed, camera_control_origin, latents, input_image, tiled, tile_size, tile_stride): + if camera_control_direction is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + camera_control_plucker_embedding = pipe.dit.control_adapter.process_camera_coordinates( + camera_control_direction, num_frames, height, width, camera_control_speed, camera_control_origin) + + control_camera_video = camera_control_plucker_embedding[:num_frames].permute([3, 0, 1, 2]).unsqueeze(0) + control_camera_latents = torch.concat( + [ + torch.repeat_interleave(control_camera_video[:, :, 0:1], repeats=4, dim=2), + control_camera_video[:, :, 1:] + ], dim=2 + ).transpose(1, 2) + b, f, c, h, w = control_camera_latents.shape + control_camera_latents = control_camera_latents.contiguous().view(b, f // 4, 4, c, h, w).transpose(2, 3) + control_camera_latents = control_camera_latents.contiguous().view(b, f // 4, c * 4, h, w).transpose(1, 2) + control_camera_latents_input = control_camera_latents.to(device=pipe.device, dtype=pipe.torch_dtype) + + input_image = input_image.resize((width, height)) + input_latents = pipe.preprocess_video([input_image]) + input_latents = pipe.vae.encode(input_latents, device=pipe.device) + y = torch.zeros_like(latents).to(pipe.device) + y[:, :, :1] = input_latents + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + + if y.shape[1] != pipe.dit.in_dim - latents.shape[1]: + image = pipe.preprocess_image(input_image.resize((width, height))).to(pipe.device) + vae_input = torch.concat([image.transpose(0, 1), torch.zeros(3, num_frames-1, height, width).to(image.device)], dim=1) + y = pipe.vae.encode([vae_input.to(dtype=pipe.torch_dtype, device=pipe.device)], device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)[0] + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + msk = torch.ones(1, num_frames, height//8, width//8, device=pipe.device) + msk[:, 1:] = 0 + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, height//8, width//8) + msk = msk.transpose(1, 2)[0] + y = torch.cat([msk,y]) + y = y.unsqueeze(0) + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"control_camera_latents_input": control_camera_latents_input, "y": y} + + + +class WanVideoUnit_SpeedControl(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("motion_bucket_id",), + output_params=("motion_bucket_id",) + ) + + def process(self, pipe: WanVideoPipeline, motion_bucket_id): + if motion_bucket_id is None: + return {} + motion_bucket_id = torch.Tensor((motion_bucket_id,)).to(dtype=pipe.torch_dtype, device=pipe.device) + return {"motion_bucket_id": motion_bucket_id} + + + +class WanVideoUnit_VACE(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("vace_video", "vace_video_mask", "vace_reference_image", "vace_scale", "height", "width", "num_frames", "tiled", "tile_size", "tile_stride"), + output_params=("vace_context", "vace_scale"), + onload_model_names=("vae",) + ) + + def process( + self, + pipe: WanVideoPipeline, + vace_video, vace_video_mask, vace_reference_image, vace_scale, + height, width, num_frames, + tiled, tile_size, tile_stride + ): + if vace_video is not None or vace_video_mask is not None or vace_reference_image is not None: + pipe.load_models_to_device(["vae"]) + if vace_video is None: + vace_video = torch.zeros((1, 3, num_frames, height, width), dtype=pipe.torch_dtype, device=pipe.device) + else: + vace_video = pipe.preprocess_video(vace_video) + + if vace_video_mask is None: + vace_video_mask = torch.ones_like(vace_video) + else: + vace_video_mask = pipe.preprocess_video(vace_video_mask, min_value=0, max_value=1) + + inactive = vace_video * (1 - vace_video_mask) + 0 * vace_video_mask + reactive = vace_video * vace_video_mask + 0 * (1 - vace_video_mask) + inactive = pipe.vae.encode(inactive, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + reactive = pipe.vae.encode(reactive, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + vace_video_latents = torch.concat((inactive, reactive), dim=1) + + vace_mask_latents = rearrange(vace_video_mask[0,0], "T (H P) (W Q) -> 1 (P Q) T H W", P=8, Q=8) + vace_mask_latents = torch.nn.functional.interpolate(vace_mask_latents, size=((vace_mask_latents.shape[2] + 3) // 4, vace_mask_latents.shape[3], vace_mask_latents.shape[4]), mode='nearest-exact') + + if vace_reference_image is None: + pass + else: + if not isinstance(vace_reference_image,list): + vace_reference_image = [vace_reference_image] + + vace_reference_image = pipe.preprocess_video(vace_reference_image) + + bs, c, f, h, w = vace_reference_image.shape + new_vace_ref_images = [] + for j in range(f): + new_vace_ref_images.append(vace_reference_image[0, :, j:j+1]) + vace_reference_image = new_vace_ref_images + + vace_reference_latents = pipe.vae.encode(vace_reference_image, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + vace_reference_latents = torch.concat((vace_reference_latents, torch.zeros_like(vace_reference_latents)), dim=1) + vace_reference_latents = [u.unsqueeze(0) for u in vace_reference_latents] + + vace_video_latents = torch.concat((*vace_reference_latents, vace_video_latents), dim=2) + vace_mask_latents = torch.concat((torch.zeros_like(vace_mask_latents[:, :, :f]), vace_mask_latents), dim=2) + + vace_context = torch.concat((vace_video_latents, vace_mask_latents), dim=1) + return {"vace_context": vace_context, "vace_scale": vace_scale} + else: + return {"vace_context": None, "vace_scale": vace_scale} + + +class WanVideoUnit_VAP(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + onload_model_names=("text_encoder", "vae", "image_encoder"), + input_params=("vap_video", "vap_prompt", "negative_vap_prompt", "end_image", "num_frames", "height", "width", "tiled", "tile_size", "tile_stride"), + output_params=("vap_clip_feature", "vap_hidden_state", "context_vap") + ) + + def encode_prompt(self, pipe: WanVideoPipeline, prompt): + ids, mask = pipe.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids = ids.to(pipe.device) + mask = mask.to(pipe.device) + seq_lens = mask.gt(0).sum(dim=1).long() + prompt_emb = pipe.text_encoder(ids, mask) + for i, v in enumerate(seq_lens): + prompt_emb[:, v:] = 0 + return prompt_emb + + def process(self, pipe: WanVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if inputs_shared.get("vap_video") is None: + return inputs_shared, inputs_posi, inputs_nega + else: + # 1. encode vap prompt + pipe.load_models_to_device(["text_encoder"]) + vap_prompt, negative_vap_prompt = inputs_posi.get("vap_prompt", ""), inputs_nega.get("negative_vap_prompt", "") + vap_prompt_emb = self.encode_prompt(pipe, vap_prompt) + negative_vap_prompt_emb = self.encode_prompt(pipe, negative_vap_prompt) + inputs_posi.update({"context_vap":vap_prompt_emb}) + inputs_nega.update({"context_vap":negative_vap_prompt_emb}) + # 2. prepare vap image clip embedding + pipe.load_models_to_device(["vae", "image_encoder"]) + vap_video, end_image = inputs_shared.get("vap_video"), inputs_shared.get("end_image") + + num_frames, height, width = inputs_shared.get("num_frames"),inputs_shared.get("height"), inputs_shared.get("width") + + image_vap = pipe.preprocess_image(vap_video[0].resize((width, height))).to(pipe.device) + + vap_clip_context = pipe.image_encoder.encode_image([image_vap]) + if end_image is not None: + vap_end_image = pipe.preprocess_image(vap_video[-1].resize((width, height))).to(pipe.device) + if pipe.dit.has_image_pos_emb: + vap_clip_context = torch.concat([vap_clip_context, pipe.image_encoder.encode_image([vap_end_image])], dim=1) + vap_clip_context = vap_clip_context.to(dtype=pipe.torch_dtype, device=pipe.device) + inputs_shared.update({"vap_clip_feature":vap_clip_context}) + + # 3. prepare vap latents + msk = torch.ones(1, num_frames, height//8, width//8, device=pipe.device) + msk[:, 1:] = 0 + if end_image is not None: + msk[:, -1:] = 1 + last_image_vap = pipe.preprocess_image(vap_video[-1].resize((width, height))).to(pipe.device) + vae_input = torch.concat([image_vap.transpose(0,1), torch.zeros(3, num_frames-2, height, width).to(image_vap.device), last_image_vap.transpose(0,1)],dim=1) + else: + vae_input = torch.concat([image_vap.transpose(0, 1), torch.zeros(3, num_frames-1, height, width).to(image_vap.device)], dim=1) + + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, height//8, width//8) + msk = msk.transpose(1, 2)[0] + + tiled,tile_size,tile_stride = inputs_shared.get("tiled"), inputs_shared.get("tile_size"), inputs_shared.get("tile_stride") + + y = pipe.vae.encode([vae_input.to(dtype=pipe.torch_dtype, device=pipe.device)], device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)[0] + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + y = torch.concat([msk, y]) + y = y.unsqueeze(0) + y = y.to(dtype=pipe.torch_dtype, device=pipe.device) + + vap_video = pipe.preprocess_video(vap_video) + vap_latent = pipe.vae.encode(vap_video, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + + vap_latent = torch.concat([vap_latent,y], dim=1).to(dtype=pipe.torch_dtype, device=pipe.device) + inputs_shared.update({"vap_hidden_state":vap_latent}) + + return inputs_shared, inputs_posi, inputs_nega + + + +class WanVideoUnit_UnifiedSequenceParallel(PipelineUnit): + def __init__(self): + super().__init__(input_params=(), output_params=("use_unified_sequence_parallel",)) + + def process(self, pipe: WanVideoPipeline): + if hasattr(pipe, "use_unified_sequence_parallel"): + if pipe.use_unified_sequence_parallel: + return {"use_unified_sequence_parallel": True} + return {} + + + +class WanVideoUnit_TeaCache(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"num_inference_steps": "num_inference_steps", "tea_cache_l1_thresh": "tea_cache_l1_thresh", "tea_cache_model_id": "tea_cache_model_id"}, + input_params_nega={"num_inference_steps": "num_inference_steps", "tea_cache_l1_thresh": "tea_cache_l1_thresh", "tea_cache_model_id": "tea_cache_model_id"}, + output_params=("tea_cache",) + ) + + def process(self, pipe: WanVideoPipeline, num_inference_steps, tea_cache_l1_thresh, tea_cache_model_id): + if tea_cache_l1_thresh is None: + return {} + return {"tea_cache": TeaCache(num_inference_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id)} + + + +class WanVideoUnit_CfgMerger(PipelineUnit): + def __init__(self): + super().__init__(take_over=True) + self.concat_tensor_names = ["context", "clip_feature", "y", "reference_latents"] + + def process(self, pipe: WanVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if not inputs_shared["cfg_merge"]: + return inputs_shared, inputs_posi, inputs_nega + for name in self.concat_tensor_names: + tensor_posi = inputs_posi.get(name) + tensor_nega = inputs_nega.get(name) + tensor_shared = inputs_shared.get(name) + if tensor_posi is not None and tensor_nega is not None: + inputs_shared[name] = torch.concat((tensor_posi, tensor_nega), dim=0) + elif tensor_shared is not None: + inputs_shared[name] = torch.concat((tensor_shared, tensor_shared), dim=0) + inputs_posi.clear() + inputs_nega.clear() + return inputs_shared, inputs_posi, inputs_nega + + +class WanVideoUnit_S2V(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + onload_model_names=("audio_encoder", "vae",), + input_params=("input_audio", "audio_embeds", "num_frames", "height", "width", "tiled", "tile_size", "tile_stride", "audio_sample_rate", "s2v_pose_video", "s2v_pose_latents", "motion_video"), + output_params=("audio_embeds", "motion_latents", "drop_motion_frames", "s2v_pose_latents"), + ) + + def process_audio(self, pipe: WanVideoPipeline, input_audio, audio_sample_rate, num_frames, fps=16, audio_embeds=None, return_all=False): + if audio_embeds is not None: + return {"audio_embeds": audio_embeds} + pipe.load_models_to_device(["audio_encoder"]) + audio_embeds = pipe.audio_encoder.get_audio_feats_per_inference(input_audio, audio_sample_rate, pipe.audio_processor, fps=fps, batch_frames=num_frames-1, dtype=pipe.torch_dtype, device=pipe.device) + if return_all: + return audio_embeds + else: + return {"audio_embeds": audio_embeds[0]} + + def process_motion_latents(self, pipe: WanVideoPipeline, height, width, tiled, tile_size, tile_stride, motion_video=None): + pipe.load_models_to_device(["vae"]) + motion_frames = 73 + kwargs = {} + if motion_video is not None and len(motion_video) > 0: + assert len(motion_video) == motion_frames, f"motion video must have {motion_frames} frames, but got {len(motion_video)}" + motion_latents = pipe.preprocess_video(motion_video) + kwargs["drop_motion_frames"] = False + else: + motion_latents = torch.zeros([1, 3, motion_frames, height, width], dtype=pipe.torch_dtype, device=pipe.device) + kwargs["drop_motion_frames"] = True + motion_latents = pipe.vae.encode(motion_latents, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + kwargs.update({"motion_latents": motion_latents}) + return kwargs + + def process_pose_cond(self, pipe: WanVideoPipeline, s2v_pose_video, num_frames, height, width, tiled, tile_size, tile_stride, s2v_pose_latents=None, num_repeats=1, return_all=False): + if s2v_pose_latents is not None: + return {"s2v_pose_latents": s2v_pose_latents} + if s2v_pose_video is None: + return {"s2v_pose_latents": None} + pipe.load_models_to_device(["vae"]) + infer_frames = num_frames - 1 + input_video = pipe.preprocess_video(s2v_pose_video)[:, :, :infer_frames * num_repeats] + # pad if not enough frames + padding_frames = infer_frames * num_repeats - input_video.shape[2] + input_video = torch.cat([input_video, -torch.ones(1, 3, padding_frames, height, width, device=input_video.device, dtype=input_video.dtype)], dim=2) + input_videos = input_video.chunk(num_repeats, dim=2) + pose_conds = [] + for r in range(num_repeats): + cond = input_videos[r] + cond = torch.cat([cond[:, :, 0:1].repeat(1, 1, 1, 1, 1), cond], dim=2) + cond_latents = pipe.vae.encode(cond, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + pose_conds.append(cond_latents[:,:,1:]) + if return_all: + return pose_conds + else: + return {"s2v_pose_latents": pose_conds[0]} + + def process(self, pipe: WanVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if (inputs_shared.get("input_audio") is None and inputs_shared.get("audio_embeds") is None) or pipe.audio_encoder is None or pipe.audio_processor is None: + return inputs_shared, inputs_posi, inputs_nega + num_frames, height, width, tiled, tile_size, tile_stride = inputs_shared.get("num_frames"), inputs_shared.get("height"), inputs_shared.get("width"), inputs_shared.get("tiled"), inputs_shared.get("tile_size"), inputs_shared.get("tile_stride") + input_audio, audio_embeds, audio_sample_rate = inputs_shared.pop("input_audio", None), inputs_shared.pop("audio_embeds", None), inputs_shared.get("audio_sample_rate", 16000) + s2v_pose_video, s2v_pose_latents, motion_video = inputs_shared.pop("s2v_pose_video", None), inputs_shared.pop("s2v_pose_latents", None), inputs_shared.pop("motion_video", None) + + audio_input_positive = self.process_audio(pipe, input_audio, audio_sample_rate, num_frames, audio_embeds=audio_embeds) + inputs_posi.update(audio_input_positive) + inputs_nega.update({"audio_embeds": 0.0 * audio_input_positive["audio_embeds"]}) + + inputs_shared.update(self.process_motion_latents(pipe, height, width, tiled, tile_size, tile_stride, motion_video)) + inputs_shared.update(self.process_pose_cond(pipe, s2v_pose_video, num_frames, height, width, tiled, tile_size, tile_stride, s2v_pose_latents=s2v_pose_latents)) + return inputs_shared, inputs_posi, inputs_nega + + @staticmethod + def pre_calculate_audio_pose(pipe: WanVideoPipeline, input_audio=None, audio_sample_rate=16000, s2v_pose_video=None, num_frames=81, height=448, width=832, fps=16, tiled=True, tile_size=(30, 52), tile_stride=(15, 26)): + assert pipe.audio_encoder is not None and pipe.audio_processor is not None, "Please load audio encoder and audio processor first." + shapes = WanVideoUnit_ShapeChecker().process(pipe, height, width, num_frames) + height, width, num_frames = shapes["height"], shapes["width"], shapes["num_frames"] + unit = WanVideoUnit_S2V() + audio_embeds = unit.process_audio(pipe, input_audio, audio_sample_rate, num_frames, fps, return_all=True) + pose_latents = unit.process_pose_cond(pipe, s2v_pose_video, num_frames, height, width, num_repeats=len(audio_embeds), return_all=True, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + pose_latents = None if s2v_pose_video is None else pose_latents + return audio_embeds, pose_latents, len(audio_embeds) + + +class WanVideoPostUnit_S2V(PipelineUnit): + def __init__(self): + super().__init__(input_params=("latents", "motion_latents", "drop_motion_frames")) + + def process(self, pipe: WanVideoPipeline, latents, motion_latents, drop_motion_frames): + if pipe.audio_encoder is None or motion_latents is None or drop_motion_frames: + return {} + latents = torch.cat([motion_latents, latents[:,:,1:]], dim=2) + return {"latents": latents} + + +class WanVideoUnit_AnimateVideoSplit(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_video", "animate_pose_video", "animate_face_video", "animate_inpaint_video", "animate_mask_video"), + output_params=("animate_pose_video", "animate_face_video", "animate_inpaint_video", "animate_mask_video") + ) + + def process(self, pipe: WanVideoPipeline, input_video, animate_pose_video, animate_face_video, animate_inpaint_video, animate_mask_video): + if input_video is None: + return {} + if animate_pose_video is not None: + animate_pose_video = animate_pose_video[:len(input_video) - 4] + if animate_face_video is not None: + animate_face_video = animate_face_video[:len(input_video) - 4] + if animate_inpaint_video is not None: + animate_inpaint_video = animate_inpaint_video[:len(input_video) - 4] + if animate_mask_video is not None: + animate_mask_video = animate_mask_video[:len(input_video) - 4] + return {"animate_pose_video": animate_pose_video, "animate_face_video": animate_face_video, "animate_inpaint_video": animate_inpaint_video, "animate_mask_video": animate_mask_video} + + +class WanVideoUnit_AnimatePoseLatents(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("animate_pose_video", "tiled", "tile_size", "tile_stride"), + output_params=("pose_latents",), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, animate_pose_video, tiled, tile_size, tile_stride): + if animate_pose_video is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + animate_pose_video = pipe.preprocess_video(animate_pose_video) + pose_latents = pipe.vae.encode(animate_pose_video, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + return {"pose_latents": pose_latents} + + +class WanVideoUnit_AnimateFacePixelValues(PipelineUnit): + def __init__(self): + super().__init__( + take_over=True, + input_params=("animate_face_video",), + output_params=("face_pixel_values"), + ) + + def process(self, pipe: WanVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if inputs_shared.get("animate_face_video", None) is None: + return inputs_shared, inputs_posi, inputs_nega + inputs_posi["face_pixel_values"] = pipe.preprocess_video(inputs_shared["animate_face_video"]) + inputs_nega["face_pixel_values"] = torch.zeros_like(inputs_posi["face_pixel_values"]) - 1 + return inputs_shared, inputs_posi, inputs_nega + + +class WanVideoUnit_AnimateInpaint(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("animate_inpaint_video", "animate_mask_video", "input_image", "tiled", "tile_size", "tile_stride"), + output_params=("y",), + onload_model_names=("vae",) + ) + + def get_i2v_mask(self, lat_t, lat_h, lat_w, mask_len=1, mask_pixel_values=None, device="cuda"): + if mask_pixel_values is None: + msk = torch.zeros(1, (lat_t-1) * 4 + 1, lat_h, lat_w, device=device) + else: + msk = mask_pixel_values.clone() + msk[:, :mask_len] = 1 + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w) + msk = msk.transpose(1, 2)[0] + return msk + + def process(self, pipe: WanVideoPipeline, animate_inpaint_video, animate_mask_video, input_image, tiled, tile_size, tile_stride): + if animate_inpaint_video is None or animate_mask_video is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + + bg_pixel_values = pipe.preprocess_video(animate_inpaint_video) + y_reft = pipe.vae.encode(bg_pixel_values, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)[0].to(dtype=pipe.torch_dtype, device=pipe.device) + _, lat_t, lat_h, lat_w = y_reft.shape + + ref_pixel_values = pipe.preprocess_video([input_image]) + ref_latents = pipe.vae.encode(ref_pixel_values, device=pipe.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=pipe.torch_dtype, device=pipe.device) + mask_ref = self.get_i2v_mask(1, lat_h, lat_w, 1, device=pipe.device) + y_ref = torch.concat([mask_ref, ref_latents[0]]).to(dtype=torch.bfloat16, device=pipe.device) + + mask_pixel_values = 1 - pipe.preprocess_video(animate_mask_video, max_value=1, min_value=0) + mask_pixel_values = rearrange(mask_pixel_values, "b c t h w -> (b t) c h w") + mask_pixel_values = torch.nn.functional.interpolate(mask_pixel_values, size=(lat_h, lat_w), mode='nearest') + mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0] + msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, 0, mask_pixel_values=mask_pixel_values, device=pipe.device) + + y_reft = torch.concat([msk_reft, y_reft]).to(dtype=torch.bfloat16, device=pipe.device) + y = torch.concat([y_ref, y_reft], dim=1).unsqueeze(0) + return {"y": y} + + +class WanVideoUnit_LongCatVideo(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("longcat_video",), + output_params=("longcat_latents",), + onload_model_names=("vae",) + ) + + def process(self, pipe: WanVideoPipeline, longcat_video): + if longcat_video is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + longcat_video = pipe.preprocess_video(longcat_video) + longcat_latents = pipe.vae.encode(longcat_video, device=pipe.device).to(dtype=pipe.torch_dtype, device=pipe.device) + return {"longcat_latents": longcat_latents} + + +class TeaCache: + def __init__(self, num_inference_steps, rel_l1_thresh, model_id): + self.num_inference_steps = num_inference_steps + self.step = 0 + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.rel_l1_thresh = rel_l1_thresh + self.previous_residual = None + self.previous_hidden_states = None + + self.coefficients_dict = { + "Wan2.1-T2V-1.3B": [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02], + "Wan2.1-T2V-14B": [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01], + "Wan2.1-I2V-14B-480P": [2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01], + "Wan2.1-I2V-14B-720P": [ 8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02], + } + if model_id not in self.coefficients_dict: + supported_model_ids = ", ".join([i for i in self.coefficients_dict]) + raise ValueError(f"{model_id} is not a supported TeaCache model id. Please choose a valid model id in ({supported_model_ids}).") + self.coefficients = self.coefficients_dict[model_id] + + def check(self, dit: WanModel, x, t_mod): + modulated_inp = t_mod.clone() + if self.step == 0 or self.step == self.num_inference_steps - 1: + should_calc = True + self.accumulated_rel_l1_distance = 0 + else: + coefficients = self.coefficients + rescale_func = np.poly1d(coefficients) + self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()) + if self.accumulated_rel_l1_distance < self.rel_l1_thresh: + should_calc = False + else: + should_calc = True + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = modulated_inp + self.step += 1 + if self.step == self.num_inference_steps: + self.step = 0 + if should_calc: + self.previous_hidden_states = x.clone() + return not should_calc + + def store(self, hidden_states): + self.previous_residual = hidden_states - self.previous_hidden_states + self.previous_hidden_states = None + + def update(self, hidden_states): + hidden_states = hidden_states + self.previous_residual + return hidden_states + + + +class TemporalTiler_BCTHW: + def __init__(self): + pass + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if border_width == 0: + return x + + shift = 0.5 + if not left_bound: + x[:border_width] = (torch.arange(border_width) + shift) / border_width + if not right_bound: + x[-border_width:] = torch.flip((torch.arange(border_width) + shift) / border_width, dims=(0,)) + return x + + def build_mask(self, data, is_bound, border_width): + _, _, T, _, _ = data.shape + t = self.build_1d_mask(T, is_bound[0], is_bound[1], border_width[0]) + mask = repeat(t, "T -> 1 1 T 1 1") + return mask + + def run(self, model_fn, sliding_window_size, sliding_window_stride, computation_device, computation_dtype, model_kwargs, tensor_names, batch_size=None): + tensor_names = [tensor_name for tensor_name in tensor_names if model_kwargs.get(tensor_name) is not None] + tensor_dict = {tensor_name: model_kwargs[tensor_name] for tensor_name in tensor_names} + B, C, T, H, W = tensor_dict[tensor_names[0]].shape + if batch_size is not None: + B *= batch_size + data_device, data_dtype = tensor_dict[tensor_names[0]].device, tensor_dict[tensor_names[0]].dtype + value = torch.zeros((B, C, T, H, W), device=data_device, dtype=data_dtype) + weight = torch.zeros((1, 1, T, 1, 1), device=data_device, dtype=data_dtype) + for t in range(0, T, sliding_window_stride): + if t - sliding_window_stride >= 0 and t - sliding_window_stride + sliding_window_size >= T: + continue + t_ = min(t + sliding_window_size, T) + model_kwargs.update({ + tensor_name: tensor_dict[tensor_name][:, :, t: t_:, :].to(device=computation_device, dtype=computation_dtype) \ + for tensor_name in tensor_names + }) + model_output = model_fn(**model_kwargs).to(device=data_device, dtype=data_dtype) + mask = self.build_mask( + model_output, + is_bound=(t == 0, t_ == T), + border_width=(sliding_window_size - sliding_window_stride,) + ).to(device=data_device, dtype=data_dtype) + value[:, :, t: t_, :, :] += model_output * mask + weight[:, :, t: t_, :, :] += mask + value /= weight + model_kwargs.update(tensor_dict) + return value + + + +def model_fn_wan_video( + dit: WanModel, + motion_controller: WanMotionControllerModel = None, + vace: VaceWanModel = None, + vap: MotWanModel = None, + animate_adapter: WanAnimateAdapter = None, + latents: torch.Tensor = None, + timestep: torch.Tensor = None, + context: torch.Tensor = None, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + reference_latents = None, + vace_context = None, + vace_scale = 1.0, + audio_embeds: Optional[torch.Tensor] = None, + motion_latents: Optional[torch.Tensor] = None, + s2v_pose_latents: Optional[torch.Tensor] = None, + vap_hidden_state = None, + vap_clip_feature = None, + context_vap = None, + drop_motion_frames: bool = True, + tea_cache: TeaCache = None, + use_unified_sequence_parallel: bool = False, + motion_bucket_id: Optional[torch.Tensor] = None, + pose_latents=None, + face_pixel_values=None, + longcat_latents=None, + sliding_window_size: Optional[int] = None, + sliding_window_stride: Optional[int] = None, + cfg_merge: bool = False, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + control_camera_latents_input = None, + fuse_vae_embedding_in_latents: bool = False, + **kwargs, +): + if sliding_window_size is not None and sliding_window_stride is not None: + model_kwargs = dict( + dit=dit, + motion_controller=motion_controller, + vace=vace, + latents=latents, + timestep=timestep, + context=context, + clip_feature=clip_feature, + y=y, + reference_latents=reference_latents, + vace_context=vace_context, + vace_scale=vace_scale, + tea_cache=tea_cache, + use_unified_sequence_parallel=use_unified_sequence_parallel, + motion_bucket_id=motion_bucket_id, + ) + return TemporalTiler_BCTHW().run( + model_fn_wan_video, + sliding_window_size, sliding_window_stride, + latents.device, latents.dtype, + model_kwargs=model_kwargs, + tensor_names=["latents", "y"], + batch_size=2 if cfg_merge else 1 + ) + # LongCat-Video + if isinstance(dit, LongCatVideoTransformer3DModel): + return model_fn_longcat_video( + dit=dit, + latents=latents, + timestep=timestep, + context=context, + longcat_latents=longcat_latents, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + ) + + # wan2.2 s2v + if audio_embeds is not None: + return model_fn_wans2v( + dit=dit, + latents=latents, + timestep=timestep, + context=context, + audio_embeds=audio_embeds, + motion_latents=motion_latents, + s2v_pose_latents=s2v_pose_latents, + drop_motion_frames=drop_motion_frames, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + use_gradient_checkpointing=use_gradient_checkpointing, + use_unified_sequence_parallel=use_unified_sequence_parallel, + ) + + if use_unified_sequence_parallel: + import torch.distributed as dist + from xfuser.core.distributed import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group) + + # Timestep + if dit.seperated_timestep and fuse_vae_embedding_in_latents: + timestep = torch.concat([ + torch.zeros((1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device), + torch.ones((latents.shape[2] - 1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device) * timestep + ]).flatten() + t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep).unsqueeze(0)) + if use_unified_sequence_parallel and dist.is_initialized() and dist.get_world_size() > 1: + t_chunks = torch.chunk(t, get_sequence_parallel_world_size(), dim=1) + t_chunks = [torch.nn.functional.pad(chunk, (0, 0, 0, t_chunks[0].shape[1]-chunk.shape[1]), value=0) for chunk in t_chunks] + t = t_chunks[get_sequence_parallel_rank()] + t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim)) + else: + t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep)) + t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim)) + + # Motion Controller + if motion_bucket_id is not None and motion_controller is not None: + t_mod = t_mod + motion_controller(motion_bucket_id).unflatten(1, (6, dit.dim)) + context = dit.text_embedding(context) + + x = latents + # Merged cfg + if x.shape[0] != context.shape[0]: + x = torch.concat([x] * context.shape[0], dim=0) + if timestep.shape[0] != context.shape[0]: + timestep = torch.concat([timestep] * context.shape[0], dim=0) + + # Image Embedding + if y is not None and dit.require_vae_embedding: + x = torch.cat([x, y], dim=1) + if clip_feature is not None and dit.require_clip_embedding: + clip_embdding = dit.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + + # Camera control + x = dit.patchify(x, control_camera_latents_input) + + # Animate + if pose_latents is not None and face_pixel_values is not None: + x, motion_vec = animate_adapter.after_patch_embedding(x, pose_latents, face_pixel_values) + + # Patchify + f, h, w = x.shape[2:] + x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous() + + # Reference image + if reference_latents is not None: + if len(reference_latents.shape) == 5: + reference_latents = reference_latents[:, :, 0] + reference_latents = dit.ref_conv(reference_latents).flatten(2).transpose(1, 2) + x = torch.concat([reference_latents, x], dim=1) + f += 1 + + freqs = torch.cat([ + dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + + # VAP + if vap is not None: + # hidden state + x_vap = vap_hidden_state + x_vap = vap.patchify(x_vap) + x_vap = rearrange(x_vap, 'b c f h w -> b (f h w) c').contiguous() + # Timestep + clean_timestep = torch.ones(timestep.shape, device=timestep.device).to(timestep.dtype) + t = vap.time_embedding(sinusoidal_embedding_1d(vap.freq_dim, clean_timestep)) + t_mod_vap = vap.time_projection(t).unflatten(1, (6, vap.dim)) + + # rope + freqs_vap = vap.compute_freqs_mot(f,h,w).to(x.device) + + # context + vap_clip_embedding = vap.img_emb(vap_clip_feature) + context_vap = vap.text_embedding(context_vap) + context_vap = torch.cat([vap_clip_embedding, context_vap], dim=1) + + # TeaCache + if tea_cache is not None: + tea_cache_update = tea_cache.check(dit, x, t_mod) + else: + tea_cache_update = False + + if vace_context is not None: + vace_hints = vace( + x, vace_context, context, t_mod, freqs, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload + ) + + # blocks + if use_unified_sequence_parallel: + if dist.is_initialized() and dist.get_world_size() > 1: + chunks = torch.chunk(x, get_sequence_parallel_world_size(), dim=1) + pad_shape = chunks[0].shape[1] - chunks[-1].shape[1] + chunks = [torch.nn.functional.pad(chunk, (0, 0, 0, chunks[0].shape[1]-chunk.shape[1]), value=0) for chunk in chunks] + x = chunks[get_sequence_parallel_rank()] + if tea_cache_update: + x = tea_cache.update(x) + else: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + def create_custom_forward_vap(block, vap): + def custom_forward(*inputs): + return vap(block, *inputs) + return custom_forward + + for block_id, block in enumerate(dit.blocks): + # Block + if vap is not None and block_id in vap.mot_layers_mapping: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x, x_vap = torch.utils.checkpoint.checkpoint( + create_custom_forward_vap(block, vap), + x, context, t_mod, freqs, x_vap, context_vap, t_mod_vap, freqs_vap, block_id, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + x, x_vap = torch.utils.checkpoint.checkpoint( + create_custom_forward_vap(block, vap), + x, context, t_mod, freqs, x_vap, context_vap, t_mod_vap, freqs_vap, block_id, + use_reentrant=False, + ) + else: + x, x_vap = vap(block, x, context, t_mod, freqs, x_vap, context_vap, t_mod_vap, freqs_vap, block_id) + else: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + # VACE + if vace_context is not None and block_id in vace.vace_layers_mapping: + current_vace_hint = vace_hints[vace.vace_layers_mapping[block_id]] + if use_unified_sequence_parallel and dist.is_initialized() and dist.get_world_size() > 1: + current_vace_hint = torch.chunk(current_vace_hint, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()] + current_vace_hint = torch.nn.functional.pad(current_vace_hint, (0, 0, 0, chunks[0].shape[1] - current_vace_hint.shape[1]), value=0) + x = x + current_vace_hint * vace_scale + + # Animate + if pose_latents is not None and face_pixel_values is not None: + x = animate_adapter.after_transformer_block(block_id, x, motion_vec) + if tea_cache is not None: + tea_cache.store(x) + + x = dit.head(x, t) + if use_unified_sequence_parallel: + if dist.is_initialized() and dist.get_world_size() > 1: + x = get_sp_group().all_gather(x, dim=1) + x = x[:, :-pad_shape] if pad_shape > 0 else x + # Remove reference latents + if reference_latents is not None: + x = x[:, reference_latents.shape[1]:] + f -= 1 + x = dit.unpatchify(x, (f, h, w)) + return x + + +def model_fn_longcat_video( + dit: LongCatVideoTransformer3DModel, + latents: torch.Tensor = None, + timestep: torch.Tensor = None, + context: torch.Tensor = None, + longcat_latents: torch.Tensor = None, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, +): + if longcat_latents is not None: + latents[:, :, :longcat_latents.shape[2]] = longcat_latents + num_cond_latents = longcat_latents.shape[2] + else: + num_cond_latents = 0 + context = context.unsqueeze(0) + encoder_attention_mask = torch.any(context != 0, dim=-1)[:, 0].to(torch.int64) + output = dit( + latents, + timestep, + context, + encoder_attention_mask, + num_cond_latents=num_cond_latents, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + ) + output = -output + output = output.to(latents.dtype) + return output + + +def model_fn_wans2v( + dit, + latents, + timestep, + context, + audio_embeds, + motion_latents, + s2v_pose_latents, + drop_motion_frames=True, + use_gradient_checkpointing_offload=False, + use_gradient_checkpointing=False, + use_unified_sequence_parallel=False, +): + if use_unified_sequence_parallel: + import torch.distributed as dist + from xfuser.core.distributed import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group) + origin_ref_latents = latents[:, :, 0:1] + x = latents[:, :, 1:] + + # context embedding + context = dit.text_embedding(context) + + # audio encode + audio_emb_global, merged_audio_emb = dit.cal_audio_emb(audio_embeds) + + # x and s2v_pose_latents + s2v_pose_latents = torch.zeros_like(x) if s2v_pose_latents is None else s2v_pose_latents + x, (f, h, w) = dit.patchify(dit.patch_embedding(x) + dit.cond_encoder(s2v_pose_latents)) + seq_len_x = seq_len_x_global = x.shape[1] # global used for unified sequence parallel + + # reference image + ref_latents, (rf, rh, rw) = dit.patchify(dit.patch_embedding(origin_ref_latents)) + grid_sizes = dit.get_grid_sizes((f, h, w), (rf, rh, rw)) + x = torch.cat([x, ref_latents], dim=1) + # mask + mask = torch.cat([torch.zeros([1, seq_len_x]), torch.ones([1, ref_latents.shape[1]])], dim=1).to(torch.long).to(x.device) + # freqs + pre_compute_freqs = rope_precompute(x.detach().view(1, x.size(1), dit.num_heads, dit.dim // dit.num_heads), grid_sizes, dit.freqs, start=None) + # motion + x, pre_compute_freqs, mask = dit.inject_motion(x, pre_compute_freqs, mask, motion_latents, drop_motion_frames=drop_motion_frames, add_last_motion=2) + + x = x + dit.trainable_cond_mask(mask).to(x.dtype) + + # tmod + timestep = torch.cat([timestep, torch.zeros([1], dtype=timestep.dtype, device=timestep.device)]) + t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep)) + t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim)).unsqueeze(2).transpose(0, 2) + + if use_unified_sequence_parallel and dist.is_initialized() and dist.get_world_size() > 1: + world_size, sp_rank = get_sequence_parallel_world_size(), get_sequence_parallel_rank() + assert x.shape[1] % world_size == 0, f"the dimension after chunk must be divisible by world size, but got {x.shape[1]} and {get_sequence_parallel_world_size()}" + x = torch.chunk(x, world_size, dim=1)[sp_rank] + seg_idxs = [0] + list(torch.cumsum(torch.tensor([x.shape[1]] * world_size), dim=0).cpu().numpy()) + seq_len_x_list = [min(max(0, seq_len_x - seg_idxs[i]), x.shape[1]) for i in range(len(seg_idxs)-1)] + seq_len_x = seq_len_x_list[sp_rank] + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block_id, block in enumerate(dit.blocks): + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, seq_len_x, pre_compute_freqs[0], + use_reentrant=False, + ) + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(lambda x: dit.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)), + x, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, seq_len_x, pre_compute_freqs[0], + use_reentrant=False, + ) + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(lambda x: dit.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)), + x, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, seq_len_x, pre_compute_freqs[0]) + x = dit.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x_global, use_unified_sequence_parallel) + + if use_unified_sequence_parallel and dist.is_initialized() and dist.get_world_size() > 1: + x = get_sp_group().all_gather(x, dim=1) + + x = x[:, :seq_len_x_global] + x = dit.head(x, t[:-1]) + x = dit.unpatchify(x, (f, h, w)) + # make compatible with wan video + x = torch.cat([origin_ref_latents, x], dim=2) + return x diff --git a/diffsynth/pipelines/z_image.py b/diffsynth/pipelines/z_image.py new file mode 100644 index 0000000000000000000000000000000000000000..f87254f35015f87924c64ac5241f257f6b150995 --- /dev/null +++ b/diffsynth/pipelines/z_image.py @@ -0,0 +1,257 @@ +import torch, math +from PIL import Image +from typing import Union +from tqdm import tqdm +from einops import rearrange +import numpy as np +from typing import Union, List, Optional, Tuple + +from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig, gradient_checkpoint_forward +from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput + +from transformers import AutoTokenizer +from ..models.z_image_text_encoder import ZImageTextEncoder +from ..models.z_image_dit import ZImageDiT +from ..models.flux_vae import FluxVAEEncoder, FluxVAEDecoder + + +class ZImagePipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, + ) + self.scheduler = FlowMatchScheduler("Z-Image") + self.text_encoder: ZImageTextEncoder = None + self.dit: ZImageDiT = None + self.vae_encoder: FluxVAEEncoder = None + self.vae_decoder: FluxVAEDecoder = None + self.tokenizer: AutoTokenizer = None + self.in_iteration_models = ("dit",) + self.units = [ + ZImageUnit_ShapeChecker(), + ZImageUnit_PromptEmbedder(), + ZImageUnit_NoiseInitializer(), + ZImageUnit_InputImageEmbedder(), + ] + self.model_fn = model_fn_z_image + + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_config: ModelConfig = ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"), + vram_limit: float = None, + ): + # Initialize pipeline + pipe = ZImagePipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder = model_pool.fetch_model("z_image_text_encoder") + pipe.dit = model_pool.fetch_model("z_image_dit") + pipe.vae_encoder = model_pool.fetch_model("flux_vae_encoder") + pipe.vae_decoder = model_pool.fetch_model("flux_vae_decoder") + if tokenizer_config is not None: + tokenizer_config.download_if_necessary() + pipe.tokenizer = AutoTokenizer.from_pretrained(tokenizer_config.path) + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: str = "", + cfg_scale: float = 1.0, + # Image + input_image: Image.Image = None, + denoising_strength: float = 1.0, + # Shape + height: int = 1024, + width: int = 1024, + # Randomness + seed: int = None, + rand_device: str = "cpu", + # Steps + num_inference_steps: int = 8, + # Progress bar + progress_bar_cmd = tqdm, + ): + # Scheduler + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength) + + # Parameters + inputs_posi = { + "prompt": prompt, + } + inputs_nega = { + "negative_prompt": negative_prompt, + } + inputs_shared = { + "cfg_scale": cfg_scale, + "input_image": input_image, "denoising_strength": denoising_strength, + "height": height, "width": width, + "seed": seed, "rand_device": rand_device, + "num_inference_steps": num_inference_steps, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + noise_pred = self.cfg_guided_model_fn( + self.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) + + # Decode + self.load_models_to_device(['vae_decoder']) + image = self.vae_decoder(inputs_shared["latents"]) + image = self.vae_output_to_image(image) + self.load_models_to_device([]) + + return image + + +class ZImageUnit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width"), + output_params=("height", "width"), + ) + + def process(self, pipe: ZImagePipeline, height, width): + height, width = pipe.check_resize_height_width(height, width) + return {"height": height, "width": width} + + +class ZImageUnit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt"}, + input_params_nega={"prompt": "negative_prompt"}, + output_params=("prompt_embeds",), + onload_model_names=("text_encoder",) + ) + + def encode_prompt( + self, + pipe, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + max_sequence_length: int = 512, + ) -> List[torch.FloatTensor]: + if isinstance(prompt, str): + prompt = [prompt] + + for i, prompt_item in enumerate(prompt): + messages = [ + {"role": "user", "content": prompt_item}, + ] + prompt_item = pipe.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + prompt[i] = prompt_item + + text_inputs = pipe.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = pipe.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + embeddings_list = [] + + for i in range(len(prompt_embeds)): + embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) + + return embeddings_list + + def process(self, pipe: ZImagePipeline, prompt): + pipe.load_models_to_device(self.onload_model_names) + prompt_embeds = self.encode_prompt(pipe, prompt, pipe.device) + return {"prompt_embeds": prompt_embeds} + + +class ZImageUnit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "seed", "rand_device"), + output_params=("noise",), + ) + + def process(self, pipe: ZImagePipeline, height, width, seed, rand_device): + noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype) + return {"noise": noise} + + +class ZImageUnit_InputImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "noise"), + output_params=("latents", "input_latents"), + onload_model_names=("vae_encoder",) + ) + + def process(self, pipe: ZImagePipeline, input_image, noise): + if input_image is None: + return {"latents": noise, "input_latents": None} + pipe.load_models_to_device(['vae']) + image = pipe.preprocess_image(input_image) + input_latents = pipe.vae_encoder(image) + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents, "input_latents": input_latents} + + +def model_fn_z_image( + dit: ZImageDiT, + latents=None, + timestep=None, + prompt_embeds=None, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + **kwargs, +): + latents = [rearrange(latents, "B C H W -> C B H W")] + timestep = (1000 - timestep) / 1000 + model_output = dit( + latents, + timestep, + prompt_embeds, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + )[0][0] + model_output = -model_output + model_output = rearrange(model_output, "C B H W -> B C H W") + return model_output diff --git a/diffsynth/utils/controlnet/__init__.py b/diffsynth/utils/controlnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..df23b6c61b99319f1f41d6448b0c52ffd03b9f25 --- /dev/null +++ b/diffsynth/utils/controlnet/__init__.py @@ -0,0 +1,2 @@ +from .controlnet_input import ControlNetInput +from .annotator import Annotator diff --git a/diffsynth/utils/controlnet/annotator.py b/diffsynth/utils/controlnet/annotator.py new file mode 100644 index 0000000000000000000000000000000000000000..06553e06d1c6d09f5a3deecfd4ea5604c5dd4352 --- /dev/null +++ b/diffsynth/utils/controlnet/annotator.py @@ -0,0 +1,62 @@ +from typing_extensions import Literal, TypeAlias + + +Processor_id: TypeAlias = Literal[ + "canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "normal", "tile", "none", "inpaint" +] + +class Annotator: + def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=None, device='cuda', skip_processor=False): + if not skip_processor: + if processor_id == "canny": + from controlnet_aux.processor import CannyDetector + self.processor = CannyDetector() + elif processor_id == "depth": + from controlnet_aux.processor import MidasDetector + self.processor = MidasDetector.from_pretrained(model_path).to(device) + elif processor_id == "softedge": + from controlnet_aux.processor import HEDdetector + self.processor = HEDdetector.from_pretrained(model_path).to(device) + elif processor_id == "lineart": + from controlnet_aux.processor import LineartDetector + self.processor = LineartDetector.from_pretrained(model_path).to(device) + elif processor_id == "lineart_anime": + from controlnet_aux.processor import LineartAnimeDetector + self.processor = LineartAnimeDetector.from_pretrained(model_path).to(device) + elif processor_id == "openpose": + from controlnet_aux.processor import OpenposeDetector + self.processor = OpenposeDetector.from_pretrained(model_path).to(device) + elif processor_id == "normal": + from controlnet_aux.processor import NormalBaeDetector + self.processor = NormalBaeDetector.from_pretrained(model_path).to(device) + elif processor_id == "tile" or processor_id == "none" or processor_id == "inpaint": + self.processor = None + else: + raise ValueError(f"Unsupported processor_id: {processor_id}") + else: + self.processor = None + + self.processor_id = processor_id + self.detect_resolution = detect_resolution + + def to(self,device): + if hasattr(self.processor,"model") and hasattr(self.processor.model,"to"): + + self.processor.model.to(device) + + def __call__(self, image, mask=None): + width, height = image.size + if self.processor_id == "openpose": + kwargs = { + "include_body": True, + "include_hand": True, + "include_face": True + } + else: + kwargs = {} + if self.processor is not None: + detect_resolution = self.detect_resolution if self.detect_resolution is not None else min(width, height) + image = self.processor(image, detect_resolution=detect_resolution, image_resolution=min(width, height), **kwargs) + image = image.resize((width, height)) + return image + diff --git a/diffsynth/utils/controlnet/controlnet_input.py b/diffsynth/utils/controlnet/controlnet_input.py new file mode 100644 index 0000000000000000000000000000000000000000..1a2949bc5fab87c4779f5f4259943aa245faa4c4 --- /dev/null +++ b/diffsynth/utils/controlnet/controlnet_input.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass +from PIL import Image + + +@dataclass +class ControlNetInput: + controlnet_id: int = 0 + scale: float = 1.0 + start: float = 1.0 + end: float = 0.0 + image: Image.Image = None + inpaint_mask: Image.Image = None + processor_id: str = None diff --git a/diffsynth/utils/data/__init__.py b/diffsynth/utils/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c6b9daa41bea9d36e012d52a1d280d1cf8d92850 --- /dev/null +++ b/diffsynth/utils/data/__init__.py @@ -0,0 +1,217 @@ +import imageio, os +import numpy as np +from PIL import Image +from tqdm import tqdm +import subprocess +import shutil + + +class LowMemoryVideo: + def __init__(self, file_name): + self.reader = imageio.get_reader(file_name) + + def __len__(self): + return self.reader.count_frames() + + def __getitem__(self, item): + return Image.fromarray(np.array(self.reader.get_data(item))).convert("RGB") + + def __del__(self): + self.reader.close() + + +def split_file_name(file_name): + result = [] + number = -1 + for i in file_name: + if ord(i)>=ord("0") and ord(i)<=ord("9"): + if number == -1: + number = 0 + number = number*10 + ord(i) - ord("0") + else: + if number != -1: + result.append(number) + number = -1 + result.append(i) + if number != -1: + result.append(number) + result = tuple(result) + return result + + +def search_for_images(folder): + file_list = [i for i in os.listdir(folder) if i.endswith(".jpg") or i.endswith(".png")] + file_list = [(split_file_name(file_name), file_name) for file_name in file_list] + file_list = [i[1] for i in sorted(file_list)] + file_list = [os.path.join(folder, i) for i in file_list] + return file_list + + +class LowMemoryImageFolder: + def __init__(self, folder, file_list=None): + if file_list is None: + self.file_list = search_for_images(folder) + else: + self.file_list = [os.path.join(folder, file_name) for file_name in file_list] + + def __len__(self): + return len(self.file_list) + + def __getitem__(self, item): + return Image.open(self.file_list[item]).convert("RGB") + + def __del__(self): + pass + + +def crop_and_resize(image, height, width): + image = np.array(image) + image_height, image_width, _ = image.shape + if image_height / image_width < height / width: + croped_width = int(image_height / height * width) + left = (image_width - croped_width) // 2 + image = image[:, left: left+croped_width] + image = Image.fromarray(image).resize((width, height)) + else: + croped_height = int(image_width / width * height) + left = (image_height - croped_height) // 2 + image = image[left: left+croped_height, :] + image = Image.fromarray(image).resize((width, height)) + return image + + +class VideoData: + def __init__(self, video_file=None, image_folder=None, height=None, width=None, **kwargs): + if video_file is not None: + self.data_type = "video" + self.data = LowMemoryVideo(video_file, **kwargs) + elif image_folder is not None: + self.data_type = "images" + self.data = LowMemoryImageFolder(image_folder, **kwargs) + else: + raise ValueError("Cannot open video or image folder") + self.length = None + self.set_shape(height, width) + + def raw_data(self): + frames = [] + for i in range(self.__len__()): + frames.append(self.__getitem__(i)) + return frames + + def set_length(self, length): + self.length = length + + def set_shape(self, height, width): + self.height = height + self.width = width + + def __len__(self): + if self.length is None: + return len(self.data) + else: + return self.length + + def shape(self): + if self.height is not None and self.width is not None: + return self.height, self.width + else: + height, width, _ = self.__getitem__(0).shape + return height, width + + def __getitem__(self, item): + frame = self.data.__getitem__(item) + width, height = frame.size + if self.height is not None and self.width is not None: + if self.height != height or self.width != width: + frame = crop_and_resize(frame, self.height, self.width) + return frame + + def __del__(self): + pass + + def save_images(self, folder): + os.makedirs(folder, exist_ok=True) + for i in tqdm(range(self.__len__()), desc="Saving images"): + frame = self.__getitem__(i) + frame.save(os.path.join(folder, f"{i}.png")) + + +def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None): + writer = imageio.get_writer(save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params) + for frame in tqdm(frames, desc="Saving video"): + frame = np.array(frame) + writer.append_data(frame) + writer.close() + +def save_frames(frames, save_path): + os.makedirs(save_path, exist_ok=True) + for i, frame in enumerate(tqdm(frames, desc="Saving images")): + frame.save(os.path.join(save_path, f"{i}.png")) + + +def merge_video_audio(video_path: str, audio_path: str): + # TODO: may need a in-python implementation to avoid subprocess dependency + """ + Merge the video and audio into a new video, with the duration set to the shorter of the two, + and overwrite the original video file. + + Parameters: + video_path (str): Path to the original video file + audio_path (str): Path to the audio file + """ + + # check + if not os.path.exists(video_path): + raise FileNotFoundError(f"video file {video_path} does not exist") + if not os.path.exists(audio_path): + raise FileNotFoundError(f"audio file {audio_path} does not exist") + + base, ext = os.path.splitext(video_path) + temp_output = f"{base}_temp{ext}" + + try: + # create ffmpeg command + command = [ + 'ffmpeg', + '-y', # overwrite + '-i', + video_path, + '-i', + audio_path, + '-c:v', + 'copy', # copy video stream + '-c:a', + 'aac', # use AAC audio encoder + '-b:a', + '192k', # set audio bitrate (optional) + '-map', + '0:v:0', # select the first video stream + '-map', + '1:a:0', # select the first audio stream + '-shortest', # choose the shortest duration + temp_output + ] + + # execute the command + result = subprocess.run( + command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + # check result + if result.returncode != 0: + error_msg = f"FFmpeg execute failed: {result.stderr}" + print(error_msg) + raise RuntimeError(error_msg) + + shutil.move(temp_output, video_path) + print(f"Merge completed, saved to {video_path}") + + except Exception as e: + if os.path.exists(temp_output): + os.remove(temp_output) + print(f"merge_video_audio failed with error: {e}") + + +def save_video_with_audio(frames, save_path, audio_path, fps=16, quality=9, ffmpeg_params=None): + save_video(frames, save_path, fps, quality, ffmpeg_params) + merge_video_audio(save_path, audio_path) diff --git a/diffsynth/utils/lora/__init__.py b/diffsynth/utils/lora/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ebbbe55443647d9812965e05e575c0717f2c921 --- /dev/null +++ b/diffsynth/utils/lora/__init__.py @@ -0,0 +1 @@ +from .general import GeneralLoRALoader diff --git a/diffsynth/utils/lora/flux.py b/diffsynth/utils/lora/flux.py new file mode 100644 index 0000000000000000000000000000000000000000..502c5fd449d619160aba29d33d3915ece1763cda --- /dev/null +++ b/diffsynth/utils/lora/flux.py @@ -0,0 +1,204 @@ +from .general import GeneralLoRALoader +import torch, math + + +class FluxLoRALoader(GeneralLoRALoader): + def __init__(self, device="cpu", torch_dtype=torch.float32): + super().__init__(device=device, torch_dtype=torch_dtype) + + self.diffusers_rename_dict = { + "transformer.single_transformer_blocks.blockid.attn.to_k.lora_A.weight":"single_blocks.blockid.a_to_k.lora_A.weight", + "transformer.single_transformer_blocks.blockid.attn.to_k.lora_B.weight":"single_blocks.blockid.a_to_k.lora_B.weight", + "transformer.single_transformer_blocks.blockid.attn.to_q.lora_A.weight":"single_blocks.blockid.a_to_q.lora_A.weight", + "transformer.single_transformer_blocks.blockid.attn.to_q.lora_B.weight":"single_blocks.blockid.a_to_q.lora_B.weight", + "transformer.single_transformer_blocks.blockid.attn.to_v.lora_A.weight":"single_blocks.blockid.a_to_v.lora_A.weight", + "transformer.single_transformer_blocks.blockid.attn.to_v.lora_B.weight":"single_blocks.blockid.a_to_v.lora_B.weight", + "transformer.single_transformer_blocks.blockid.norm.linear.lora_A.weight":"single_blocks.blockid.norm.linear.lora_A.weight", + "transformer.single_transformer_blocks.blockid.norm.linear.lora_B.weight":"single_blocks.blockid.norm.linear.lora_B.weight", + "transformer.single_transformer_blocks.blockid.proj_mlp.lora_A.weight":"single_blocks.blockid.proj_in_besides_attn.lora_A.weight", + "transformer.single_transformer_blocks.blockid.proj_mlp.lora_B.weight":"single_blocks.blockid.proj_in_besides_attn.lora_B.weight", + "transformer.single_transformer_blocks.blockid.proj_out.lora_A.weight":"single_blocks.blockid.proj_out.lora_A.weight", + "transformer.single_transformer_blocks.blockid.proj_out.lora_B.weight":"single_blocks.blockid.proj_out.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_A.weight":"blocks.blockid.attn.b_to_k.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_B.weight":"blocks.blockid.attn.b_to_k.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_A.weight":"blocks.blockid.attn.b_to_q.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_B.weight":"blocks.blockid.attn.b_to_q.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_A.weight":"blocks.blockid.attn.b_to_v.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_B.weight":"blocks.blockid.attn.b_to_v.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.to_add_out.lora_A.weight":"blocks.blockid.attn.b_to_out.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.to_add_out.lora_B.weight":"blocks.blockid.attn.b_to_out.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.to_k.lora_A.weight":"blocks.blockid.attn.a_to_k.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.to_k.lora_B.weight":"blocks.blockid.attn.a_to_k.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.to_out.0.lora_A.weight":"blocks.blockid.attn.a_to_out.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.to_out.0.lora_B.weight":"blocks.blockid.attn.a_to_out.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.to_q.lora_A.weight":"blocks.blockid.attn.a_to_q.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.to_q.lora_B.weight":"blocks.blockid.attn.a_to_q.lora_B.weight", + "transformer.transformer_blocks.blockid.attn.to_v.lora_A.weight":"blocks.blockid.attn.a_to_v.lora_A.weight", + "transformer.transformer_blocks.blockid.attn.to_v.lora_B.weight":"blocks.blockid.attn.a_to_v.lora_B.weight", + "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_A.weight":"blocks.blockid.ff_a.0.lora_A.weight", + "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_B.weight":"blocks.blockid.ff_a.0.lora_B.weight", + "transformer.transformer_blocks.blockid.ff.net.2.lora_A.weight":"blocks.blockid.ff_a.2.lora_A.weight", + "transformer.transformer_blocks.blockid.ff.net.2.lora_B.weight":"blocks.blockid.ff_a.2.lora_B.weight", + "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_A.weight":"blocks.blockid.ff_b.0.lora_A.weight", + "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_B.weight":"blocks.blockid.ff_b.0.lora_B.weight", + "transformer.transformer_blocks.blockid.ff_context.net.2.lora_A.weight":"blocks.blockid.ff_b.2.lora_A.weight", + "transformer.transformer_blocks.blockid.ff_context.net.2.lora_B.weight":"blocks.blockid.ff_b.2.lora_B.weight", + "transformer.transformer_blocks.blockid.norm1.linear.lora_A.weight":"blocks.blockid.norm1_a.linear.lora_A.weight", + "transformer.transformer_blocks.blockid.norm1.linear.lora_B.weight":"blocks.blockid.norm1_a.linear.lora_B.weight", + "transformer.transformer_blocks.blockid.norm1_context.linear.lora_A.weight":"blocks.blockid.norm1_b.linear.lora_A.weight", + "transformer.transformer_blocks.blockid.norm1_context.linear.lora_B.weight":"blocks.blockid.norm1_b.linear.lora_B.weight", + } + + self.civitai_rename_dict = { + "lora_unet_double_blocks_blockid_img_mod_lin.lora_down.weight": "blocks.blockid.norm1_a.linear.lora_A.weight", + "lora_unet_double_blocks_blockid_img_mod_lin.lora_up.weight": "blocks.blockid.norm1_a.linear.lora_B.weight", + "lora_unet_double_blocks_blockid_txt_mod_lin.lora_down.weight": "blocks.blockid.norm1_b.linear.lora_A.weight", + "lora_unet_double_blocks_blockid_txt_mod_lin.lora_up.weight": "blocks.blockid.norm1_b.linear.lora_B.weight", + "lora_unet_double_blocks_blockid_img_attn_qkv.lora_down.weight": "blocks.blockid.attn.a_to_qkv.lora_A.weight", + "lora_unet_double_blocks_blockid_img_attn_qkv.lora_up.weight": "blocks.blockid.attn.a_to_qkv.lora_B.weight", + "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_down.weight": "blocks.blockid.attn.b_to_qkv.lora_A.weight", + "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_up.weight": "blocks.blockid.attn.b_to_qkv.lora_B.weight", + "lora_unet_double_blocks_blockid_img_attn_proj.lora_down.weight": "blocks.blockid.attn.a_to_out.lora_A.weight", + "lora_unet_double_blocks_blockid_img_attn_proj.lora_up.weight": "blocks.blockid.attn.a_to_out.lora_B.weight", + "lora_unet_double_blocks_blockid_txt_attn_proj.lora_down.weight": "blocks.blockid.attn.b_to_out.lora_A.weight", + "lora_unet_double_blocks_blockid_txt_attn_proj.lora_up.weight": "blocks.blockid.attn.b_to_out.lora_B.weight", + "lora_unet_double_blocks_blockid_img_mlp_0.lora_down.weight": "blocks.blockid.ff_a.0.lora_A.weight", + "lora_unet_double_blocks_blockid_img_mlp_0.lora_up.weight": "blocks.blockid.ff_a.0.lora_B.weight", + "lora_unet_double_blocks_blockid_img_mlp_2.lora_down.weight": "blocks.blockid.ff_a.2.lora_A.weight", + "lora_unet_double_blocks_blockid_img_mlp_2.lora_up.weight": "blocks.blockid.ff_a.2.lora_B.weight", + "lora_unet_double_blocks_blockid_txt_mlp_0.lora_down.weight": "blocks.blockid.ff_b.0.lora_A.weight", + "lora_unet_double_blocks_blockid_txt_mlp_0.lora_up.weight": "blocks.blockid.ff_b.0.lora_B.weight", + "lora_unet_double_blocks_blockid_txt_mlp_2.lora_down.weight": "blocks.blockid.ff_b.2.lora_A.weight", + "lora_unet_double_blocks_blockid_txt_mlp_2.lora_up.weight": "blocks.blockid.ff_b.2.lora_B.weight", + "lora_unet_single_blocks_blockid_modulation_lin.lora_down.weight": "single_blocks.blockid.norm.linear.lora_A.weight", + "lora_unet_single_blocks_blockid_modulation_lin.lora_up.weight": "single_blocks.blockid.norm.linear.lora_B.weight", + "lora_unet_single_blocks_blockid_linear1.lora_down.weight": "single_blocks.blockid.to_qkv_mlp.lora_A.weight", + "lora_unet_single_blocks_blockid_linear1.lora_up.weight": "single_blocks.blockid.to_qkv_mlp.lora_B.weight", + "lora_unet_single_blocks_blockid_linear2.lora_down.weight": "single_blocks.blockid.proj_out.lora_A.weight", + "lora_unet_single_blocks_blockid_linear2.lora_up.weight": "single_blocks.blockid.proj_out.lora_B.weight", + } + + def fuse_lora_to_base_model(self, model: torch.nn.Module, state_dict_lora, alpha=1.0): + super().fuse_lora_to_base_model(model, state_dict_lora, alpha) + + def convert_state_dict(self, state_dict): + + def guess_block_id(name,model_resource): + if model_resource == 'civitai': + names = name.split("_") + for i in names: + if i.isdigit(): + return i, name.replace(f"_{i}_", "_blockid_") + if model_resource == 'diffusers': + names = name.split(".") + for i in names: + if i.isdigit(): + return i, name.replace(f"transformer_blocks.{i}.", "transformer_blocks.blockid.") + return None, None + + def guess_resource(state_dict): + for k in state_dict: + if "lora_unet_" in k: + return 'civitai' + elif k.startswith("transformer."): + return 'diffusers' + else: + None + + model_resource = guess_resource(state_dict) + if model_resource is None: + return state_dict + + rename_dict = self.diffusers_rename_dict if model_resource == 'diffusers' else self.civitai_rename_dict + def guess_alpha(state_dict): + for name, param in state_dict.items(): + if ".alpha" in name: + for suffix in [".lora_down.weight", ".lora_A.weight"]: + name_ = name.replace(".alpha", suffix) + if name_ in state_dict: + lora_alpha = param.item() / state_dict[name_].shape[0] + lora_alpha = math.sqrt(lora_alpha) + return lora_alpha + + return 1 + + alpha = guess_alpha(state_dict) + + state_dict_ = {} + for name, param in state_dict.items(): + block_id, source_name = guess_block_id(name,model_resource) + if alpha != 1: + param *= alpha + if source_name in rename_dict: + target_name = rename_dict[source_name] + target_name = target_name.replace(".blockid.", f".{block_id}.") + state_dict_[target_name] = param + else: + state_dict_[name] = param + + if model_resource == 'diffusers': + for name in list(state_dict_.keys()): + if "single_blocks." in name and ".a_to_q." in name: + mlp = state_dict_.get(name.replace(".a_to_q.", ".proj_in_besides_attn."), None) + if mlp is None: + dim = 4 + if 'lora_A' in name: + dim = 1 + mlp = torch.zeros(dim * state_dict_[name].shape[0], + *state_dict_[name].shape[1:], + dtype=state_dict_[name].dtype) + else: + state_dict_.pop(name.replace(".a_to_q.", ".proj_in_besides_attn.")) + if 'lora_A' in name: + param = torch.concat([ + state_dict_.pop(name), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")), + mlp, + ], dim=0) + elif 'lora_B' in name: + d, r = state_dict_[name].shape + param = torch.zeros((3*d+mlp.shape[0], 3*r+mlp.shape[1]), dtype=state_dict_[name].dtype, device=state_dict_[name].device) + param[:d, :r] = state_dict_.pop(name) + param[d:2*d, r:2*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")) + param[2*d:3*d, 2*r:3*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")) + param[3*d:, 3*r:] = mlp + else: + param = torch.concat([ + state_dict_.pop(name), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")), + state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")), + mlp, + ], dim=0) + name_ = name.replace(".a_to_q.", ".to_qkv_mlp.") + state_dict_[name_] = param + for name in list(state_dict_.keys()): + for component in ["a", "b"]: + if f".{component}_to_q." in name: + name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.") + concat_dim = 0 + if 'lora_A' in name: + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + elif 'lora_B' in name: + origin = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")] + d, r = origin.shape + # print(d, r) + param = torch.zeros((3*d, 3*r), dtype=origin.dtype, device=origin.device) + param[:d, :r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")] + param[d:2*d, r:2*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")] + param[2*d:3*d, 2*r:3*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")] + else: + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v.")) + return state_dict_ diff --git a/diffsynth/utils/lora/general.py b/diffsynth/utils/lora/general.py new file mode 100644 index 0000000000000000000000000000000000000000..624549d518fb8f2a43b04625b268cbab4441a21a --- /dev/null +++ b/diffsynth/utils/lora/general.py @@ -0,0 +1,62 @@ +import torch + + +class GeneralLoRALoader: + def __init__(self, device="cpu", torch_dtype=torch.float32): + self.device = device + self.torch_dtype = torch_dtype + + + def get_name_dict(self, lora_state_dict): + lora_name_dict = {} + for key in lora_state_dict: + if ".lora_up." in key: + lora_A_key = "lora_down" + lora_B_key = "lora_up" + else: + lora_A_key = "lora_A" + lora_B_key = "lora_B" + if lora_B_key not in key: + continue + keys = key.split(".") + if len(keys) > keys.index(lora_B_key) + 2: + keys.pop(keys.index(lora_B_key) + 1) + keys.pop(keys.index(lora_B_key)) + if keys[0] == "diffusion_model": + keys.pop(0) + keys.pop(-1) + target_name = ".".join(keys) + lora_name_dict[target_name] = (key, key.replace(lora_B_key, lora_A_key)) + return lora_name_dict + + + def convert_state_dict(self, state_dict, suffix=".weight"): + name_dict = self.get_name_dict(state_dict) + state_dict_ = {} + for name in name_dict: + weight_up = state_dict[name_dict[name][0]] + weight_down = state_dict[name_dict[name][1]] + state_dict_[name + f".lora_B{suffix}"] = weight_up + state_dict_[name + f".lora_A{suffix}"] = weight_down + return state_dict_ + + + def fuse_lora_to_base_model(self, model: torch.nn.Module, state_dict, alpha=1.0): + updated_num = 0 + state_dict = self.convert_state_dict(state_dict) + lora_layer_names = set([i.replace(".lora_B.weight", "") for i in state_dict if i.endswith(".lora_B.weight")]) + for name, module in model.named_modules(): + if name in lora_layer_names: + weight_up = state_dict[name + ".lora_B.weight"].to(device=self.device, dtype=self.torch_dtype) + weight_down = state_dict[name + ".lora_A.weight"].to(device=self.device, dtype=self.torch_dtype) + if len(weight_up.shape) == 4: + weight_up = weight_up.squeeze(3).squeeze(2) + weight_down = weight_down.squeeze(3).squeeze(2) + weight_lora = alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) + else: + weight_lora = alpha * torch.mm(weight_up, weight_down) + state_dict_base = module.state_dict() + state_dict_base["weight"] = state_dict_base["weight"].to(device=self.device, dtype=self.torch_dtype) + weight_lora + module.load_state_dict(state_dict_base) + updated_num += 1 + print(f"{updated_num} tensors are fused by LoRA. Fused LoRA layers cannot be cleared by `pipe.clear_lora()`.") diff --git a/diffsynth/utils/lora/merge.py b/diffsynth/utils/lora/merge.py new file mode 100644 index 0000000000000000000000000000000000000000..61904ff4bcebc6c344c23f26073aec292355217c --- /dev/null +++ b/diffsynth/utils/lora/merge.py @@ -0,0 +1,20 @@ +import torch +from typing import Dict, List + + +def merge_lora_weight(tensors_A, tensors_B): + lora_A = torch.concat(tensors_A, dim=0) + lora_B = torch.concat(tensors_B, dim=1) + return lora_A, lora_B + + +def merge_lora(loras: List[Dict[str, torch.Tensor]], alpha=1): + lora_merged = {} + keys = [i for i in loras[0].keys() if ".lora_A." in i] + for key in keys: + tensors_A = [lora[key] for lora in loras] + tensors_B = [lora[key.replace(".lora_A.", ".lora_B.")] for lora in loras] + lora_A, lora_B = merge_lora_weight(tensors_A, tensors_B) + lora_merged[key] = lora_A * alpha + lora_merged[key.replace(".lora_A.", ".lora_B.")] = lora_B + return lora_merged diff --git a/diffsynth/utils/state_dict_converters/__init__.py b/diffsynth/utils/state_dict_converters/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/diffsynth/utils/state_dict_converters/flux2_text_encoder.py b/diffsynth/utils/state_dict_converters/flux2_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..0975e62a35021c697192ad054f0e3aff42289292 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux2_text_encoder.py @@ -0,0 +1,17 @@ +def Flux2TextEncoderStateDictConverter(state_dict): + rename_dict = { + "multi_modal_projector.linear_1.weight": "model.multi_modal_projector.linear_1.weight", + "multi_modal_projector.linear_2.weight": "model.multi_modal_projector.linear_2.weight", + "multi_modal_projector.norm.weight": "model.multi_modal_projector.norm.weight", + "multi_modal_projector.patch_merger.merging_layer.weight": "model.multi_modal_projector.patch_merger.merging_layer.weight", + "language_model.lm_head.weight": "lm_head.weight", + } + state_dict_ = {} + for k in state_dict: + k_ = k + k_ = k_.replace("language_model.model", "model.language_model") + k_ = k_.replace("vision_tower", "model.vision_tower") + if k_ in rename_dict: + k_ = rename_dict[k_] + state_dict_[k_] = state_dict[k] + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/flux_controlnet.py b/diffsynth/utils/state_dict_converters/flux_controlnet.py new file mode 100644 index 0000000000000000000000000000000000000000..15f9447d22bc0ebc2dbb3d2eac8dbf0bd78e4151 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_controlnet.py @@ -0,0 +1,103 @@ +import torch + + +def FluxControlNetStateDictConverter(state_dict): + global_rename_dict = { + "context_embedder": "context_embedder", + "x_embedder": "x_embedder", + "time_text_embed.timestep_embedder.linear_1": "time_embedder.timestep_embedder.0", + "time_text_embed.timestep_embedder.linear_2": "time_embedder.timestep_embedder.2", + "time_text_embed.guidance_embedder.linear_1": "guidance_embedder.timestep_embedder.0", + "time_text_embed.guidance_embedder.linear_2": "guidance_embedder.timestep_embedder.2", + "time_text_embed.text_embedder.linear_1": "pooled_text_embedder.0", + "time_text_embed.text_embedder.linear_2": "pooled_text_embedder.2", + "norm_out.linear": "final_norm_out.linear", + "proj_out": "final_proj_out", + } + rename_dict = { + "proj_out": "proj_out", + "norm1.linear": "norm1_a.linear", + "norm1_context.linear": "norm1_b.linear", + "attn.to_q": "attn.a_to_q", + "attn.to_k": "attn.a_to_k", + "attn.to_v": "attn.a_to_v", + "attn.to_out.0": "attn.a_to_out", + "attn.add_q_proj": "attn.b_to_q", + "attn.add_k_proj": "attn.b_to_k", + "attn.add_v_proj": "attn.b_to_v", + "attn.to_add_out": "attn.b_to_out", + "ff.net.0.proj": "ff_a.0", + "ff.net.2": "ff_a.2", + "ff_context.net.0.proj": "ff_b.0", + "ff_context.net.2": "ff_b.2", + "attn.norm_q": "attn.norm_q_a", + "attn.norm_k": "attn.norm_k_a", + "attn.norm_added_q": "attn.norm_q_b", + "attn.norm_added_k": "attn.norm_k_b", + } + rename_dict_single = { + "attn.to_q": "a_to_q", + "attn.to_k": "a_to_k", + "attn.to_v": "a_to_v", + "attn.norm_q": "norm_q_a", + "attn.norm_k": "norm_k_a", + "norm.linear": "norm.linear", + "proj_mlp": "proj_in_besides_attn", + "proj_out": "proj_out", + } + state_dict_ = {} + + for name in state_dict: + param = state_dict[name] + if name.endswith(".weight") or name.endswith(".bias"): + suffix = ".weight" if name.endswith(".weight") else ".bias" + prefix = name[:-len(suffix)] + if prefix in global_rename_dict: + state_dict_[global_rename_dict[prefix] + suffix] = param + elif prefix.startswith("transformer_blocks."): + names = prefix.split(".") + names[0] = "blocks" + middle = ".".join(names[2:]) + if middle in rename_dict: + name_ = ".".join(names[:2] + [rename_dict[middle]] + [suffix[1:]]) + state_dict_[name_] = param + elif prefix.startswith("single_transformer_blocks."): + names = prefix.split(".") + names[0] = "single_blocks" + middle = ".".join(names[2:]) + if middle in rename_dict_single: + name_ = ".".join(names[:2] + [rename_dict_single[middle]] + [suffix[1:]]) + state_dict_[name_] = param + else: + state_dict_[name] = param + else: + state_dict_[name] = param + for name in list(state_dict_.keys()): + if ".proj_in_besides_attn." in name: + name_ = name.replace(".proj_in_besides_attn.", ".to_qkv_mlp.") + param = torch.concat([ + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_q.")], + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_k.")], + state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_v.")], + state_dict_[name], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_q.")) + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_k.")) + state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_v.")) + state_dict_.pop(name) + for name in list(state_dict_.keys()): + for component in ["a", "b"]: + if f".{component}_to_q." in name: + name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.") + param = torch.concat([ + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")], + state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")], + ], dim=0) + state_dict_[name_] = param + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k.")) + state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v.")) + + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/flux_dit.py b/diffsynth/utils/state_dict_converters/flux_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..0fbe460e2554d4be71a5e55ff2b38f000c4cc041 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_dit.py @@ -0,0 +1,92 @@ +import torch + + +def FluxDiTStateDictConverter(state_dict): + is_nexus_gen = sum([key.startswith("pipe.dit.") for key in state_dict]) > 0 + if is_nexus_gen: + dit_state_dict = {} + for key in state_dict: + if key.startswith('pipe.dit.'): + param = state_dict[key] + new_key = key.replace("pipe.dit.", "") + if new_key.startswith("final_norm_out.linear."): + param = torch.concat([param[3072:], param[:3072]], dim=0) + dit_state_dict[new_key] = param + return dit_state_dict + + rename_dict = { + "time_in.in_layer.bias": "time_embedder.timestep_embedder.0.bias", + "time_in.in_layer.weight": "time_embedder.timestep_embedder.0.weight", + "time_in.out_layer.bias": "time_embedder.timestep_embedder.2.bias", + "time_in.out_layer.weight": "time_embedder.timestep_embedder.2.weight", + "txt_in.bias": "context_embedder.bias", + "txt_in.weight": "context_embedder.weight", + "vector_in.in_layer.bias": "pooled_text_embedder.0.bias", + "vector_in.in_layer.weight": "pooled_text_embedder.0.weight", + "vector_in.out_layer.bias": "pooled_text_embedder.2.bias", + "vector_in.out_layer.weight": "pooled_text_embedder.2.weight", + "final_layer.linear.bias": "final_proj_out.bias", + "final_layer.linear.weight": "final_proj_out.weight", + "guidance_in.in_layer.bias": "guidance_embedder.timestep_embedder.0.bias", + "guidance_in.in_layer.weight": "guidance_embedder.timestep_embedder.0.weight", + "guidance_in.out_layer.bias": "guidance_embedder.timestep_embedder.2.bias", + "guidance_in.out_layer.weight": "guidance_embedder.timestep_embedder.2.weight", + "img_in.bias": "x_embedder.bias", + "img_in.weight": "x_embedder.weight", + "final_layer.adaLN_modulation.1.weight": "final_norm_out.linear.weight", + "final_layer.adaLN_modulation.1.bias": "final_norm_out.linear.bias", + } + suffix_rename_dict = { + "img_attn.norm.key_norm.scale": "attn.norm_k_a.weight", + "img_attn.norm.query_norm.scale": "attn.norm_q_a.weight", + "img_attn.proj.bias": "attn.a_to_out.bias", + "img_attn.proj.weight": "attn.a_to_out.weight", + "img_attn.qkv.bias": "attn.a_to_qkv.bias", + "img_attn.qkv.weight": "attn.a_to_qkv.weight", + "img_mlp.0.bias": "ff_a.0.bias", + "img_mlp.0.weight": "ff_a.0.weight", + "img_mlp.2.bias": "ff_a.2.bias", + "img_mlp.2.weight": "ff_a.2.weight", + "img_mod.lin.bias": "norm1_a.linear.bias", + "img_mod.lin.weight": "norm1_a.linear.weight", + "txt_attn.norm.key_norm.scale": "attn.norm_k_b.weight", + "txt_attn.norm.query_norm.scale": "attn.norm_q_b.weight", + "txt_attn.proj.bias": "attn.b_to_out.bias", + "txt_attn.proj.weight": "attn.b_to_out.weight", + "txt_attn.qkv.bias": "attn.b_to_qkv.bias", + "txt_attn.qkv.weight": "attn.b_to_qkv.weight", + "txt_mlp.0.bias": "ff_b.0.bias", + "txt_mlp.0.weight": "ff_b.0.weight", + "txt_mlp.2.bias": "ff_b.2.bias", + "txt_mlp.2.weight": "ff_b.2.weight", + "txt_mod.lin.bias": "norm1_b.linear.bias", + "txt_mod.lin.weight": "norm1_b.linear.weight", + + "linear1.bias": "to_qkv_mlp.bias", + "linear1.weight": "to_qkv_mlp.weight", + "linear2.bias": "proj_out.bias", + "linear2.weight": "proj_out.weight", + "modulation.lin.bias": "norm.linear.bias", + "modulation.lin.weight": "norm.linear.weight", + "norm.key_norm.scale": "norm_k_a.weight", + "norm.query_norm.scale": "norm_q_a.weight", + } + state_dict_ = {} + for name in state_dict: + original_name = name + if name.startswith("model.diffusion_model."): + name = name[len("model.diffusion_model."):] + names = name.split(".") + if name in rename_dict: + rename = rename_dict[name] + state_dict_[rename] = state_dict[original_name] + elif names[0] == "double_blocks": + rename = f"blocks.{names[1]}." + suffix_rename_dict[".".join(names[2:])] + state_dict_[rename] = state_dict[original_name] + elif names[0] == "single_blocks": + if ".".join(names[2:]) in suffix_rename_dict: + rename = f"single_blocks.{names[1]}." + suffix_rename_dict[".".join(names[2:])] + state_dict_[rename] = state_dict[original_name] + else: + pass + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/flux_infiniteyou.py b/diffsynth/utils/state_dict_converters/flux_infiniteyou.py new file mode 100644 index 0000000000000000000000000000000000000000..7025b392d54c5b4844ed3b3387bd010217897f4a --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_infiniteyou.py @@ -0,0 +1,2 @@ +def FluxInfiniteYouImageProjectorStateDictConverter(state_dict): + return state_dict['image_proj'] \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/flux_ipadapter.py b/diffsynth/utils/state_dict_converters/flux_ipadapter.py new file mode 100644 index 0000000000000000000000000000000000000000..86dfb133655fbe9c33c84b419706a103cec96b1b --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_ipadapter.py @@ -0,0 +1,32 @@ +def FluxIpAdapterStateDictConverter(state_dict): + state_dict_ = {} + + if "ip_adapter" in state_dict and isinstance(state_dict["ip_adapter"], dict): + for name, param in state_dict["ip_adapter"].items(): + name_ = 'ipadapter_modules.' + name + state_dict_[name_] = param + + if "image_proj" in state_dict: + for name, param in state_dict["image_proj"].items(): + name_ = "image_proj." + name + state_dict_[name_] = param + return state_dict_ + + for key, value in state_dict.items(): + if key.startswith("image_proj."): + state_dict_[key] = value + elif key.startswith("ip_adapter."): + new_key = key.replace("ip_adapter.", "ipadapter_modules.") + state_dict_[new_key] = value + else: + pass + + return state_dict_ + + +def SiglipStateDictConverter(state_dict): + new_state_dict = {} + for key in state_dict: + if key.startswith("vision_model."): + new_state_dict[key] = state_dict[key] + return new_state_dict \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/flux_text_encoder_clip.py b/diffsynth/utils/state_dict_converters/flux_text_encoder_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..aa018aa5c570cc67f4856002e8f1f83f18998e07 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_text_encoder_clip.py @@ -0,0 +1,31 @@ +def FluxTextEncoderClipStateDictConverter(state_dict): + rename_dict = { + "text_model.embeddings.token_embedding.weight": "token_embedding.weight", + "text_model.embeddings.position_embedding.weight": "position_embeds", + "text_model.final_layer_norm.weight": "final_layer_norm.weight", + "text_model.final_layer_norm.bias": "final_layer_norm.bias", + } + attn_rename_dict = { + "self_attn.q_proj": "attn.to_q", + "self_attn.k_proj": "attn.to_k", + "self_attn.v_proj": "attn.to_v", + "self_attn.out_proj": "attn.to_out", + "layer_norm1": "layer_norm1", + "layer_norm2": "layer_norm2", + "mlp.fc1": "fc1", + "mlp.fc2": "fc2", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + if name == "text_model.embeddings.position_embedding.weight": + param = param.reshape((1, param.shape[0], param.shape[1])) + state_dict_[rename_dict[name]] = param + elif name.startswith("text_model.encoder.layers."): + param = state_dict[name] + names = name.split(".") + layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1] + name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail]) + state_dict_[name_] = param + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/flux_text_encoder_t5.py b/diffsynth/utils/state_dict_converters/flux_text_encoder_t5.py new file mode 100644 index 0000000000000000000000000000000000000000..d35eb831d2a7b1d48eee747d251d6cfb6ad508ef --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_text_encoder_t5.py @@ -0,0 +1,4 @@ +def FluxTextEncoderT5StateDictConverter(state_dict): + state_dict_ = {i: state_dict[i] for i in state_dict} + state_dict_["encoder.embed_tokens.weight"] = state_dict["shared.weight"] + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/flux_vae.py b/diffsynth/utils/state_dict_converters/flux_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..6547f18f1e1cfe69d0cf4ef43860702812d25fab --- /dev/null +++ b/diffsynth/utils/state_dict_converters/flux_vae.py @@ -0,0 +1,382 @@ +def FluxVAEEncoderStateDictConverter(state_dict): + rename_dict = { + "encoder.conv_in.bias": "conv_in.bias", + "encoder.conv_in.weight": "conv_in.weight", + "encoder.conv_out.bias": "conv_out.bias", + "encoder.conv_out.weight": "conv_out.weight", + "encoder.down.0.block.0.conv1.bias": "blocks.0.conv1.bias", + "encoder.down.0.block.0.conv1.weight": "blocks.0.conv1.weight", + "encoder.down.0.block.0.conv2.bias": "blocks.0.conv2.bias", + "encoder.down.0.block.0.conv2.weight": "blocks.0.conv2.weight", + "encoder.down.0.block.0.norm1.bias": "blocks.0.norm1.bias", + "encoder.down.0.block.0.norm1.weight": "blocks.0.norm1.weight", + "encoder.down.0.block.0.norm2.bias": "blocks.0.norm2.bias", + "encoder.down.0.block.0.norm2.weight": "blocks.0.norm2.weight", + "encoder.down.0.block.1.conv1.bias": "blocks.1.conv1.bias", + "encoder.down.0.block.1.conv1.weight": "blocks.1.conv1.weight", + "encoder.down.0.block.1.conv2.bias": "blocks.1.conv2.bias", + "encoder.down.0.block.1.conv2.weight": "blocks.1.conv2.weight", + "encoder.down.0.block.1.norm1.bias": "blocks.1.norm1.bias", + "encoder.down.0.block.1.norm1.weight": "blocks.1.norm1.weight", + "encoder.down.0.block.1.norm2.bias": "blocks.1.norm2.bias", + "encoder.down.0.block.1.norm2.weight": "blocks.1.norm2.weight", + "encoder.down.0.downsample.conv.bias": "blocks.2.conv.bias", + "encoder.down.0.downsample.conv.weight": "blocks.2.conv.weight", + "encoder.down.1.block.0.conv1.bias": "blocks.3.conv1.bias", + "encoder.down.1.block.0.conv1.weight": "blocks.3.conv1.weight", + "encoder.down.1.block.0.conv2.bias": "blocks.3.conv2.bias", + "encoder.down.1.block.0.conv2.weight": "blocks.3.conv2.weight", + "encoder.down.1.block.0.nin_shortcut.bias": "blocks.3.conv_shortcut.bias", + "encoder.down.1.block.0.nin_shortcut.weight": "blocks.3.conv_shortcut.weight", + "encoder.down.1.block.0.norm1.bias": "blocks.3.norm1.bias", + "encoder.down.1.block.0.norm1.weight": "blocks.3.norm1.weight", + "encoder.down.1.block.0.norm2.bias": "blocks.3.norm2.bias", + "encoder.down.1.block.0.norm2.weight": "blocks.3.norm2.weight", + "encoder.down.1.block.1.conv1.bias": "blocks.4.conv1.bias", + "encoder.down.1.block.1.conv1.weight": "blocks.4.conv1.weight", + "encoder.down.1.block.1.conv2.bias": "blocks.4.conv2.bias", + "encoder.down.1.block.1.conv2.weight": "blocks.4.conv2.weight", + "encoder.down.1.block.1.norm1.bias": "blocks.4.norm1.bias", + "encoder.down.1.block.1.norm1.weight": "blocks.4.norm1.weight", + "encoder.down.1.block.1.norm2.bias": "blocks.4.norm2.bias", + "encoder.down.1.block.1.norm2.weight": "blocks.4.norm2.weight", + "encoder.down.1.downsample.conv.bias": "blocks.5.conv.bias", + "encoder.down.1.downsample.conv.weight": "blocks.5.conv.weight", + "encoder.down.2.block.0.conv1.bias": "blocks.6.conv1.bias", + "encoder.down.2.block.0.conv1.weight": "blocks.6.conv1.weight", + "encoder.down.2.block.0.conv2.bias": "blocks.6.conv2.bias", + "encoder.down.2.block.0.conv2.weight": "blocks.6.conv2.weight", + "encoder.down.2.block.0.nin_shortcut.bias": "blocks.6.conv_shortcut.bias", + "encoder.down.2.block.0.nin_shortcut.weight": "blocks.6.conv_shortcut.weight", + "encoder.down.2.block.0.norm1.bias": "blocks.6.norm1.bias", + "encoder.down.2.block.0.norm1.weight": "blocks.6.norm1.weight", + "encoder.down.2.block.0.norm2.bias": "blocks.6.norm2.bias", + "encoder.down.2.block.0.norm2.weight": "blocks.6.norm2.weight", + "encoder.down.2.block.1.conv1.bias": "blocks.7.conv1.bias", + "encoder.down.2.block.1.conv1.weight": "blocks.7.conv1.weight", + "encoder.down.2.block.1.conv2.bias": "blocks.7.conv2.bias", + "encoder.down.2.block.1.conv2.weight": "blocks.7.conv2.weight", + "encoder.down.2.block.1.norm1.bias": "blocks.7.norm1.bias", + "encoder.down.2.block.1.norm1.weight": "blocks.7.norm1.weight", + "encoder.down.2.block.1.norm2.bias": "blocks.7.norm2.bias", + "encoder.down.2.block.1.norm2.weight": "blocks.7.norm2.weight", + "encoder.down.2.downsample.conv.bias": "blocks.8.conv.bias", + "encoder.down.2.downsample.conv.weight": "blocks.8.conv.weight", + "encoder.down.3.block.0.conv1.bias": "blocks.9.conv1.bias", + "encoder.down.3.block.0.conv1.weight": "blocks.9.conv1.weight", + "encoder.down.3.block.0.conv2.bias": "blocks.9.conv2.bias", + "encoder.down.3.block.0.conv2.weight": "blocks.9.conv2.weight", + "encoder.down.3.block.0.norm1.bias": "blocks.9.norm1.bias", + "encoder.down.3.block.0.norm1.weight": "blocks.9.norm1.weight", + "encoder.down.3.block.0.norm2.bias": "blocks.9.norm2.bias", + "encoder.down.3.block.0.norm2.weight": "blocks.9.norm2.weight", + "encoder.down.3.block.1.conv1.bias": "blocks.10.conv1.bias", + "encoder.down.3.block.1.conv1.weight": "blocks.10.conv1.weight", + "encoder.down.3.block.1.conv2.bias": "blocks.10.conv2.bias", + "encoder.down.3.block.1.conv2.weight": "blocks.10.conv2.weight", + "encoder.down.3.block.1.norm1.bias": "blocks.10.norm1.bias", + "encoder.down.3.block.1.norm1.weight": "blocks.10.norm1.weight", + "encoder.down.3.block.1.norm2.bias": "blocks.10.norm2.bias", + "encoder.down.3.block.1.norm2.weight": "blocks.10.norm2.weight", + "encoder.mid.attn_1.k.bias": "blocks.12.transformer_blocks.0.to_k.bias", + "encoder.mid.attn_1.k.weight": "blocks.12.transformer_blocks.0.to_k.weight", + "encoder.mid.attn_1.norm.bias": "blocks.12.norm.bias", + "encoder.mid.attn_1.norm.weight": "blocks.12.norm.weight", + "encoder.mid.attn_1.proj_out.bias": "blocks.12.transformer_blocks.0.to_out.bias", + "encoder.mid.attn_1.proj_out.weight": "blocks.12.transformer_blocks.0.to_out.weight", + "encoder.mid.attn_1.q.bias": "blocks.12.transformer_blocks.0.to_q.bias", + "encoder.mid.attn_1.q.weight": "blocks.12.transformer_blocks.0.to_q.weight", + "encoder.mid.attn_1.v.bias": "blocks.12.transformer_blocks.0.to_v.bias", + "encoder.mid.attn_1.v.weight": "blocks.12.transformer_blocks.0.to_v.weight", + "encoder.mid.block_1.conv1.bias": "blocks.11.conv1.bias", + "encoder.mid.block_1.conv1.weight": "blocks.11.conv1.weight", + "encoder.mid.block_1.conv2.bias": "blocks.11.conv2.bias", + "encoder.mid.block_1.conv2.weight": "blocks.11.conv2.weight", + "encoder.mid.block_1.norm1.bias": "blocks.11.norm1.bias", + "encoder.mid.block_1.norm1.weight": "blocks.11.norm1.weight", + "encoder.mid.block_1.norm2.bias": "blocks.11.norm2.bias", + "encoder.mid.block_1.norm2.weight": "blocks.11.norm2.weight", + "encoder.mid.block_2.conv1.bias": "blocks.13.conv1.bias", + "encoder.mid.block_2.conv1.weight": "blocks.13.conv1.weight", + "encoder.mid.block_2.conv2.bias": "blocks.13.conv2.bias", + "encoder.mid.block_2.conv2.weight": "blocks.13.conv2.weight", + "encoder.mid.block_2.norm1.bias": "blocks.13.norm1.bias", + "encoder.mid.block_2.norm1.weight": "blocks.13.norm1.weight", + "encoder.mid.block_2.norm2.bias": "blocks.13.norm2.bias", + "encoder.mid.block_2.norm2.weight": "blocks.13.norm2.weight", + "encoder.norm_out.bias": "conv_norm_out.bias", + "encoder.norm_out.weight": "conv_norm_out.weight", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + state_dict_[rename_dict[name]] = param + return state_dict_ + + +def FluxVAEDecoderStateDictConverter(state_dict): + rename_dict = { + "decoder.conv_in.bias": "conv_in.bias", + "decoder.conv_in.weight": "conv_in.weight", + "decoder.conv_out.bias": "conv_out.bias", + "decoder.conv_out.weight": "conv_out.weight", + "decoder.mid.attn_1.k.bias": "blocks.1.transformer_blocks.0.to_k.bias", + "decoder.mid.attn_1.k.weight": "blocks.1.transformer_blocks.0.to_k.weight", + "decoder.mid.attn_1.norm.bias": "blocks.1.norm.bias", + "decoder.mid.attn_1.norm.weight": "blocks.1.norm.weight", + "decoder.mid.attn_1.proj_out.bias": "blocks.1.transformer_blocks.0.to_out.bias", + "decoder.mid.attn_1.proj_out.weight": "blocks.1.transformer_blocks.0.to_out.weight", + "decoder.mid.attn_1.q.bias": "blocks.1.transformer_blocks.0.to_q.bias", + "decoder.mid.attn_1.q.weight": "blocks.1.transformer_blocks.0.to_q.weight", + "decoder.mid.attn_1.v.bias": "blocks.1.transformer_blocks.0.to_v.bias", + "decoder.mid.attn_1.v.weight": "blocks.1.transformer_blocks.0.to_v.weight", + "decoder.mid.block_1.conv1.bias": "blocks.0.conv1.bias", + "decoder.mid.block_1.conv1.weight": "blocks.0.conv1.weight", + "decoder.mid.block_1.conv2.bias": "blocks.0.conv2.bias", + "decoder.mid.block_1.conv2.weight": "blocks.0.conv2.weight", + "decoder.mid.block_1.norm1.bias": "blocks.0.norm1.bias", + "decoder.mid.block_1.norm1.weight": "blocks.0.norm1.weight", + "decoder.mid.block_1.norm2.bias": "blocks.0.norm2.bias", + "decoder.mid.block_1.norm2.weight": "blocks.0.norm2.weight", + "decoder.mid.block_2.conv1.bias": "blocks.2.conv1.bias", + "decoder.mid.block_2.conv1.weight": "blocks.2.conv1.weight", + "decoder.mid.block_2.conv2.bias": "blocks.2.conv2.bias", + "decoder.mid.block_2.conv2.weight": "blocks.2.conv2.weight", + "decoder.mid.block_2.norm1.bias": "blocks.2.norm1.bias", + "decoder.mid.block_2.norm1.weight": "blocks.2.norm1.weight", + "decoder.mid.block_2.norm2.bias": "blocks.2.norm2.bias", + "decoder.mid.block_2.norm2.weight": "blocks.2.norm2.weight", + "decoder.norm_out.bias": "conv_norm_out.bias", + "decoder.norm_out.weight": "conv_norm_out.weight", + "decoder.up.0.block.0.conv1.bias": "blocks.15.conv1.bias", + "decoder.up.0.block.0.conv1.weight": "blocks.15.conv1.weight", + "decoder.up.0.block.0.conv2.bias": "blocks.15.conv2.bias", + "decoder.up.0.block.0.conv2.weight": "blocks.15.conv2.weight", + "decoder.up.0.block.0.nin_shortcut.bias": "blocks.15.conv_shortcut.bias", + "decoder.up.0.block.0.nin_shortcut.weight": "blocks.15.conv_shortcut.weight", + "decoder.up.0.block.0.norm1.bias": "blocks.15.norm1.bias", + "decoder.up.0.block.0.norm1.weight": "blocks.15.norm1.weight", + "decoder.up.0.block.0.norm2.bias": "blocks.15.norm2.bias", + "decoder.up.0.block.0.norm2.weight": "blocks.15.norm2.weight", + "decoder.up.0.block.1.conv1.bias": "blocks.16.conv1.bias", + "decoder.up.0.block.1.conv1.weight": "blocks.16.conv1.weight", + "decoder.up.0.block.1.conv2.bias": "blocks.16.conv2.bias", + "decoder.up.0.block.1.conv2.weight": "blocks.16.conv2.weight", + "decoder.up.0.block.1.norm1.bias": "blocks.16.norm1.bias", + "decoder.up.0.block.1.norm1.weight": "blocks.16.norm1.weight", + "decoder.up.0.block.1.norm2.bias": "blocks.16.norm2.bias", + "decoder.up.0.block.1.norm2.weight": "blocks.16.norm2.weight", + "decoder.up.0.block.2.conv1.bias": "blocks.17.conv1.bias", + "decoder.up.0.block.2.conv1.weight": "blocks.17.conv1.weight", + "decoder.up.0.block.2.conv2.bias": "blocks.17.conv2.bias", + "decoder.up.0.block.2.conv2.weight": "blocks.17.conv2.weight", + "decoder.up.0.block.2.norm1.bias": "blocks.17.norm1.bias", + "decoder.up.0.block.2.norm1.weight": "blocks.17.norm1.weight", + "decoder.up.0.block.2.norm2.bias": "blocks.17.norm2.bias", + "decoder.up.0.block.2.norm2.weight": "blocks.17.norm2.weight", + "decoder.up.1.block.0.conv1.bias": "blocks.11.conv1.bias", + "decoder.up.1.block.0.conv1.weight": "blocks.11.conv1.weight", + "decoder.up.1.block.0.conv2.bias": "blocks.11.conv2.bias", + "decoder.up.1.block.0.conv2.weight": "blocks.11.conv2.weight", + "decoder.up.1.block.0.nin_shortcut.bias": "blocks.11.conv_shortcut.bias", + "decoder.up.1.block.0.nin_shortcut.weight": "blocks.11.conv_shortcut.weight", + "decoder.up.1.block.0.norm1.bias": "blocks.11.norm1.bias", + "decoder.up.1.block.0.norm1.weight": "blocks.11.norm1.weight", + "decoder.up.1.block.0.norm2.bias": "blocks.11.norm2.bias", + "decoder.up.1.block.0.norm2.weight": "blocks.11.norm2.weight", + "decoder.up.1.block.1.conv1.bias": "blocks.12.conv1.bias", + "decoder.up.1.block.1.conv1.weight": "blocks.12.conv1.weight", + "decoder.up.1.block.1.conv2.bias": "blocks.12.conv2.bias", + "decoder.up.1.block.1.conv2.weight": "blocks.12.conv2.weight", + "decoder.up.1.block.1.norm1.bias": "blocks.12.norm1.bias", + "decoder.up.1.block.1.norm1.weight": "blocks.12.norm1.weight", + "decoder.up.1.block.1.norm2.bias": "blocks.12.norm2.bias", + "decoder.up.1.block.1.norm2.weight": "blocks.12.norm2.weight", + "decoder.up.1.block.2.conv1.bias": "blocks.13.conv1.bias", + "decoder.up.1.block.2.conv1.weight": "blocks.13.conv1.weight", + "decoder.up.1.block.2.conv2.bias": "blocks.13.conv2.bias", + "decoder.up.1.block.2.conv2.weight": "blocks.13.conv2.weight", + "decoder.up.1.block.2.norm1.bias": "blocks.13.norm1.bias", + "decoder.up.1.block.2.norm1.weight": "blocks.13.norm1.weight", + "decoder.up.1.block.2.norm2.bias": "blocks.13.norm2.bias", + "decoder.up.1.block.2.norm2.weight": "blocks.13.norm2.weight", + "decoder.up.1.upsample.conv.bias": "blocks.14.conv.bias", + "decoder.up.1.upsample.conv.weight": "blocks.14.conv.weight", + "decoder.up.2.block.0.conv1.bias": "blocks.7.conv1.bias", + "decoder.up.2.block.0.conv1.weight": "blocks.7.conv1.weight", + "decoder.up.2.block.0.conv2.bias": "blocks.7.conv2.bias", + "decoder.up.2.block.0.conv2.weight": "blocks.7.conv2.weight", + "decoder.up.2.block.0.norm1.bias": "blocks.7.norm1.bias", + "decoder.up.2.block.0.norm1.weight": "blocks.7.norm1.weight", + "decoder.up.2.block.0.norm2.bias": "blocks.7.norm2.bias", + "decoder.up.2.block.0.norm2.weight": "blocks.7.norm2.weight", + "decoder.up.2.block.1.conv1.bias": "blocks.8.conv1.bias", + "decoder.up.2.block.1.conv1.weight": "blocks.8.conv1.weight", + "decoder.up.2.block.1.conv2.bias": "blocks.8.conv2.bias", + "decoder.up.2.block.1.conv2.weight": "blocks.8.conv2.weight", + "decoder.up.2.block.1.norm1.bias": "blocks.8.norm1.bias", + "decoder.up.2.block.1.norm1.weight": "blocks.8.norm1.weight", + "decoder.up.2.block.1.norm2.bias": "blocks.8.norm2.bias", + "decoder.up.2.block.1.norm2.weight": "blocks.8.norm2.weight", + "decoder.up.2.block.2.conv1.bias": "blocks.9.conv1.bias", + "decoder.up.2.block.2.conv1.weight": "blocks.9.conv1.weight", + "decoder.up.2.block.2.conv2.bias": "blocks.9.conv2.bias", + "decoder.up.2.block.2.conv2.weight": "blocks.9.conv2.weight", + "decoder.up.2.block.2.norm1.bias": "blocks.9.norm1.bias", + "decoder.up.2.block.2.norm1.weight": "blocks.9.norm1.weight", + "decoder.up.2.block.2.norm2.bias": "blocks.9.norm2.bias", + "decoder.up.2.block.2.norm2.weight": "blocks.9.norm2.weight", + "decoder.up.2.upsample.conv.bias": "blocks.10.conv.bias", + "decoder.up.2.upsample.conv.weight": "blocks.10.conv.weight", + "decoder.up.3.block.0.conv1.bias": "blocks.3.conv1.bias", + "decoder.up.3.block.0.conv1.weight": "blocks.3.conv1.weight", + "decoder.up.3.block.0.conv2.bias": "blocks.3.conv2.bias", + "decoder.up.3.block.0.conv2.weight": "blocks.3.conv2.weight", + "decoder.up.3.block.0.norm1.bias": "blocks.3.norm1.bias", + "decoder.up.3.block.0.norm1.weight": "blocks.3.norm1.weight", + "decoder.up.3.block.0.norm2.bias": "blocks.3.norm2.bias", + "decoder.up.3.block.0.norm2.weight": "blocks.3.norm2.weight", + "decoder.up.3.block.1.conv1.bias": "blocks.4.conv1.bias", + "decoder.up.3.block.1.conv1.weight": "blocks.4.conv1.weight", + "decoder.up.3.block.1.conv2.bias": "blocks.4.conv2.bias", + "decoder.up.3.block.1.conv2.weight": "blocks.4.conv2.weight", + "decoder.up.3.block.1.norm1.bias": "blocks.4.norm1.bias", + "decoder.up.3.block.1.norm1.weight": "blocks.4.norm1.weight", + "decoder.up.3.block.1.norm2.bias": "blocks.4.norm2.bias", + "decoder.up.3.block.1.norm2.weight": "blocks.4.norm2.weight", + "decoder.up.3.block.2.conv1.bias": "blocks.5.conv1.bias", + "decoder.up.3.block.2.conv1.weight": "blocks.5.conv1.weight", + "decoder.up.3.block.2.conv2.bias": "blocks.5.conv2.bias", + "decoder.up.3.block.2.conv2.weight": "blocks.5.conv2.weight", + "decoder.up.3.block.2.norm1.bias": "blocks.5.norm1.bias", + "decoder.up.3.block.2.norm1.weight": "blocks.5.norm1.weight", + "decoder.up.3.block.2.norm2.bias": "blocks.5.norm2.bias", + "decoder.up.3.block.2.norm2.weight": "blocks.5.norm2.weight", + "decoder.up.3.upsample.conv.bias": "blocks.6.conv.bias", + "decoder.up.3.upsample.conv.weight": "blocks.6.conv.weight", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + param = state_dict[name] + state_dict_[rename_dict[name]] = param + return state_dict_ + + +def FluxVAEEncoderStateDictConverterDiffusers(state_dict): + # architecture + block_types = [ + 'ResnetBlock', 'ResnetBlock', 'DownSampler', + 'ResnetBlock', 'ResnetBlock', 'DownSampler', + 'ResnetBlock', 'ResnetBlock', 'DownSampler', + 'ResnetBlock', 'ResnetBlock', + 'ResnetBlock', 'VAEAttentionBlock', 'ResnetBlock' + ] + + # Rename each parameter + local_rename_dict = { + "quant_conv": "quant_conv", + "encoder.conv_in": "conv_in", + "encoder.mid_block.attentions.0.group_norm": "blocks.12.norm", + "encoder.mid_block.attentions.0.to_q": "blocks.12.transformer_blocks.0.to_q", + "encoder.mid_block.attentions.0.to_k": "blocks.12.transformer_blocks.0.to_k", + "encoder.mid_block.attentions.0.to_v": "blocks.12.transformer_blocks.0.to_v", + "encoder.mid_block.attentions.0.to_out.0": "blocks.12.transformer_blocks.0.to_out", + "encoder.mid_block.resnets.0.norm1": "blocks.11.norm1", + "encoder.mid_block.resnets.0.conv1": "blocks.11.conv1", + "encoder.mid_block.resnets.0.norm2": "blocks.11.norm2", + "encoder.mid_block.resnets.0.conv2": "blocks.11.conv2", + "encoder.mid_block.resnets.1.norm1": "blocks.13.norm1", + "encoder.mid_block.resnets.1.conv1": "blocks.13.conv1", + "encoder.mid_block.resnets.1.norm2": "blocks.13.norm2", + "encoder.mid_block.resnets.1.conv2": "blocks.13.conv2", + "encoder.conv_norm_out": "conv_norm_out", + "encoder.conv_out": "conv_out", + } + name_list = sorted([name for name in state_dict]) + rename_dict = {} + block_id = {"ResnetBlock": -1, "DownSampler": -1, "UpSampler": -1} + last_block_type_with_id = {"ResnetBlock": "", "DownSampler": "", "UpSampler": ""} + for name in name_list: + names = name.split(".") + name_prefix = ".".join(names[:-1]) + if name_prefix in local_rename_dict: + rename_dict[name] = local_rename_dict[name_prefix] + "." + names[-1] + elif name.startswith("encoder.down_blocks"): + block_type = {"resnets": "ResnetBlock", "downsamplers": "DownSampler", "upsamplers": "UpSampler"}[names[3]] + block_type_with_id = ".".join(names[:5]) + if block_type_with_id != last_block_type_with_id[block_type]: + block_id[block_type] += 1 + last_block_type_with_id[block_type] = block_type_with_id + while block_id[block_type] < len(block_types) and block_types[block_id[block_type]] != block_type: + block_id[block_type] += 1 + block_type_with_id = ".".join(names[:5]) + names = ["blocks", str(block_id[block_type])] + names[5:] + rename_dict[name] = ".".join(names) + + # Convert state_dict + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + state_dict_[rename_dict[name]] = state_dict[name] + return state_dict_ + + +def FluxVAEDecoderStateDictConverterDiffusers(state_dict): + # architecture + block_types = [ + 'ResnetBlock', 'VAEAttentionBlock', 'ResnetBlock', + 'ResnetBlock', 'ResnetBlock', 'ResnetBlock', 'UpSampler', + 'ResnetBlock', 'ResnetBlock', 'ResnetBlock', 'UpSampler', + 'ResnetBlock', 'ResnetBlock', 'ResnetBlock', 'UpSampler', + 'ResnetBlock', 'ResnetBlock', 'ResnetBlock' + ] + + # Rename each parameter + local_rename_dict = { + "post_quant_conv": "post_quant_conv", + "decoder.conv_in": "conv_in", + "decoder.mid_block.attentions.0.group_norm": "blocks.1.norm", + "decoder.mid_block.attentions.0.to_q": "blocks.1.transformer_blocks.0.to_q", + "decoder.mid_block.attentions.0.to_k": "blocks.1.transformer_blocks.0.to_k", + "decoder.mid_block.attentions.0.to_v": "blocks.1.transformer_blocks.0.to_v", + "decoder.mid_block.attentions.0.to_out.0": "blocks.1.transformer_blocks.0.to_out", + "decoder.mid_block.resnets.0.norm1": "blocks.0.norm1", + "decoder.mid_block.resnets.0.conv1": "blocks.0.conv1", + "decoder.mid_block.resnets.0.norm2": "blocks.0.norm2", + "decoder.mid_block.resnets.0.conv2": "blocks.0.conv2", + "decoder.mid_block.resnets.1.norm1": "blocks.2.norm1", + "decoder.mid_block.resnets.1.conv1": "blocks.2.conv1", + "decoder.mid_block.resnets.1.norm2": "blocks.2.norm2", + "decoder.mid_block.resnets.1.conv2": "blocks.2.conv2", + "decoder.conv_norm_out": "conv_norm_out", + "decoder.conv_out": "conv_out", + } + name_list = sorted([name for name in state_dict]) + rename_dict = {} + block_id = {"ResnetBlock": 2, "DownSampler": 2, "UpSampler": 2} + last_block_type_with_id = {"ResnetBlock": "", "DownSampler": "", "UpSampler": ""} + for name in name_list: + names = name.split(".") + name_prefix = ".".join(names[:-1]) + if name_prefix in local_rename_dict: + rename_dict[name] = local_rename_dict[name_prefix] + "." + names[-1] + elif name.startswith("decoder.up_blocks"): + block_type = {"resnets": "ResnetBlock", "downsamplers": "DownSampler", "upsamplers": "UpSampler"}[names[3]] + block_type_with_id = ".".join(names[:5]) + if block_type_with_id != last_block_type_with_id[block_type]: + block_id[block_type] += 1 + last_block_type_with_id[block_type] = block_type_with_id + while block_id[block_type] < len(block_types) and block_types[block_id[block_type]] != block_type: + block_id[block_type] += 1 + block_type_with_id = ".".join(names[:5]) + names = ["blocks", str(block_id[block_type])] + names[5:] + rename_dict[name] = ".".join(names) + + # Convert state_dict + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + state_dict_[rename_dict[name]] = state_dict[name] + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/nexus_gen.py b/diffsynth/utils/state_dict_converters/nexus_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..aff853d0e76dd1f130ce44241462f61af84370db --- /dev/null +++ b/diffsynth/utils/state_dict_converters/nexus_gen.py @@ -0,0 +1,6 @@ +def NexusGenAutoregressiveModelStateDictConverter(state_dict): + new_state_dict = {} + for key in state_dict: + value = state_dict[key] + new_state_dict["model." + key] = value + return new_state_dict \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/nexus_gen_projector.py b/diffsynth/utils/state_dict_converters/nexus_gen_projector.py new file mode 100644 index 0000000000000000000000000000000000000000..c9a44665551ba4a97d063de94f6025c8c48989fd --- /dev/null +++ b/diffsynth/utils/state_dict_converters/nexus_gen_projector.py @@ -0,0 +1,15 @@ +def NexusGenMergerStateDictConverter(state_dict): + merger_state_dict = {} + for key in state_dict: + if key.startswith('embedding_merger.'): + value = state_dict[key] + new_key = key.replace("embedding_merger.", "") + merger_state_dict[new_key] = value + return merger_state_dict + +def NexusGenAdapterStateDictConverter(state_dict): + adapter_state_dict = {} + for key in state_dict: + if key.startswith('adapter.'): + adapter_state_dict[key] = state_dict[key] + return adapter_state_dict \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/qwen_image_text_encoder.py b/diffsynth/utils/state_dict_converters/qwen_image_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..e8192a1f2a959685cf1fa5af40824bd896454141 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/qwen_image_text_encoder.py @@ -0,0 +1,10 @@ +def QwenImageTextEncoderStateDictConverter(state_dict): + state_dict_ = {} + for k in state_dict: + v = state_dict[k] + if k.startswith("visual."): + k = "model." + k + elif k.startswith("model."): + k = k.replace("model.", "model.language_model.") + state_dict_[k] = v + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/step1x_connector.py b/diffsynth/utils/state_dict_converters/step1x_connector.py new file mode 100644 index 0000000000000000000000000000000000000000..35a2a4167b16ea5cc16aaa1b0f20575bc2918bbf --- /dev/null +++ b/diffsynth/utils/state_dict_converters/step1x_connector.py @@ -0,0 +1,7 @@ +def Qwen2ConnectorStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("connector."): + name_ = name[len("connector."):] + state_dict_[name_] = state_dict[name] + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/wan_video_animate_adapter.py b/diffsynth/utils/state_dict_converters/wan_video_animate_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..8ea69f4e6696bbef6de197abaa031ea8cc5b398e --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_animate_adapter.py @@ -0,0 +1,6 @@ +def WanAnimateAdapterStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("pose_patch_embedding.") or name.startswith("face_adapter") or name.startswith("face_encoder") or name.startswith("motion_encoder"): + state_dict_[name] = state_dict[name] + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/wan_video_dit.py b/diffsynth/utils/state_dict_converters/wan_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..c7716dad52e42ebf76f98dd85511ac0a04b3d3b3 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_dit.py @@ -0,0 +1,83 @@ +def WanVideoDiTFromDiffusers(state_dict): + rename_dict = { + "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight", + "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight", + "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias", + "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight", + "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias", + "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight", + "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias", + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias", + "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight", + "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight", + "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight", + "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias", + "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight", + "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias", + "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight", + "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias", + "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight", + "blocks.0.attn2.add_k_proj.bias":"blocks.0.cross_attn.k_img.bias", + "blocks.0.attn2.add_k_proj.weight":"blocks.0.cross_attn.k_img.weight", + "blocks.0.attn2.add_v_proj.bias":"blocks.0.cross_attn.v_img.bias", + "blocks.0.attn2.add_v_proj.weight":"blocks.0.cross_attn.v_img.weight", + "blocks.0.attn2.norm_added_k.weight":"blocks.0.cross_attn.norm_k_img.weight", + "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.bias": "blocks.0.norm3.bias", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "blocks.0.scale_shift_table": "blocks.0.modulation", + "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias", + "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight", + "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias", + "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight", + "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias", + "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight", + "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias", + "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight", + "condition_embedder.time_proj.bias": "time_projection.1.bias", + "condition_embedder.time_proj.weight": "time_projection.1.weight", + "condition_embedder.image_embedder.ff.net.0.proj.bias":"img_emb.proj.1.bias", + "condition_embedder.image_embedder.ff.net.0.proj.weight":"img_emb.proj.1.weight", + "condition_embedder.image_embedder.ff.net.2.bias":"img_emb.proj.3.bias", + "condition_embedder.image_embedder.ff.net.2.weight":"img_emb.proj.3.weight", + "condition_embedder.image_embedder.norm1.bias":"img_emb.proj.0.bias", + "condition_embedder.image_embedder.norm1.weight":"img_emb.proj.0.weight", + "condition_embedder.image_embedder.norm2.bias":"img_emb.proj.4.bias", + "condition_embedder.image_embedder.norm2.weight":"img_emb.proj.4.weight", + "patch_embedding.bias": "patch_embedding.bias", + "patch_embedding.weight": "patch_embedding.weight", + "scale_shift_table": "head.modulation", + "proj_out.bias": "head.head.bias", + "proj_out.weight": "head.head.weight", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + state_dict_[rename_dict[name]] = state_dict[name] + else: + name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) + if name_ in rename_dict: + name_ = rename_dict[name_] + name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:]) + state_dict_[name_] = state_dict[name] + return state_dict_ + + +def WanVideoDiTStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("vace"): + continue + if name.split(".")[0] in ["pose_patch_embedding", "face_adapter", "face_encoder", "motion_encoder"]: + continue + name_ = name + if name_.startswith("model."): + name_ = name_[len("model."):] + state_dict_[name_] = state_dict[name] + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/wan_video_image_encoder.py b/diffsynth/utils/state_dict_converters/wan_video_image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ecb7e9bfce50e88601f8876341ac56645a8e5913 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_image_encoder.py @@ -0,0 +1,8 @@ +def WanImageEncoderStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("textual."): + continue + name_ = "model." + name + state_dict_[name_] = state_dict[name] + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/wan_video_mot.py b/diffsynth/utils/state_dict_converters/wan_video_mot.py new file mode 100644 index 0000000000000000000000000000000000000000..12b42d7db752fca1cb24c0f16217deab925916f5 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_mot.py @@ -0,0 +1,78 @@ +def WanVideoMotStateDictConverter(state_dict): + rename_dict = { + "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight", + "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight", + "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias", + "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight", + "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias", + "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight", + "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias", + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias", + "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight", + "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight", + "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight", + "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias", + "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight", + "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias", + "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight", + "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias", + "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight", + "blocks.0.attn2.add_k_proj.bias":"blocks.0.cross_attn.k_img.bias", + "blocks.0.attn2.add_k_proj.weight":"blocks.0.cross_attn.k_img.weight", + "blocks.0.attn2.add_v_proj.bias":"blocks.0.cross_attn.v_img.bias", + "blocks.0.attn2.add_v_proj.weight":"blocks.0.cross_attn.v_img.weight", + "blocks.0.attn2.norm_added_k.weight":"blocks.0.cross_attn.norm_k_img.weight", + "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.bias": "blocks.0.norm3.bias", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "blocks.0.scale_shift_table": "blocks.0.modulation", + "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias", + "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight", + "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias", + "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight", + "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias", + "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight", + "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias", + "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight", + "condition_embedder.time_proj.bias": "time_projection.1.bias", + "condition_embedder.time_proj.weight": "time_projection.1.weight", + "condition_embedder.image_embedder.ff.net.0.proj.bias":"img_emb.proj.1.bias", + "condition_embedder.image_embedder.ff.net.0.proj.weight":"img_emb.proj.1.weight", + "condition_embedder.image_embedder.ff.net.2.bias":"img_emb.proj.3.bias", + "condition_embedder.image_embedder.ff.net.2.weight":"img_emb.proj.3.weight", + "condition_embedder.image_embedder.norm1.bias":"img_emb.proj.0.bias", + "condition_embedder.image_embedder.norm1.weight":"img_emb.proj.0.weight", + "condition_embedder.image_embedder.norm2.bias":"img_emb.proj.4.bias", + "condition_embedder.image_embedder.norm2.weight":"img_emb.proj.4.weight", + "patch_embedding.bias": "patch_embedding.bias", + "patch_embedding.weight": "patch_embedding.weight", + "scale_shift_table": "head.modulation", + "proj_out.bias": "head.head.bias", + "proj_out.weight": "head.head.weight", + } + mot_layers = (0, 4, 8, 12, 16, 20, 24, 28, 32, 36) + mot_layers_mapping = {i:n for n, i in enumerate(mot_layers)} + state_dict_ = {} + for name in state_dict: + if "_mot_ref" not in name: + continue + param = state_dict[name] + name = name.replace("_mot_ref", "") + if name in rename_dict: + state_dict_[rename_dict[name]] = param + else: + if name.split(".")[1].isdigit(): + block_id = int(name.split(".")[1]) + name = name.replace(str(block_id), str(mot_layers_mapping[block_id])) + name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) + if name_ in rename_dict: + name_ = rename_dict[name_] + name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:]) + state_dict_[name_] = param + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/wan_video_vace.py b/diffsynth/utils/state_dict_converters/wan_video_vace.py new file mode 100644 index 0000000000000000000000000000000000000000..cdfef6998f47ac7d3640b28b99f109e7f04baeba --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_vace.py @@ -0,0 +1,3 @@ +def VaceWanModelDictConverter(state_dict): + state_dict_ = {name: state_dict[name] for name in state_dict if name.startswith("vace")} + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/wan_video_vae.py b/diffsynth/utils/state_dict_converters/wan_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..76a430e1bd4575e0ae06234de23b620d4877566f --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wan_video_vae.py @@ -0,0 +1,7 @@ +def WanVideoVAEStateDictConverter(state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ \ No newline at end of file diff --git a/diffsynth/utils/state_dict_converters/wans2v_audio_encoder.py b/diffsynth/utils/state_dict_converters/wans2v_audio_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..eaa12c0d4ff7fc166eea1f804cb645be9aa28776 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/wans2v_audio_encoder.py @@ -0,0 +1,12 @@ +def WanS2VAudioEncoderStateDictConverter(state_dict): + rename_dict = { + "model.wav2vec2.encoder.pos_conv_embed.conv.weight_g": "model.wav2vec2.encoder.pos_conv_embed.conv.parametrizations.weight.original0", + "model.wav2vec2.encoder.pos_conv_embed.conv.weight_v": "model.wav2vec2.encoder.pos_conv_embed.conv.parametrizations.weight.original1", + } + state_dict_ = {} + for name in state_dict: + name_ = "model." + name + if name_ in rename_dict: + name_ = rename_dict[name_] + state_dict_[name_] = state_dict[name] + return state_dict_ diff --git a/diffsynth/utils/xfuser/__init__.py b/diffsynth/utils/xfuser/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..13dd178e2d47bf58de1bfca6d21052d0563c70ca --- /dev/null +++ b/diffsynth/utils/xfuser/__init__.py @@ -0,0 +1 @@ +from .xdit_context_parallel import usp_attn_forward, usp_dit_forward, get_sequence_parallel_world_size, initialize_usp diff --git a/diffsynth/utils/xfuser/xdit_context_parallel.py b/diffsynth/utils/xfuser/xdit_context_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..11733132f95dea86e6fbe34900de40c40b92ba60 --- /dev/null +++ b/diffsynth/utils/xfuser/xdit_context_parallel.py @@ -0,0 +1,145 @@ +import torch +from typing import Optional +from einops import rearrange +from xfuser.core.distributed import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group) +from xfuser.core.long_ctx_attention import xFuserLongContextAttention + + +def initialize_usp(): + import torch.distributed as dist + from xfuser.core.distributed import initialize_model_parallel, init_distributed_environment + dist.init_process_group(backend="nccl", init_method="env://") + init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size()) + initialize_model_parallel( + sequence_parallel_degree=dist.get_world_size(), + ring_degree=1, + ulysses_degree=dist.get_world_size(), + ) + torch.cuda.set_device(dist.get_rank()) + + +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer(position.type(torch.float64), torch.pow( + 10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + +def pad_freqs(original_tensor, target_len): + seq_len, s1, s2 = original_tensor.shape + pad_size = target_len - seq_len + padding_tensor = torch.ones( + pad_size, + s1, + s2, + dtype=original_tensor.dtype, + device=original_tensor.device) + padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0) + return padded_tensor + +def rope_apply(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + s_per_rank = x.shape[1] + + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + + sp_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + freqs = pad_freqs(freqs, s_per_rank * sp_size) + freqs_rank = freqs[(sp_rank * s_per_rank):((sp_rank + 1) * s_per_rank), :, :] + + x_out = torch.view_as_real(x_out * freqs_rank).flatten(2) + return x_out.to(x.dtype) + +def usp_dit_forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + + x, (f, h, w) = self.patchify(x) + + freqs = torch.cat([ + self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + # Context Parallel + chunks = torch.chunk(x, get_sequence_parallel_world_size(), dim=1) + pad_shape = chunks[0].shape[1] - chunks[-1].shape[1] + chunks = [torch.nn.functional.pad(chunk, (0, 0, 0, chunks[0].shape[1]-chunk.shape[1]), value=0) for chunk in chunks] + x = chunks[get_sequence_parallel_rank()] + + for block in self.blocks: + if self.training and use_gradient_checkpointing: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + x = self.head(x, t) + + # Context Parallel + x = get_sp_group().all_gather(x, dim=1) + x = x[:, :-pad_shape] if pad_shape > 0 else x + + # unpatchify + x = self.unpatchify(x, (f, h, w)) + return x + + +def usp_attn_forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads) + + x = xFuserLongContextAttention()( + None, + query=q, + key=k, + value=v, + ) + x = x.flatten(2) + + del q, k, v + torch.cuda.empty_cache() + return self.o(x) \ No newline at end of file diff --git a/examples/1.png b/examples/1.png new file mode 100644 index 0000000000000000000000000000000000000000..d27424667f511c33c575d6a3930e2c4fca6db4c3 --- /dev/null +++ b/examples/1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6885e3458c9a5aa5d43afee50df957762b4a7ee5627e02bf4fad64aa1acff3f0 +size 534109 diff --git a/examples/12.png b/examples/12.png new file mode 100644 index 0000000000000000000000000000000000000000..5a3dea25547d097f9ee37aab949e51a56ac68ca3 --- /dev/null +++ b/examples/12.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3ca7297c6c1c8ef45c3d49c1390ba0acad30a314c7f31f8cd62e02fd69ec6824 +size 822464 diff --git a/examples/5.png b/examples/5.png new file mode 100644 index 0000000000000000000000000000000000000000..922d934fa4912629025626ca8927942937409c9b --- /dev/null +++ b/examples/5.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3cdfc7f7592aa7e0ed4b9d73d9eac792cd12160f574e0cb75f243dc55d47576e +size 460457 diff --git a/examples/9.png b/examples/9.png new file mode 100644 index 0000000000000000000000000000000000000000..bb5f4106e6a78a4f9408a2d99464d346f5dd72f0 --- /dev/null +++ b/examples/9.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:277b33dd7ad4709b9a720d53879287cf3162261d6c3e3e01d995e9106036389a +size 730234 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..da17863e58df7ca2f1f43b1bdc530c4cd2c95c2c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,20 @@ +torchvision +transformers>=4.51.0 +diffusers>=0.32.0 +matplotlib +accelerate +peft +safetensors +einops +sentencepiece +protobuf +ftfy +imageio +imageio[ffmpeg] +pandas +modelscope +omegaconf +addict +opencv-python-headless +evo +numpy<2 diff --git a/src/MetaView_dit.py b/src/MetaView_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..319f4658364c4176c126c94e5dae43b28533442f --- /dev/null +++ b/src/MetaView_dit.py @@ -0,0 +1,797 @@ +# Adapted from https://github.com/modelscope/DiffSynth-Studio + +import torch, math +import torch.nn as nn +from typing import Tuple, Optional, Union, List +from einops import rearrange +import matplotlib.pyplot as plt +from diffsynth.models.general_modules import TimestepEmbeddings, RMSNorm, AdaLayerNorm + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +from src.PRoPE import PropeDotProductAttention +from src.lora import LoRALinearLayer + +def visualize_attention_heads(attn_probs, name=None, layer=None): + seq_len, _ = attn_probs.shape + + plt.figure(figsize=(4, 3)) + + attn_map = attn_probs.clone().to(torch.float32).detach().cpu().numpy() + + im = plt.imshow(attn_map, cmap='viridis', aspect='auto', vmin=0, vmax=0.005) + plt.title(f'Attention Maps {name} layer {layer}') + plt.xlabel('Key Position') + plt.ylabel('Query Position') + + cbar = plt.colorbar(im, fraction=0.046, pad=0.04) + #cbar.set_ticks([0, 0.001, 0.002, 0.003, 0.004, 0.005]) + #cbar.set_ticklabels(['0.0', '0.001', '0.002', '0.003', '0.004', '0.005']) + # plt.tight_layout() + + save_path = f"./attn_maps/attn_map_merge/{name}_layer_{layer:02d}" + if save_path: + plt.savefig(save_path, dpi=100, bbox_inches='tight') + + + +def qwen_image_flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, attention_mask = None, enable_fp8_attention: bool = False, seq_order=None, layer=None): + if FLASH_ATTN_3_AVAILABLE and attention_mask is None: + print("flash attn!!!") + if not enable_fp8_attention: + q = rearrange(q, "b n s d -> b s n d", n=num_heads) + k = rearrange(k, "b n s d -> b s n d", n=num_heads) + v = rearrange(v, "b n s d -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x, tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + else: + origin_dtype = q.dtype + q_std, k_std, v_std = q.std(), k.std(), v.std() + q, k, v = (q / q_std).to(torch.float8_e4m3fn), (k / k_std).to(torch.float8_e4m3fn), (v / v_std).to(torch.float8_e4m3fn) + q = rearrange(q, "b n s d -> b s n d", n=num_heads) + k = rearrange(k, "b n s d -> b s n d", n=num_heads) + v = rearrange(v, "b n s d -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=q_std * k_std / math.sqrt(q.size(-1))) + if isinstance(x, tuple): + x = x[0] + x = x.to(origin_dtype) * v_std + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + else: + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + + if 0 and seq_order is not None and layer % 10 ==0: + img_seq = seq_order[0] // 2 + _3D_seq = seq_order[1] + # print(seq_order) + + batch_size, _, seq_len, head_dim = q.shape + + attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + + attn_probs = torch.softmax(attn_scores, dim=-1) # b n s s + # print(attn_probs.shape) + # lantent_attn_map = attn_probs[0, 0, :img_seq, :] + # src_attn_map = attn_probs[0, 0, img_seq : 2*img_seq, :] + # _3D_attn_map = attn_probs[0, 0, 2*img_seq: , :] # n s s + + lantent_attn_map = attn_probs[0, :, :img_seq, :] + src_attn_map = attn_probs[0, :, img_seq : 2*img_seq, :] + _3D_attn_map = attn_probs[0, :, 2*img_seq: , :] # n s s + + lantent_attn_map = torch.mean(lantent_attn_map, dim=0) # s s + src_attn_map = torch.mean(src_attn_map, dim=0) + _3D_attn_map = torch.mean(_3D_attn_map, dim=0) + + visualize_attention_heads(lantent_attn_map, name="latent", layer=layer) + visualize_attention_heads(src_attn_map, name="src", layer=layer) + visualize_attention_heads(_3D_attn_map, name="3D", layer=layer) + + return x + + +class ApproximateGELU(nn.Module): + def __init__(self, dim_in: int, dim_out: int, bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + return x * torch.sigmoid(1.702 * x) + +def apply_rotary_emb_qwen( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] +): + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + return x_out.type_as(x) + + +class QwenEmbedRope(nn.Module): + def __init__(self, theta: int, axes_dim: list[int], scale_rope=False): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + pos_index = torch.arange(4096) + neg_index = torch.arange(4096).flip(0) * -1 - 1 + self.pos_freqs = torch.cat([ + self.rope_params(pos_index, self.axes_dim[0], self.theta), + self.rope_params(pos_index, self.axes_dim[1], self.theta), + self.rope_params(pos_index, self.axes_dim[2], self.theta), + ], dim=1) + self.neg_freqs = torch.cat([ + self.rope_params(neg_index, self.axes_dim[0], self.theta), + self.rope_params(neg_index, self.axes_dim[1], self.theta), + self.rope_params(neg_index, self.axes_dim[2], self.theta), + ], dim=1) + self.rope_cache = {} + self.scale_rope = scale_rope + + def rope_params(self, index, dim, theta=10000): + """ + Args: + index: [0, 1, 2, 3] 1D Tensor representing the position index of the token + """ + assert dim % 2 == 0 + freqs = torch.outer( + index, + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)) + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + + def _expand_pos_freqs_if_needed(self, video_fhw, txt_seq_lens): + if isinstance(video_fhw, list): + video_fhw = tuple(max([i[j] for i in video_fhw]) for j in range(3)) + _, height, width = video_fhw + if self.scale_rope: + max_vid_index = max(height // 2, width // 2) + else: + max_vid_index = max(height, width) + required_len = max_vid_index + max(txt_seq_lens) + cur_max_len = self.pos_freqs.shape[0] + if required_len <= cur_max_len: + return + + new_max_len = math.ceil(required_len / 512) * 512 + pos_index = torch.arange(new_max_len) + neg_index = torch.arange(new_max_len).flip(0) * -1 - 1 + self.pos_freqs = torch.cat([ + self.rope_params(pos_index, self.axes_dim[0], self.theta), + self.rope_params(pos_index, self.axes_dim[1], self.theta), + self.rope_params(pos_index, self.axes_dim[2], self.theta), + ], dim=1) + self.neg_freqs = torch.cat([ + self.rope_params(neg_index, self.axes_dim[0], self.theta), + self.rope_params(neg_index, self.axes_dim[1], self.theta), + self.rope_params(neg_index, self.axes_dim[2], self.theta), + ], dim=1) + return + + + def forward(self, video_fhw, txt_seq_lens, device): + self._expand_pos_freqs_if_needed(video_fhw, txt_seq_lens) + if self.pos_freqs.device != device: + self.pos_freqs = self.pos_freqs.to(device) + self.neg_freqs = self.neg_freqs.to(device) + + vid_freqs = [] + max_vid_index = 0 + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + rope_key = f"{idx}_{height}_{width}" + + if rope_key not in self.rope_cache: + seq_lens = frame * height * width + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + if self.scale_rope: + freqs_height = torch.cat( + [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0 + ) + freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0) + freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1) + + else: + freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1) + + freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1) + self.rope_cache[rope_key] = freqs.clone().contiguous() + vid_freqs.append(self.rope_cache[rope_key]) + + if self.scale_rope: + max_vid_index = max(height // 2, width // 2, max_vid_index) + else: + max_vid_index = max(height, width, max_vid_index) + + max_len = max(txt_seq_lens) + txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + vid_freqs = torch.cat(vid_freqs, dim=0) + + return vid_freqs, txt_freqs + + + def forward_sampling(self, video_fhw, txt_seq_lens, device): + self._expand_pos_freqs_if_needed(video_fhw, txt_seq_lens) + if self.pos_freqs.device != device: + self.pos_freqs = self.pos_freqs.to(device) + self.neg_freqs = self.neg_freqs.to(device) + + vid_freqs = [] + max_vid_index = 0 + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + rope_key = f"{idx}_{height}_{width}" + if idx > 0 and f"{0}_{height}_{width}" not in self.rope_cache: + frame_0, height_0, width_0 = video_fhw[0] + + rope_key_0 = f"0_{height_0}_{width_0}" + spatial_freqs_0 = self.rope_cache[rope_key_0].reshape(frame_0, height_0, width_0, -1) + h_indices = torch.linspace(0, height_0 - 1, height).long() + w_indices = torch.linspace(0, width_0 - 1, width).long() + h_grid, w_grid = torch.meshgrid(h_indices, w_indices, indexing='ij') + sampled_rope = spatial_freqs_0[:, h_grid, w_grid, :] + + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + sampled_rope[:, :, :, :freqs_frame.shape[-1]] = freqs_frame + + seq_lens = frame * height * width + self.rope_cache[rope_key] = sampled_rope.reshape(seq_lens, -1).clone() + if rope_key not in self.rope_cache: + seq_lens = frame * height * width + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) + if self.scale_rope: + freqs_height = torch.cat( + [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0 + ) + freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0) + freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1) + + else: + freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1) + freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1) + + freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1) + self.rope_cache[rope_key] = freqs.clone() + vid_freqs.append(self.rope_cache[rope_key].contiguous()) + + if self.scale_rope: + max_vid_index = max(height // 2, width // 2, max_vid_index) + else: + max_vid_index = max(height, width, max_vid_index) + + max_len = max(txt_seq_lens) + txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + vid_freqs = torch.cat(vid_freqs, dim=0) + + return vid_freqs, txt_freqs + + +class QwenFeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + dropout: float = 0.0, + ): + super().__init__() + inner_dim = int(dim * 4) + self.net = nn.ModuleList([]) + self.net.append(ApproximateGELU(dim, inner_dim)) + self.net.append(nn.Dropout(dropout)) + self.net.append(nn.Linear(inner_dim, dim_out)) + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +class MetaViewSelfAttention3D(nn.Module): + def __init__( + self, + dim_a, # 3072 + dim_b, + num_heads, + head_dim, + merge_3D = False, + ): + super().__init__() + self.merge_3D = merge_3D + + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = nn.Linear(dim_a, dim_a) + self.to_k = nn.Linear(dim_a, dim_a) + self.to_v = nn.Linear(dim_a, dim_a) + self.to_q_3D = nn.Linear(dim_b, dim_a) + self.to_k_3D = nn.Linear(dim_b, dim_a) + self.to_v_3D = nn.Linear(dim_b, dim_a) + + self.norm_q = RMSNorm(head_dim, eps=1e-6) + self.norm_k = RMSNorm(head_dim, eps=1e-6) + + self.norm_added_q = RMSNorm(head_dim, eps=1e-6) + self.norm_added_k = RMSNorm(head_dim, eps=1e-6) + + # zero init linear out + self.to_out = torch.nn.Sequential(nn.Linear(dim_a, dim_a)) + nn.init.zeros_(self.to_out[0].weight) + nn.init.zeros_(self.to_out[0].bias) + + def forward( + self, + image: torch.FloatTensor, + feat_3D: torch.FloatTensor, # (1, 20, 36, 1536) + attention_mask: Optional[torch.FloatTensor] = None, + enable_fp8_attention: bool = False, + prope: Optional[PropeDotProductAttention] = None, + add_prope: Optional[PropeDotProductAttention] = None, + block_id=None, + ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + # feat_3D = rearrange(feat_3D, 'b h w d -> b (h w) d') + + img_q, img_k, img_v = self.to_q(image), self.to_k(image), self.to_v(image) + _3D_q, _3D_k, _3D_v = self.to_q_3D(feat_3D), self.to_k_3D(feat_3D), self.to_v_3D(feat_3D) + seq_img = img_q.shape[1] + seq_3D = _3D_q.shape[1] + + img_q = rearrange(img_q, 'b s (h d) -> b h s d', h=self.num_heads) + img_k = rearrange(img_k, 'b s (h d) -> b h s d', h=self.num_heads) + img_v = rearrange(img_v, 'b s (h d) -> b h s d', h=self.num_heads) + img_q, img_k = self.norm_q(img_q), self.norm_k(img_k) + + _3D_q = rearrange(_3D_q, 'b s (h d) -> b h s d', h=self.num_heads) + _3D_k = rearrange(_3D_k, 'b s (h d) -> b h s d', h=self.num_heads) + _3D_v = rearrange(_3D_v, 'b s (h d) -> b h s d', h=self.num_heads) + _3D_q, _3D_k = self.norm_added_q(_3D_q), self.norm_added_k(_3D_k) + + if prope is not None: # standard prope in q k v out + # print("Inject Prope in q k v out in added attn layer! ") + img_q = prope._apply_to_q(img_q) + img_k = prope._apply_to_kv(img_k) + img_v = prope._apply_to_kv(img_v) + + _3D_q = add_prope._apply_to_q(_3D_q) + _3D_k = add_prope._apply_to_kv(_3D_k) + _3D_v = add_prope._apply_to_kv(_3D_v) + + joint_q = torch.cat([img_q, _3D_q], dim=2) + joint_k = torch.cat([img_k, _3D_k], dim=2) + joint_v = torch.cat([img_v, _3D_v], dim=2) + + num_heads = img_q.shape[1] + joint_attn_out = qwen_image_flash_attention(joint_q, joint_k, joint_v, num_heads=num_heads, layer=block_id, seq_order=[seq_img, seq_3D], attention_mask=attention_mask, enable_fp8_attention=enable_fp8_attention).to(img_q.dtype) + # [b, s, (n d)] + + img_attn_output = joint_attn_out[:, :seq_img, :] # discard 3D tokens + _3D_attn_output = joint_attn_out[:, seq_img:, :] + + # reshape back and apply prope + if prope is not None: + img_attn_output = rearrange(img_attn_output, "b s (n d) -> b n s d", n=num_heads) + img_attn_output = prope._apply_to_o(img_attn_output) + img_attn_output = rearrange(img_attn_output, "b n s d -> b s (n d) ", n=num_heads) + + img_attn_output = self.to_out(img_attn_output) + + if add_prope is not None and self.merge_3D: + _3D_attn_output = rearrange(_3D_attn_output, "b s (n d) -> b n s d", n=num_heads) + _3D_attn_output = add_prope._apply_to_o(_3D_attn_output) + _3D_attn_output = rearrange(_3D_attn_output, "b n s d -> b s (n d) ", n=num_heads) + + if self.merge_3D: + _3D_attn_output = self.to_out(_3D_attn_output) + + if self.merge_3D: + return img_attn_output, _3D_attn_output + else: + return img_attn_output, None + + +class QwenDoubleStreamAttention(nn.Module): + def __init__( + self, + dim_a, + dim_b, + num_heads, + head_dim, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.to_q = nn.Linear(dim_a, dim_a) + self.to_k = nn.Linear(dim_a, dim_a) + self.to_v = nn.Linear(dim_a, dim_a) + self.norm_q = RMSNorm(head_dim, eps=1e-6) + self.norm_k = RMSNorm(head_dim, eps=1e-6) + + self.add_q_proj = nn.Linear(dim_b, dim_b) + self.add_k_proj = nn.Linear(dim_b, dim_b) + self.add_v_proj = nn.Linear(dim_b, dim_b) + self.norm_added_q = RMSNorm(head_dim, eps=1e-6) + self.norm_added_k = RMSNorm(head_dim, eps=1e-6) + + self.to_out = torch.nn.Sequential(nn.Linear(dim_a, dim_a)) + self.to_add_out = nn.Linear(dim_b, dim_b) + + def forward( + self, + image: torch.FloatTensor, + text: torch.FloatTensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + enable_fp8_attention: bool = False, + prope: Optional[PropeDotProductAttention] = None, + ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + img_q, img_k, img_v = self.to_q(image), self.to_k(image), self.to_v(image) + txt_q, txt_k, txt_v = self.add_q_proj(text), self.add_k_proj(text), self.add_v_proj(text) + seq_txt = txt_q.shape[1] + + img_q = rearrange(img_q, 'b s (h d) -> b h s d', h=self.num_heads) + img_k = rearrange(img_k, 'b s (h d) -> b h s d', h=self.num_heads) + img_v = rearrange(img_v, 'b s (h d) -> b h s d', h=self.num_heads) + + txt_q = rearrange(txt_q, 'b s (h d) -> b h s d', h=self.num_heads) + txt_k = rearrange(txt_k, 'b s (h d) -> b h s d', h=self.num_heads) + txt_v = rearrange(txt_v, 'b s (h d) -> b h s d', h=self.num_heads) + + img_q, img_k = self.norm_q(img_q), self.norm_k(img_k) + txt_q, txt_k = self.norm_added_q(txt_q), self.norm_added_k(txt_k) + + if prope is not None and image_rotary_emb is not None: # prope with original Qwen backbone + # print("!! modified PRoPE in original Qwen backbone") + _, txt_freqs = image_rotary_emb + txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs) + txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs) + img_q = prope._apply_to_q(img_q) + img_k = prope._apply_to_kv(img_k) + # img_v = prope._apply_to_kv(img_v) + elif image_rotary_emb is not None: + # print("!!! No PRoPE in original Qwen backbone !!") + img_freqs, txt_freqs = image_rotary_emb + img_q = apply_rotary_emb_qwen(img_q, img_freqs) + img_k = apply_rotary_emb_qwen(img_k, img_freqs) + txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs) + txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs) + + joint_q = torch.cat([txt_q, img_q], dim=2) + joint_k = torch.cat([txt_k, img_k], dim=2) + joint_v = torch.cat([txt_v, img_v], dim=2) + + num_heads = joint_q.shape[1] + joint_attn_out = qwen_image_flash_attention(joint_q, joint_k, joint_v, num_heads=joint_q.shape[1], attention_mask=attention_mask, enable_fp8_attention=enable_fp8_attention).to(joint_q.dtype) + # [b, s, (n d)] + + txt_attn_output = joint_attn_out[:, :seq_txt, :] + img_attn_output = joint_attn_out[:, seq_txt:, :] + + # reshape back and apply prope + # if prope is not None: + # img_attn_output = rearrange(img_attn_output, "b s (n d) -> b n s d", n=num_heads) + # img_attn_output = prope._apply_to_o(img_attn_output) + # img_attn_output = rearrange(img_attn_output, "b n s d -> b s (n d) ", n=num_heads) + + img_attn_output = self.to_out(img_attn_output) + txt_attn_output = self.to_add_out(txt_attn_output) + + return img_attn_output, txt_attn_output + + +class MetaViewTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + add_attn = False, + add_attn_type = None, + add_in_dim = None, + lora_rank = None, + merge_3D = False, + ): + super().__init__() + + self.merge_3D = merge_3D + + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + self.img_mod = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, 6 * dim), + ) + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = QwenDoubleStreamAttention( + dim_a=dim, + dim_b=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + ) + if add_attn: + self.prope_attn = MetaViewSelfAttention3D( + dim_a=dim, + dim_b=add_in_dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + merge_3D=merge_3D, + ) + + + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenFeedForward(dim=dim, dim_out=dim) + + self.txt_mod = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, 6 * dim, bias=True), + ) + self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_mlp = QwenFeedForward(dim=dim, dim_out=dim) + + def _modulate(self, x, mod_params): + shift, scale, gate = mod_params.chunk(3, dim=-1) + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1) + + def forward( + self, + image: torch.Tensor, + text: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + enable_fp8_attention = False, + prope: Optional[PropeDotProductAttention] = None, + add_prope: Optional[PropeDotProductAttention] = None, + add_attn = False, + feat_3D = None, # (1, 20, 36, 1536) + block_id = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # print("feat_3D:", feat_3D.shape) + # print("image: ", image.shape) + img_mod_attn, img_mod_mlp = self.img_mod(temb).chunk(2, dim=-1) # [B, 3*dim] each + txt_mod_attn, txt_mod_mlp = self.txt_mod(temb).chunk(2, dim=-1) # [B, 3*dim] each + + img_normed = self.img_norm1(image) + img_modulated, img_gate = self._modulate(img_normed, img_mod_attn) + + txt_normed = self.txt_norm1(text) + txt_modulated, txt_gate = self._modulate(txt_normed, txt_mod_attn) + + if self.merge_3D: # ref to edit image (no noise), uncond by timestep + # feat_3D_mod_attn, feat_3D_mod_mlp = self.img_mod(temb).chunk(2, dim=-1) + feat_3D_normed = self.img_norm1(feat_3D) + # feat_3D_modulated, feat_3D_gate = self._modulate(feat_3D_normed, feat_3D_mod_attn) + feat_3D_modulated = feat_3D_normed + else: + feat_3D_modulated = feat_3D + + + if add_attn and prope is not None: + img_prope_out, _3D_prope_out = self.prope_attn( # self_attn_3D + image=img_modulated, + feat_3D=feat_3D_modulated, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + prope=prope, + add_prope=add_prope, + block_id=block_id, + ) + + img_attn_out, txt_attn_out = self.attn( + image=img_modulated, + text=txt_modulated, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + prope=None, + ) + img_attn_out = img_attn_out + img_prope_out + elif not add_attn or prope is None: + img_attn_out, txt_attn_out = self.attn( + image=img_modulated, + text=txt_modulated, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + prope=prope, + ) + + image = image + img_gate * img_attn_out + text = text + txt_gate * txt_attn_out + + img_normed_2 = self.img_norm2(image) + img_modulated_2, img_gate_2 = self._modulate(img_normed_2, img_mod_mlp) + + txt_normed_2 = self.txt_norm2(text) + txt_modulated_2, txt_gate_2 = self._modulate(txt_normed_2, txt_mod_mlp) + + img_mlp_out = self.img_mlp(img_modulated_2) + txt_mlp_out = self.txt_mlp(txt_modulated_2) + + image = image + img_gate_2 * img_mlp_out + text = text + txt_gate_2 * txt_mlp_out + + if self.merge_3D: # ref to edit image (no noise), uncond by timestep + feat_3D = feat_3D + _3D_prope_out + feat_3D_normed_2 = self.img_norm2(feat_3D) + feat_3D_mlp_out = self.img_mlp(feat_3D_normed_2) + feat_3D = feat_3D + feat_3D_mlp_out + + return text, image, feat_3D + else: + return text, image + + +class MetaViewDiT(torch.nn.Module): + def __init__( + self, + num_layers: int = 60, + add_attn_type=None, + add_in_dim=None, + lora_rank=None, + merge_3D=False, + decode_3D=False, + _3d_dim=None, + ): + super().__init__() + + self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=[16,56,56], scale_rope=True) # only for text embed + + self.time_text_embed = TimestepEmbeddings(256, 3072, diffusers_compatible_format=True, scale=1000, align_dtype_to_timestep=True) + self.txt_norm = RMSNorm(3584, eps=1e-6) + + self.img_in = nn.Linear(64, 3072) + self.txt_in = nn.Linear(3584, 3072) + + self.transformer_blocks = nn.ModuleList( + [ + MetaViewTransformerBlock( + dim=3072, + num_attention_heads=24, + attention_head_dim=128, + ## add args + add_attn=True, + add_attn_type=add_attn_type, + add_in_dim=add_in_dim, + lora_rank=lora_rank, + merge_3D=merge_3D, + ) + for _ in range(num_layers) + ] + ) + self.norm_out = AdaLayerNorm(3072, single=True) + self.proj_out = nn.Linear(3072, 64) + + if merge_3D: + self._3D_in = nn.Linear(_3d_dim, add_in_dim) + if decode_3D: + self.norm_3D_out = nn.LayerNorm(add_in_dim, elementwise_affine=False, eps=1e-6) + self.proj_3D_out = nn.Linear(add_in_dim, _3d_dim) + + self.PRoPE = None + self.add_PRoPE = None + self.add_attn = True + + + def process_entity_masks(self, latents, prompt_emb, prompt_emb_mask, entity_prompt_emb, entity_prompt_emb_mask, entity_masks, height, width, image, img_shapes): + # prompt_emb + all_prompt_emb = entity_prompt_emb + [prompt_emb] + all_prompt_emb = [self.txt_in(self.txt_norm(local_prompt_emb)) for local_prompt_emb in all_prompt_emb] + all_prompt_emb = torch.cat(all_prompt_emb, dim=1) + + # image_rotary_emb + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + entity_seq_lens = [emb_mask.sum(dim=1).tolist() for emb_mask in entity_prompt_emb_mask] + entity_rotary_emb = [self.pos_embed(img_shapes, entity_seq_len, device=latents.device)[1] for entity_seq_len in entity_seq_lens] + txt_rotary_emb = torch.cat(entity_rotary_emb + [image_rotary_emb[1]], dim=0) + image_rotary_emb = (image_rotary_emb[0], txt_rotary_emb) + + # attention_mask + repeat_dim = latents.shape[1] + max_masks = entity_masks.shape[1] + entity_masks = entity_masks.repeat(1, 1, repeat_dim, 1, 1) + entity_masks = [entity_masks[:, i, None].squeeze(1) for i in range(max_masks)] + global_mask = torch.ones_like(entity_masks[0]).to(device=latents.device, dtype=latents.dtype) + entity_masks = entity_masks + [global_mask] + + N = len(entity_masks) + batch_size = entity_masks[0].shape[0] + seq_lens = [mask_.sum(dim=1).item() for mask_ in entity_prompt_emb_mask] + [prompt_emb_mask.sum(dim=1).item()] + total_seq_len = sum(seq_lens) + image.shape[1] + patched_masks = [] + for i in range(N): + patched_mask = rearrange(entity_masks[i], "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + patched_masks.append(patched_mask) + attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), dtype=torch.bool).to(device=entity_masks[0].device) + + # prompt-image attention mask + image_start = sum(seq_lens) + image_end = total_seq_len + cumsum = [0] + single_image_seq = image_end - image_start + for length in seq_lens: + cumsum.append(cumsum[-1] + length) + for i in range(N): + prompt_start = cumsum[i] + prompt_end = cumsum[i+1] + image_mask = torch.sum(patched_masks[i], dim=-1) > 0 + image_mask = image_mask.unsqueeze(1).repeat(1, seq_lens[i], 1) + # repeat image mask to match the single image sequence length + repeat_time = single_image_seq // image_mask.shape[-1] + image_mask = image_mask.repeat(1, 1, repeat_time) + # prompt update with image + attention_mask[:, prompt_start:prompt_end, image_start:image_end] = image_mask + # image update with prompt + attention_mask[:, image_start:image_end, prompt_start:prompt_end] = image_mask.transpose(1, 2) + # prompt-prompt attention mask, let the prompt tokens not attend to each other + for i in range(N): + for j in range(N): + if i == j: + continue + start_i, end_i = cumsum[i], cumsum[i+1] + start_j, end_j = cumsum[j], cumsum[j+1] + attention_mask[:, start_i:end_i, start_j:end_j] = False + + attention_mask = attention_mask.float() + attention_mask[attention_mask == 0] = float('-inf') + attention_mask[attention_mask == 1] = 0 + attention_mask = attention_mask.to(device=latents.device, dtype=latents.dtype).unsqueeze(1) + + return all_prompt_emb, image_rotary_emb, attention_mask + + + def forward( + self, + latents=None, + timestep=None, + prompt_emb=None, + prompt_emb_mask=None, + height=None, + width=None, + ): + img_shapes = [(latents.shape[0], latents.shape[2]//2, latents.shape[3]//2)] + # print(latentimg_shapes) + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + + image = rearrange(latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + image = self.img_in(image) + text = self.txt_in(self.txt_norm(prompt_emb)) + + conditioning = self.time_text_embed(timestep, image.dtype) + + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + + for block in self.transformer_blocks: + text, image = block( + image=image, + text=text, + temb=conditioning, + image_rotary_emb=image_rotary_emb, + ) + + image = self.norm_out(image, conditioning) + image = self.proj_out(image) + + latents = rearrange(image, "B (H W) (C P Q) -> B C (H P) (W Q)", H=height//16, W=width//16, P=2, Q=2) + return image diff --git a/src/MetaView_pipeline.py b/src/MetaView_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..b70d4c9e45514cc51d278540cf7b75054dfb83ac --- /dev/null +++ b/src/MetaView_pipeline.py @@ -0,0 +1,613 @@ +# Adapted from https://github.com/modelscope/DiffSynth-Studio + +import torch, math +from PIL import Image +from typing import Union +from tqdm import tqdm +from einops import rearrange +import numpy as np + +from diffsynth.diffusion import FlowMatchScheduler +from diffsynth.core import ModelConfig, gradient_checkpoint_forward +from diffsynth.diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput + +from diffsynth.models.qwen_image_text_encoder import QwenImageTextEncoder +from diffsynth.models.qwen_image_vae import QwenImageVAE +from diffsynth.models.qwen_image_controlnet import QwenImageBlockWiseControlNet + +from src.PRoPE import PropeDotProductAttention +from src.MetaView_dit import MetaViewDiT + +import torch.nn.functional as F + +class MetaViewPipeline(BasePipeline): + + def __init__(self, device="cuda", torch_dtype=torch.bfloat16): + super().__init__( + device=device, torch_dtype=torch_dtype, + height_division_factor=16, width_division_factor=16, + ) + from transformers import Qwen2Tokenizer, Qwen2VLProcessor + + self.scheduler = FlowMatchScheduler("Qwen-Image") + self.text_encoder: QwenImageTextEncoder = None + self.dit: MetaViewDiT = None + self.vae: QwenImageVAE = None + self.blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None + self.tokenizer: Qwen2Tokenizer = None + self.processor: Qwen2VLProcessor = None + self.in_iteration_models = ("dit", "blockwise_controlnet") + self.units = [ + MetaViewUnit_ShapeChecker(), + MetaViewUnit_NoiseInitializer(), + MetaViewUnit_InputImageEmbedder(), + MetaViewUnit_EditImageEmbedder(), + MetaViewUnit_PromptEmbedder(), + ] + self.model_fn = model_fn_MetaView + + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = "cuda", + model_configs: list[ModelConfig] = [], + tokenizer_config: ModelConfig = ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), + processor_config: ModelConfig = None, + vram_limit: float = None, + ): + # Initialize pipeline + pipe = MetaViewPipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + + # Fetch models + pipe.text_encoder = model_pool.fetch_model("qwen_image_text_encoder") + pipe.dit = model_pool.fetch_model("metaview_dit") + pipe.vae = model_pool.fetch_model("qwen_image_vae") + pipe.blockwise_controlnet = QwenImageBlockwiseMultiControlNet(model_pool.fetch_model("qwen_image_blockwise_controlnet", index="all")) + if tokenizer_config is not None: + tokenizer_config.download_if_necessary() + from transformers import Qwen2Tokenizer + pipe.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_config.path) + if processor_config is not None: + processor_config.download_if_necessary() + from transformers import Qwen2VLProcessor + pipe.processor = Qwen2VLProcessor.from_pretrained(processor_config.path) + + # VRAM Management + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + + @torch.no_grad() + def __call__( + self, + # Prompt + prompt: str, + negative_prompt: str = "", + cfg_scale: float = 4.0, + # Image + input_image: Image.Image = None, + denoising_strength: float = 1.0, + # Inpaint + inpaint_mask: Image.Image = None, + inpaint_blur_size: int = None, + inpaint_blur_sigma: float = None, + # Shape + height: int = 1328, + width: int = 1328, + # Randomness + seed: int = None, + rand_device: str = "cpu", + # Steps + num_inference_steps: int = 30, + exponential_shift_mu: float = None, + # Blockwise ControlNet + blockwise_controlnet_inputs: list[ControlNetInput] = None, + # EliGen + eligen_entity_prompts: list[str] = None, + eligen_entity_masks: list[Image.Image] = None, + eligen_enable_on_negative: bool = False, + # Qwen-Image-Edit + edit_image: Image.Image = None, + edit_image_auto_resize: bool = True, + edit_rope_interpolation: bool = False, + # In-context control + context_image: Image.Image = None, + # Tile + tiled: bool = False, + tile_size: int = 128, + tile_stride: int = 64, + # Progress bar + progress_bar_cmd = tqdm, + # added prope + viewmats = None, # [b, 2, 4, 4] order (target, edit) + Ks = None, # [b, 2, 3, 3] + prope_dim_arrange = [16, 56, 56], + add_attn = True, + add_3D = False, + feat_3D = None, + depth = None, + merge_3D = False, + val = False, + batch_size = 1, + ): + # Scheduler + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, dynamic_shift_len=(height // 16) * (width // 16), exponential_shift_mu=exponential_shift_mu) + + # Parameters + inputs_posi = { + "prompt": prompt, + } + inputs_nega = { + "negative_prompt": [negative_prompt], + } + inputs_shared = { + "cfg_scale": cfg_scale, + "input_image": input_image, "denoising_strength": denoising_strength, + "inpaint_mask": inpaint_mask, "inpaint_blur_size": inpaint_blur_size, "inpaint_blur_sigma": inpaint_blur_sigma, + "height": height, "width": width, + "seed": seed, "rand_device": rand_device, + "num_inference_steps": num_inference_steps, + "blockwise_controlnet_inputs": blockwise_controlnet_inputs, + "tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride, + "eligen_entity_prompts": eligen_entity_prompts, "eligen_entity_masks": eligen_entity_masks, "eligen_enable_on_negative": eligen_enable_on_negative, + "edit_image": edit_image, "edit_image_auto_resize": edit_image_auto_resize, "edit_rope_interpolation": edit_rope_interpolation, + "context_image": context_image, + # add camera param + "viewmats": viewmats, + "Ks": Ks, + "prope_dim_arrange": prope_dim_arrange, + "add_attn": add_attn, + "add_3D": add_3D, + "feat_3D": feat_3D, + "depth": depth, + "merge_3D": merge_3D, + "val": val, + "batch_size": batch_size, + } + for unit in self.units: + inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + + # Denoise + self.load_models_to_device(self.in_iteration_models) + models = {name: getattr(self, name) for name in self.in_iteration_models} + for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + noise_pred = self.cfg_guided_model_fn( + self.model_fn, cfg_scale, + inputs_shared, inputs_posi, inputs_nega, + **models, timestep=timestep, progress_id=progress_id + ) + inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) + # print(inputs_shared["latents"]) + + # Decode + self.load_models_to_device(['vae']) + image = self.vae.decode(inputs_shared["latents"], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + image = self.vae_output_to_image(image) + self.load_models_to_device([]) + + return image + + +class QwenImageBlockwiseMultiControlNet(torch.nn.Module): + def __init__(self, models: list[QwenImageBlockWiseControlNet]): + super().__init__() + if not isinstance(models, list): + models = [models] + self.models = torch.nn.ModuleList(models) + for model in models: + if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"): + self.vram_management_enabled = True + + def preprocess(self, controlnet_inputs: list[ControlNetInput], conditionings: list[torch.Tensor], **kwargs): + processed_conditionings = [] + for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): + conditioning = rearrange(conditioning, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) + model_output = self.models[controlnet_input.controlnet_id].process_controlnet_conditioning(conditioning) + processed_conditionings.append(model_output) + return processed_conditionings + + def blockwise_forward(self, image, conditionings: list[torch.Tensor], controlnet_inputs: list[ControlNetInput], progress_id, num_inference_steps, block_id, **kwargs): + res = 0 + for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): + progress = (num_inference_steps - 1 - progress_id) / max(num_inference_steps - 1, 1) + if progress > controlnet_input.start + (1e-4) or progress < controlnet_input.end - (1e-4): + continue + model_output = self.models[controlnet_input.controlnet_id].blockwise_forward(image, conditioning, block_id) + res = res + model_output * controlnet_input.scale + return res + + +class MetaViewUnit_ShapeChecker(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width"), + output_params=("height", "width"), + ) + + def process(self, pipe: MetaViewPipeline, height, width): + height, width = pipe.check_resize_height_width(height, width) + return {"height": height, "width": width} + + + +class MetaViewUnit_NoiseInitializer(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("height", "width", "seed", "rand_device", "batch_size"), + output_params=("noise",), + ) + + def process(self, pipe: MetaViewPipeline, height, width, seed, rand_device, batch_size): + noise = pipe.generate_noise((batch_size, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype) + return {"noise": noise} + + + +class MetaViewUnit_InputImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("input_image", "noise", "tiled", "tile_size", "tile_stride"), + output_params=("latents", "input_latents"), + onload_model_names=("vae",) + ) + + def process(self, pipe: MetaViewPipeline, input_image, noise, tiled, tile_size, tile_stride): + if input_image is None: + return {"latents": noise, "input_latents": None} + pipe.load_models_to_device(['vae']) + + if isinstance(input_image, list): + input_latents = [] + for input_img in input_image: + img = pipe.preprocess_image(input_img).to(device=pipe.device, dtype=pipe.torch_dtype) + input_latent = pipe.vae.encode(img, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + input_latents.append(input_latent) + input_latents = torch.cat(input_latents, dim=0) # B C H W + else: + # single PIL img, ret [1, c, h, w] + image = pipe.preprocess_image(input_image).to(device=pipe.device, dtype=pipe.torch_dtype) + input_latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + + assert noise.shape[0] == input_latents.shape[0] + + if pipe.scheduler.training: + return {"latents": noise, "input_latents": input_latents} + else: + latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) + return {"latents": latents, "input_latents": input_latents} + + +class MetaViewUnit_EditImageEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + input_params=("edit_image", "tiled", "tile_size", "tile_stride", "edit_image_auto_resize"), + output_params=("edit_latents", "edit_image"), + onload_model_names=("vae",) + ) + + + def calculate_dimensions(self, target_area, ratio): + import math + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + + def edit_image_auto_resize(self, edit_image): + calculated_width, calculated_height = self.calculate_dimensions(1024 * 1024, edit_image.size[0] / edit_image.size[1]) + return edit_image.resize((calculated_width, calculated_height)) + + + def process(self, pipe: MetaViewPipeline, edit_image, tiled, tile_size, tile_stride, edit_image_auto_resize=False): + if edit_image is None: + return {} + pipe.load_models_to_device(self.onload_model_names) + if isinstance(edit_image, Image.Image): + # resized_edit_image = self.edit_image_auto_resize(edit_image) if edit_image_auto_resize else edit_image + resized_edit_image = edit_image # skip resize + edit_image = pipe.preprocess_image(resized_edit_image).to(device=pipe.device, dtype=pipe.torch_dtype) + edit_latents = pipe.vae.encode(edit_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + else: + resized_edit_image, edit_latents = [], [] + for image in edit_image: + # if edit_image_auto_resize: + # image = self.edit_image_auto_resize(image) + resized_edit_image.append(image) + image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) + latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + edit_latents.append(latents) + edit_latents = torch.cat(edit_latents, dim=0) # B C H W + return {"edit_latents": edit_latents, "edit_image": resized_edit_image} + + + +class MetaViewUnit_PromptEmbedder(PipelineUnit): + def __init__(self): + super().__init__( + seperate_cfg=True, + input_params_posi={"prompt": "prompt"}, + input_params_nega={"prompt": "negative_prompt"}, + input_params=("edit_image",), + output_params=("prompt_emb", "prompt_emb_mask"), + onload_model_names=("text_encoder",) + ) + + def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + return split_result + + def calculate_dimensions(self, target_area, ratio): + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + def resize_image(self, image, target_area=384*384): + width, height = self.calculate_dimensions(target_area, image.size[0] / image.size[1]) + return image.resize((width, height)) + + def encode_prompt(self, pipe: MetaViewPipeline, prompt): + template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 34 + txt = [template.format(e) for e in prompt] + model_inputs = pipe.tokenizer(txt, max_length=4096+drop_idx, padding=True, truncation=True, return_tensors="pt").to(pipe.device) + if model_inputs.input_ids.shape[1] >= 1024: + print(f"Warning!!! QwenImage model was trained on prompts up to 512 tokens. Current prompt requires {model_inputs['input_ids'].shape[1] - drop_idx} tokens, which may lead to unpredictable behavior.") + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def encode_prompt_edit(self, pipe: MetaViewPipeline, prompt, edit_image): + template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 64 + txt = [template.format(e) for e in prompt] + # print(txt) + model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def encode_prompt_edit_batch(self, pipe: MetaViewPipeline, prompt, edit_image): + # list batch + template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 64 + txt = [template.format(e) for e in prompt] + split_hidden_states_list = [] + for i in range(len(prompt)): + model_inputs = pipe.processor(text=[txt[i]], images=[edit_image[i]], padding=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) #tuple (1) + # print(type(split_hidden_states[0])) + # print(len(split_hidden_states[0])) + split_hidden_states_list.append(split_hidden_states[0]) + + split_hidden_states = [e[drop_idx:] for e in split_hidden_states_list] + + return split_hidden_states + + def encode_prompt_edit_multi(self, pipe: MetaViewPipeline, prompt, edit_image): + template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + drop_idx = 64 + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + base_img_prompt = "".join([img_prompt_template.format(i + 1) for i in range(len(edit_image))]) + txt = [template.format(base_img_prompt + e) for e in prompt] + edit_image = [self.resize_image(image) for image in edit_image] + model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) + hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] + split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + return split_hidden_states + + def process(self, pipe: MetaViewPipeline, prompt, edit_image=None) -> dict: + #prompt [n] str list + pipe.load_models_to_device(self.onload_model_names) + if pipe.text_encoder is not None: + # prompt = [prompt] + if edit_image is None: + split_hidden_states = self.encode_prompt(pipe, prompt) + elif isinstance(edit_image, Image.Image): + split_hidden_states = self.encode_prompt_edit(pipe, prompt, edit_image) + elif isinstance(edit_image, list): # batch + split_hidden_states = self.encode_prompt_edit_batch(pipe, prompt, edit_image) + # else: + # split_hidden_states = self.encode_prompt_edit_multi(pipe, prompt, edit_image) + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) + encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) + prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) + return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask} + else: + return {} + + + +def model_fn_MetaView( + dit: MetaViewDiT = None, + blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None, + latents=None, + timestep=None, + prompt_emb=None, + prompt_emb_mask=None, + height=None, + width=None, + blockwise_controlnet_conditioning=None, + blockwise_controlnet_inputs=None, + progress_id=0, + num_inference_steps=1, + entity_prompt_emb=None, + entity_prompt_emb_mask=None, + entity_masks=None, + edit_latents=None, + context_latents=None, + enable_fp8_attention=False, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + edit_rope_interpolation=False, + viewmats=None, # camera param + Ks=None, + feat_3D=None, + prope_dim_arrange=None, + add_attn=False, + add_3D=False, + depth=None, + merge_3D=False, + decode_3D=False, + val=False, + **kwargs +): + img_shapes = [(1, latents.shape[2]//2, latents.shape[3]//2)] + txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() + timestep = timestep / 1000 + + image = rearrange(latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) + image_seq_len = image.shape[1] + + + if edit_latents is not None: # only single edit imgß + e = edit_latents # B C H W + img_shapes += [(1, e.shape[2]//2, e.shape[3]//2)] + edit_image = [rearrange(e, "B C (H P) (W Q) -> B (H W) (C P Q)", H=e.shape[2]//2, W=e.shape[3]//2, P=2, Q=2)] + image = torch.cat([image] + edit_image, dim=1) + + # print(img_shapes) + # print(image.shape) + # print(prompt_emb.shape) + # print(txt_seq_lens) + + # order tgt(latent, gt), src(edit_image ref) + # resize to 1024*1024 + # print("image ",image.shape) # ([1, 8184 (62 * 66 * 2), 64]) + # print("latents ",latents.shape) #[1, 16, 124, 132] + + # [(1, 33, 60), (1, 33, 60)] + # 960 528 + # [(1, 33, 60), (1, 33, 60)] + # 960 528 + # print(img_shapes) # (1, 62, 66), (1, 62, 66) + # print(width, height) + + image = dit.img_in(image) + conditioning = dit.time_text_embed(timestep, image.dtype) + + text = dit.txt_in(dit.txt_norm(prompt_emb)) + if edit_rope_interpolation: + image_rotary_emb = dit.pos_embed.forward_sampling(img_shapes, txt_seq_lens, device=latents.device) + else: + image_rotary_emb = dit.pos_embed(img_shapes, txt_seq_lens, device=latents.device) + # add prope + if viewmats is not None: + if depth is not None: # b n h w + depth = F.interpolate(depth, size=(height // 16, width // 16), mode='bilinear', align_corners=False) + depth = depth.to(image.device) + # print("depth:", depth.shape) + # print("image:", image.shape) + + # depth_np = depth[0, 1].detach().to(torch.float).cpu().numpy() + # depth_min, depth_max = depth_np.min(), depth_np.max() + # depth_norm = (depth_np - depth_min) / (depth_max - depth_min) * 255.0 + # depth_norm = depth_norm.astype(np.uint8) + # depth_save = Image.fromarray(depth_norm, 'L') + # depth_save.save(f"tmp/{depth_max}.png") + + + dit.PRoPE = PropeDotProductAttention( + head_dim=128, + patches_x=width // 16, + patches_y=height // 16, + image_width=width, + image_height=height, + freq_base=10000, #TODO 100? + dim_arrange=prope_dim_arrange, + ) + dit.PRoPE = dit.PRoPE.to(image.device) + dit.PRoPE._precompute_and_cache_apply_fns(viewmats.to(image.device), Ks.to(image.device), depth) # b, frames, h, w + + if feat_3D is not None: + dit.add_PRoPE = PropeDotProductAttention( + head_dim=128, + patches_x=width // 16, + patches_y=height // 16, + image_width=width, + image_height=height, + freq_base=10000, + dim_arrange=prope_dim_arrange, + ) + dit.add_PRoPE = dit.add_PRoPE.to(image.device) + if depth is not None: + dit.add_PRoPE._precompute_and_cache_apply_fns(viewmats[:, 1:2, :, :].to(image.device), Ks[:, 1:2, :, :].to(image.device), depth[:, 1:2, :, :]) + else: + dit.add_PRoPE._precompute_and_cache_apply_fns(viewmats[:, 1:2, :, :].to(image.device), Ks[:, 1:2, :, :].to(image.device)) + attention_mask = None + + if feat_3D is not None: + h_3D, w_3D = feat_3D.shape[1], feat_3D.shape[2] + feat_3D = rearrange(feat_3D, 'b h w d -> b (h w) d') + + if merge_3D: + feat_3D = dit._3D_in(feat_3D) + + for block_id, block in enumerate(dit.transformer_blocks): + if merge_3D: + text, image, feat_3D = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + image=image, + text=text, + temb=conditioning, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + prope=dit.PRoPE, # prope + add_prope=dit.add_PRoPE, + add_attn=add_attn, + feat_3D=feat_3D, + block_id=block_id, + ) + else: + text, image = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + image=image, + text=text, + temb=conditioning, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + enable_fp8_attention=enable_fp8_attention, + prope=dit.PRoPE, # prope + add_prope=dit.add_PRoPE, + add_attn=add_attn, + feat_3D=feat_3D, + block_id=block_id, + ) + + image = dit.norm_out(image, conditioning) + image = dit.proj_out(image) + image = image[:, :image_seq_len] + + latents = rearrange(image, "B (H W) (C P Q) -> B C (H P) (W Q)", H=height//16, W=width//16, P=2, Q=2) + + if val: + return latents + + if decode_3D: + feat_3D = dit.norm_3D_out(feat_3D) + feat_3D = dit.proj_3D_out(feat_3D) + latents_3D = feat_3D.unsqueeze(0).unsqueeze(0) + latents_3D = list(torch.chunk(latents_3D, chunks=4, dim=-1)) + return latents, latents_3D + + return latents, None diff --git a/src/PRoPE.py b/src/PRoPE.py new file mode 100644 index 0000000000000000000000000000000000000000..99aadd6563dac08ac84fe381228621fefece209e --- /dev/null +++ b/src/PRoPE.py @@ -0,0 +1,625 @@ +# MIT License +# +# Adapted from the official implementation of PRoPE +# "Cameras as Relative Positional Encoding" https://arxiv.org/pdf/2507.10496 +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from functools import partial +from typing import Callable, Optional, Tuple, List + +import torch +import torch.nn.functional as F + +from einops import rearrange + + +class PropeDotProductAttention(torch.nn.Module): + """PRoPE attention with precomputed RoPE coefficients.""" + + coeffs_x_0: torch.Tensor + coeffs_x_1: torch.Tensor + coeffs_y_0: torch.Tensor + coeffs_y_1: torch.Tensor + + def __init__( + self, + head_dim: int, + patches_x: int, + patches_y: int, + image_width: int, + image_height: int, + freq_base: float = 10000.0, # qwen 10000 + freq_scale: float = 1.0, + dim_arrange = [16, 56, 56], # (frame ,height, width) default for qwen. + depth = None, + ): + super().__init__() + self.head_dim = head_dim + self.patches_x = patches_x + self.patches_y = patches_y + self.image_width = image_width + self.image_height = image_height + self.freq_base = freq_base + self.freq_scale = freq_scale + + self.use_PRoPE = False + self.dim_arrange = dim_arrange + + # fit Qwen scale-rope + pos_index_x = torch.arange(patches_x) + neg_index_x = torch.arange(patches_x).flip(0) * -1 - 1 + index_x = torch.cat([neg_index_x[-(patches_x - patches_x // 2) :], pos_index_x[: patches_x // 2]], dim=0) + # print(index_x) + ## Qwen rope apply order is frame, height, width! + # coeffs_x + coeffs_y: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( # + # torch.tile(torch.arange(patches_x), (patches_y,)), + torch.tile(index_x, (patches_y,)), + freq_base=freq_base, + freq_scale=freq_scale, + # feat_dim=head_dim // 4, + feat_dim=dim_arrange[2], + ) + + # fit Qwen scale-rope + pos_index_y = torch.arange(patches_y) + neg_index_y = torch.arange(patches_y).flip(0) * -1 - 1 + index_y = torch.cat([neg_index_y[-(patches_y - patches_y // 2) :], pos_index_y[: patches_y // 2]], dim=0) + # print(index_y) + ## Qwen rope apply order is frame, height, width! + # coeffs_y + coeffs_x: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( + # torch.repeat_interleave(torch.arange(patches_y), patches_x), + torch.repeat_interleave(index_y, patches_x), + freq_base=freq_base, + freq_scale=freq_scale, + # feat_dim=head_dim // 4, + feat_dim=dim_arrange[1], + ) + + # Do not save coeffs to checkpoint as `cameras` might change during testing. + self.register_buffer("coeffs_x_0", coeffs_x[0], persistent=False) + self.register_buffer("coeffs_x_1", coeffs_x[1], persistent=False) + self.register_buffer("coeffs_y_0", coeffs_y[0], persistent=False) + self.register_buffer("coeffs_y_1", coeffs_y[1], persistent=False) + + + + + # override load_state_dict to not load coeffs if they exist (for backward compatibility) + def load_state_dict(self, state_dict, strict=True): + # remove coeffs from state_dict + state_dict.pop("coeffs_x_0", None) + state_dict.pop("coeffs_x_1", None) + state_dict.pop("coeffs_y_0", None) + state_dict.pop("coeffs_y_1", None) + super().load_state_dict(state_dict, strict) + + def forward( + self, + q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + **kwargs, + ) -> torch.Tensor: + return prope_dot_product_attention( + q, + k, + v, + viewmats=viewmats, + Ks=Ks, + patches_x=self.patches_x, + patches_y=self.patches_y, + image_width=self.image_width, + image_height=self.image_height, + coeffs_x=(self.coeffs_x_0, self.coeffs_x_1), + coeffs_y=(self.coeffs_y_0, self.coeffs_y_1), + **kwargs, + ) + + def _precompute_and_cache_apply_fns( + self, + viewmats: torch.Tensor, + Ks: Optional[torch.Tensor], + depth = None, + ): + (batch, cameras, _, _) = viewmats.shape + assert viewmats.shape == (batch, cameras, 4, 4) + assert Ks is None or Ks.shape == (batch, cameras, 3, 3) + self.cameras = cameras + self.use_PRoPE = True + + self.apply_fn_q, self.apply_fn_kv, self.apply_fn_o = _prepare_apply_fns( + head_dim=self.head_dim, + viewmats=viewmats, + Ks=Ks, + patches_x=self.patches_x, + patches_y=self.patches_y, + image_width=self.image_width, + image_height=self.image_height, + coeffs_x=(self.coeffs_x_0, self.coeffs_x_1), + coeffs_y=(self.coeffs_y_0, self.coeffs_y_1), + dim_arrange=self.dim_arrange, + freq_base=self.freq_base, + freq_scale=self.freq_scale, + depth=depth + ) + + def _apply_to_q(self, q: torch.Tensor) -> torch.Tensor: + (batch, num_heads, seqlen, head_dim) = q.shape + # print("!!!", q.shape) + # print(self.cameras, self.patches_x, self.patches_y) + assert seqlen == self.cameras * self.patches_x * self.patches_y, f"seqlen:{seqlen}, {self.cameras}, {self.patches_x}, {self.patches_y}" + assert head_dim == self.head_dim + assert q.shape == (batch, num_heads, seqlen, head_dim) + assert self.apply_fn_q is not None + return self.apply_fn_q(q) + + def _apply_to_kv(self, kv: torch.Tensor) -> torch.Tensor: + (batch, num_heads, seqlen, head_dim) = kv.shape + assert seqlen == self.cameras * self.patches_x * self.patches_y, f"seqlen:{seqlen}, {self.cameras}, {self.patches_x}, {self.patches_y}" + assert head_dim == self.head_dim + assert kv.shape == (batch, num_heads, seqlen, head_dim) + assert self.apply_fn_kv is not None + return self.apply_fn_kv(kv) + + def _apply_to_o(self, o: torch.Tensor) -> torch.Tensor: + (batch, num_heads, seqlen, head_dim) = o.shape + assert seqlen == self.cameras * self.patches_x * self.patches_y + assert head_dim == self.head_dim + assert o.shape == (batch, num_heads, seqlen, head_dim) + assert self.apply_fn_o is not None + return self.apply_fn_o(o) + + +def prope_dot_product_attention( + q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + *, + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int, # How many patches wide is each image? + patches_y: int, # How many patches tall is each image? + image_width: int, # Width of the image. Used to normalize intrinsics. + image_height: int, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, +) -> torch.Tensor: + """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style + positional encoding. + + Currently, we assume that the sequence length is equal to: + + cameras * patches_x * patches_y + + And token ordering allows the `(seqlen,)` axis to be reshaped into + `(cameras, patches_x, patches_y)`. + """ + # We're going to assume self-attention: all inputs are the same shape. + (batch, num_heads, seqlen, head_dim) = q.shape + cameras = viewmats.shape[1] + assert q.shape == k.shape == v.shape + assert viewmats.shape == (batch, cameras, 4, 4) + assert Ks is None or Ks.shape == (batch, cameras, 3, 3) + assert seqlen == cameras * patches_x * patches_y + + apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns( + head_dim=head_dim, + viewmats=viewmats, + Ks=Ks, + patches_x=patches_x, + patches_y=patches_y, + image_width=image_width, + image_height=image_height, + coeffs_x=coeffs_x, + coeffs_y=coeffs_y, + ) + + out = F.scaled_dot_product_attention( + query=apply_fn_q(q), + key=apply_fn_kv(k), + value=apply_fn_kv(v), + **kwargs, + ) + out = apply_fn_o(out) + assert out.shape == (batch, num_heads, seqlen, head_dim) + return out + + +def _prepare_apply_fns( + head_dim: int, # Q/K/V will have this last dimension + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int, # How many patches wide is each image? + patches_y: int, # How many patches tall is each image? + image_width: int, # Width of the image. Used to normalize intrinsics. + image_height: int, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[torch.Tensor] = None, + coeffs_y: Optional[torch.Tensor] = None, + coeffs_z: Optional[torch.Tensor] = None, + dim_arrange = None, + freq_base = None, + freq_scale = None, + depth = None, +) -> Tuple[ + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], +]: + """Prepare transforms for PRoPE-style positional encoding.""" + device = viewmats.device + (batch, cameras, _, _) = viewmats.shape + + viewmats = viewmats.to(torch.float32) + Ks = Ks.to(torch.float32) + + # Normalize camera intrinsics. + if Ks is not None: + # Ks has been normalized in the dataset getitem !! + Ks_norm = Ks + # Compute the camera projection matrices we use in PRoPE. + # - K is an `image<-camera` transform. + # - viewmats is a `camera<-world` transform. + # - P = lift(K) @ viewmats is an `image<-world` transform. + P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats) + P_T = P.transpose(-1, -2) + P_inv = torch.einsum( + "...ij,...jk->...ik", + _invert_SE3(viewmats), + _lift_K(_invert_K(Ks_norm)), + ) + + else: + # GTA formula. P is `camera<-world` transform. + P = viewmats + P_T = P.transpose(-1, -2) + P_inv = _invert_SE3(viewmats) + + assert P.shape == P_inv.shape == (batch, cameras, 4, 4) + + # Precompute cos/sin terms for RoPE. We use tiles/repeats for 'row-major' + # broadcasting. + + assert coeffs_x is not None + if coeffs_x is None: + coeffs_x = _rope_precompute_coeffs( + torch.tile(torch.arange(patches_x, device=device), (patches_y * cameras,)), + freq_base=100.0, + freq_scale=1.0, + # feat_dim=head_dim // 4, + feat_dim=dim_arrange[1], + ) + assert coeffs_y is not None + if coeffs_y is None: + coeffs_y = _rope_precompute_coeffs( + torch.tile( + torch.repeat_interleave( + torch.arange(patches_y, device=device), patches_x + ), + (cameras,), + ), + freq_base=100.0, + freq_scale=1.0, + # feat_dim=head_dim // 4, + feat_dim=dim_arrange[2], + ) + + if torch.isnan(P_inv).any(): + print("!!P_inv has NaN!!!") + exit(0) + if torch.isnan(P_T).any(): + print("!!P_T has NaN!!!") + exit(0) + if torch.isnan(coeffs_x[0]).any() or torch.isnan(coeffs_x[1]).any(): + print("!!coeffs_x has NaN!!!") + exit(0) + if torch.isnan(coeffs_y[0]).any() or torch.isnan(coeffs_y[1]).any(): + print("!!coeffs_y has NaN!!!") + exit(0) + + + # Block-diagonal transforms to the inputs and outputs of the attention operator. + assert head_dim % 4 == 0 + + transforms_q = [ + (partial(_apply_tiled_projmat, matrix=P_T), dim_arrange[0]), + (partial(_rope_apply_coeffs, coeffs=coeffs_x), dim_arrange[1]), + (partial(_rope_apply_coeffs, coeffs=coeffs_y), dim_arrange[2]), + ] + transforms_kv = [ + (partial(_apply_tiled_projmat, matrix=P_inv), dim_arrange[0]), + (partial(_rope_apply_coeffs, coeffs=coeffs_x), dim_arrange[1]), + (partial(_rope_apply_coeffs, coeffs=coeffs_y), dim_arrange[2]), + ] + transforms_o = [ + (partial(_apply_tiled_projmat, matrix=P), dim_arrange[0]), + (partial(_rope_apply_coeffs, coeffs=coeffs_x, inverse=True), dim_arrange[1]), + (partial(_rope_apply_coeffs, coeffs=coeffs_y, inverse=True), dim_arrange[2]), + ] + + if len(dim_arrange) == 4: + index_z = rearrange(depth, 'b n h w -> b n (h w)') # (batch, frame, seq_len) + coeffs_z: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs_z( + index_z, + freq_base=freq_base, + freq_scale=freq_scale, + feat_dim=dim_arrange[3], + ) + coeffs_z_0 = coeffs_z[0] + coeffs_z_1 = coeffs_z[1] + coeffs_z = (coeffs_z_0, coeffs_z_1) + + transforms_q += [(partial(_rope_apply_coeffs_z, coeffs=coeffs_z), dim_arrange[3])] + transforms_kv += [(partial(_rope_apply_coeffs_z, coeffs=coeffs_z), dim_arrange[3])] + transforms_o += [(partial(_rope_apply_coeffs_z, coeffs=coeffs_z, inverse=True), dim_arrange[3])] + + apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) + return apply_fn_q, apply_fn_kv, apply_fn_o + + +def _apply_tiled_projmat( + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + matrix: torch.Tensor, # (batch, cameras, D, D) +) -> torch.Tensor: + """Apply projection matrix to features.""" + # - seqlen => (cameras, patches_x * patches_y) + # - feat_dim => (feat_dim // 4, 4) + + matrix = matrix.to(feats.dtype) + + (batch, num_heads, seqlen, feat_dim) = feats.shape + cameras = matrix.shape[1] + assert seqlen > cameras and seqlen % cameras == 0 + D = matrix.shape[-1] + assert matrix.shape == (batch, cameras, D, D) + assert feat_dim % D == 0 + # print(matrix.device, feats.device) + return torch.einsum( + "bcij,bncpkj->bncpki", + matrix, + feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)), + ).reshape(feats.shape) + + +def _rope_apply_coeffs_z( + feats: torch.Tensor, # (batch, num_heads, seqlen_total, feat_dim) + coeffs: Tuple[torch.Tensor, torch.Tensor], # (batch, 1, frame, seqlen_img, num_freqs) + inverse: bool = False, +) -> torch.Tensor: + """Apply RoPE coefficients to features. We adopt a 'split' ordering + convention. (in contrast to 'interleaved')""" + + # print("Inject z rope!!") + #TODO change to interleaved same as Qwen? + cos, sin = coeffs + + batch, num_heads, total_seq_len, feat_dim = feats.shape + _, __, frames, seq_len_per_img, num_freqs = cos.shape + + cos = cos.to(feats.dtype) + sin = sin.to(feats.dtype) + # We allow (cos, sin) to be either with shape (1, 1, seqlen, feat_dim // 2), + # or (1, 1, seqlen_per_image, feat_dim // 2) and we repeat it to + # match the shape of feats. + feats = feats.reshape((batch, num_heads, frames, seq_len_per_img, feat_dim)) + assert feats.shape[3] * frames == total_seq_len + # if cos.shape[2] != feats.shape[2]: + # n_repeats = feats.shape[2] // cos.shape[2] + # cos = cos.repeat(1, 1, n_repeats, 1) + # sin = sin.repeat(1, 1, n_repeats, 1) + assert len(cos.shape) == len(sin.shape) == len(feats.shape) == 5 + assert cos.shape[-1] == sin.shape[-1] == feats.shape[-1] // 2 + + # cos (batch, 1, frame, seqlen_img, feat_dim) + x_in = feats[..., ::2] # even # (batch, num_heads, frames, seqlen_img, feat_dim) + y_in = feats[..., 1::2] + + if inverse == False: # for qkv + x_out = cos * x_in - sin * y_in # broadcast on "num_heads" + y_out = sin * x_in + cos * y_in + else: # for out + x_out = cos * x_in + sin * y_in + y_out = -sin * x_in + cos * y_in + + res = torch.stack((x_out, y_out), dim=-1).flatten(start_dim=-2) + res = rearrange(res, 'b n f s d -> b n (f s) d') + # print(res.shape) + return res + +def _rope_precompute_coeffs_z( + positions: torch.Tensor, # (batch, frame, seq_len) + freq_base: float, + freq_scale: float, + feat_dim: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Precompute RoPE coefficients.""" + assert len(positions.shape) == 3 + assert feat_dim % 2 == 0 + num_freqs = feat_dim // 2 + freqs = freq_scale * ( + freq_base + ** ( + -torch.arange(num_freqs, device=positions.device)[None, None, None, :] + / num_freqs + ) + ) + # print(freqs.shape) + # print(positions[:128]) + angles = positions[:, None, :, :, None] * freqs + # Shape should be: `(batch, num_heads, frame, seqlen, num_freqs)`; we're + # broadcasting across `num_heads`. + assert angles.shape == (positions.shape[0], 1, positions.shape[1], positions.shape[2], num_freqs) + return torch.cos(angles), torch.sin(angles) + + +def _rope_precompute_coeffs( + positions: torch.Tensor, # (seqlen,) + freq_base: float, + freq_scale: float, + feat_dim: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Precompute RoPE coefficients.""" + assert len(positions.shape) == 1 + assert feat_dim % 2 == 0 + num_freqs = feat_dim // 2 + freqs = freq_scale * ( + freq_base + ** ( + -torch.arange(num_freqs, device=positions.device)[None, None, None, :] + / num_freqs + ) + ) + # print(freqs.shape) + # print(positions[:128]) + angles = positions[None, None, :, None] * freqs + # Shape should be: `(batch, num_heads, seqlen, num_freqs)`; we're + # broadcasting across `batch` and `num_heads`. + assert angles.shape == (1, 1, positions.shape[0], num_freqs) + return torch.cos(angles), torch.sin(angles) + + +if __name__ == '__main__': + + patches_x = 64 + patches_y = 64 + freq_base = 1 + freq_scale = 10000 + head_dim = 128 + + pos_index = torch.arange(patches_x) + neg_index = torch.arange(patches_x).flip(0) * -1 - 1 + index = torch.cat([neg_index[-(patches_x - patches_x // 2) :], pos_index[: patches_x // 2]], dim=0) + print(index) + print(torch.arange(patches_x)) + coeffs_x: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( + # torch.tile(torch.arange(patches_x), (patches_y,)), + torch.tile(index, (patches_y,)), + freq_base=freq_base, + freq_scale=freq_scale, + # feat_dim=head_dim // 4, + feat_dim=56, + ) + + coeffs_y: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( + torch.repeat_interleave(torch.arange(patches_y), patches_x), + freq_base=freq_base, + freq_scale=freq_scale, + # feat_dim=head_dim // 4, + feat_dim=56, + ) + +def _rope_apply_coeffs( + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + coeffs: Tuple[torch.Tensor, torch.Tensor], + inverse: bool = False, +) -> torch.Tensor: + """Apply RoPE coefficients to features. We adopt a 'split' ordering + convention. (in contrast to 'interleaved')""" + + #TODO change to interleaved same as Qwen? + cos, sin = coeffs + + cos = cos.to(feats.dtype) + sin = sin.to(feats.dtype) + # We allow (cos, sin) to be either with shape (1, 1, seqlen, feat_dim // 2), + # or (1, 1, seqlen_per_image, feat_dim // 2) and we repeat it to + # match the shape of feats. + if cos.shape[2] != feats.shape[2]: + n_repeats = feats.shape[2] // cos.shape[2] + cos = cos.repeat(1, 1, n_repeats, 1) + sin = sin.repeat(1, 1, n_repeats, 1) + assert len(feats.shape) == len(cos.shape) == len(sin.shape) == 4 + assert cos.shape[-1] == sin.shape[-1] == feats.shape[-1] // 2 + + x_in = feats[..., ::2] # even # (batch, num_heads, seqlen, feat_dim) + y_in = feats[..., 1::2] + + if inverse == False: # for qkv + x_out = cos * x_in - sin * y_in + y_out = sin * x_in + cos * y_in + else: # for out + x_out = cos * x_in + sin * y_in + y_out = -sin * x_in + cos * y_in + + res = torch.stack((x_out, y_out), dim=-1).flatten(start_dim=-2) + # print(res.shape) + return res + + + +def _apply_block_diagonal( + feats: torch.Tensor, # (..., dim) + func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Apply a block-diagonal function to an input array. + + Each function is specified as a tuple with form: + + ((Tensor) -> Tensor, int) + + Where the integer is the size of the input to the function. + """ + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes) + x_blocks = torch.split(feats, block_sizes, dim=-1) + out = torch.cat( + [f(x_block) for f, x_block in zip(funcs, x_blocks)], + dim=-1, + ) + assert out.shape == feats.shape, "Input/output shapes should match." + return out + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix.""" + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +def _lift_K(Ks: torch.Tensor) -> torch.Tensor: + """Lift 3x3 matrices to homogeneous 4x4 matrices.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device) + out[..., :3, :3] = Ks + out[..., 3, 3] = 1.0 + return out + + +def _invert_K(Ks: torch.Tensor) -> torch.Tensor: + """Invert 3x3 intrinsics matrices. Assumes no skew.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros_like(Ks) + out[..., 0, 0] = 1.0 / Ks[..., 0, 0] + out[..., 1, 1] = 1.0 / Ks[..., 1, 1] + out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0] + out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1] + out[..., 2, 2] = 1.0 + return out \ No newline at end of file diff --git a/src/inference.py b/src/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..e7aa670dcfbbf4d061f75ec8af01edf1b5ae8805 --- /dev/null +++ b/src/inference.py @@ -0,0 +1,226 @@ +import torch, os, sys, glob +import argparse +import math +import numpy as np +from PIL import Image + +sys.path.append(os.getcwd()) +sys.path.append("./DepthAnything3/src") + +from DepthAnything3.src.depth_anything_3.api import DepthAnything3 + +from diffsynth.core import ModelConfig +from diffsynth import load_state_dict +from src.MetaView_pipeline import MetaViewPipeline +import torch.nn.functional as F + +def compute_target_extrinsic(yaw_deg, pitch_deg, radius): + """ + Compute the camera extrinsic matrix (World-to-Camera) for rotation + around a sphere center in front of the camera. + Supports simultaneous yaw (left-right) and pitch (up-down) angles. + + Args: + yaw_deg (float): Yaw angle in degrees. + pitch_deg (float): Pitch angle in degrees. + radius (float): Distance from the rotation sphere center to the camera + (typically the depth of the target object). + Returns: + numpy.ndarray: 4x4 extrinsic matrix. + """ + yaw = np.radians(yaw_deg) + pitch = np.radians(pitch_deg) + + # Rotation matrix around Y axis (Yaw) + R_y = np.array([ + [np.cos(yaw), 0, np.sin(yaw)], + [0, 1, 0 ], + [-np.sin(yaw), 0, np.cos(yaw)] + ]) + + # Rotation matrix around X axis (Pitch) + R_x = np.array([ + [1, 0, 0 ], + [0, np.cos(pitch), -np.sin(pitch)], + [0, np.sin(pitch), np.cos(pitch) ] + ]) + + # Combined rotation (pitch first, then yaw) + R = R_y @ R_x + + # Set sphere center coordinates C + C = np.array([0.0, 0.0, radius]) + + # Compute translation vector t = C - R * C + t = C - R @ C + + # Construct 4x4 extrinsic matrix + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + + return T + +def main(): + parser = argparse.ArgumentParser(description="MetaView Interactive Inference CLI") + + # Core interactive parameters + parser.add_argument("--image_path", type=str, required=True, help="Path to the input image") + parser.add_argument("--output_path", type=str, default="./output_novel_view.png", help="Path to save the generated image") + parser.add_argument("--yaw", type=float, default=0.0, help="Yaw angle in degrees (e.g., 60 for right, -60 for left)") + parser.add_argument("--pitch", type=float, default=0.0, help="Pitch angle in degrees (e.g., 30 for top, -30 for bottom)") + parser.add_argument("--radius", type=float, default=None, help="Rotation radius. If None, auto-calculated from center depth.") + + # Model path parameters + parser.add_argument("--da3_giant_path", type=str, default="../../Depth-Anything-3/model/DA3-GIANT-1.1", help="Path to DA3 Giant") + parser.add_argument("--da3_depth_path", type=str, default="../../Depth-Anything-3/model/DA3NESTED-GIANT-LARGE-1.1", help="Path to DA3 Depth model") + parser.add_argument("--qwen_path", type=str, default=None, help="Base path to Qwen-Image-Edit models") + parser.add_argument("--ckpt_path", type=str, required=True, help="Path to the trained MetaView checkpoint (.safetensors)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Global parameter configuration + export_3D_feat_layers = [19, 27, 33, 39] + prope_dim_arrange = [64, 20, 20, 24] + add_depth = (len(prope_dim_arrange) == 4) + merge_3D = True + prompt = ["镜头视角转到指定位置"] + + # 1. Load input image + print(f"[*] Loading input image from {args.image_path}...") + original_image = Image.open(args.image_path).convert("RGB") + edit_image = original_image.resize((960, 528)) + + # ========================================================================= + # 2. Depth and geometry prior extraction (Depth & Intrinsics) + # ========================================================================= + print("[*] Loading DepthAnything3 Prior Models...") + + with torch.inference_mode(): + # Load feature extraction model (GIANT) + model_3D = DepthAnything3.from_pretrained(args.da3_giant_path).to(device=device) + + print(" -> Extracting 3D Features and Intrinsics...") + feat_3D_output = model_3D.inference([edit_image], export_feat_layers=export_3D_feat_layers, process_res=840) + + # Process intrinsics Ks + intri = feat_3D_output.intrinsics[0] + width = intri[0, 2] * 2 + height = intri[1, 2] * 2 + Ks_matrix = [ + [intri[0, 0] / width, 0.0, 0.0], + [0.0, intri[1, 1] / height, 0.0], + [0.0, 0.0, 1.0], + ] + Ks = torch.Tensor(Ks_matrix) + Ks = torch.stack([Ks, Ks], dim=0).unsqueeze(0) # Shape: (1, 2, 3, 3) + + # Process features feat_3D + feats = [torch.from_numpy(feat_3D_output.aux[f"feat_layer_{layer}"]) for layer in export_3D_feat_layers] + feat_3D = torch.cat(feats, dim=-1).to(dtype=torch.bfloat16, device=device) + + # Release feature extraction model + model_3D.to("cpu") + del model_3D + torch.cuda.empty_cache() + + # Load depth extraction model (NESTED) + print(" -> Extracting Depth Map...") + model_depth = DepthAnything3.from_pretrained(args.da3_depth_path).to(device=device) + prediction = model_depth.inference([edit_image], process_res=840) + + depth_edit = torch.Tensor(prediction.depth).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0).unsqueeze(0) # Shape: (1, 2, H, W) + + # Release depth model + model_depth.to("cpu") + del model_depth + torch.cuda.empty_cache() + + # ========================================================================= + # 3. Target pose calculation + # ========================================================================= + # Auto-derive radius: if user did not specify radius, use center depth from the depth map + if args.radius is None: + depth_squeeze = depth[0, 1] # Get real depth channel + z_c = depth_squeeze[depth_squeeze.shape[0]//2, depth_squeeze.shape[1]//2].item() + args.radius = z_c + print(f"[*] Auto-calculated rotation radius from center depth: {args.radius:.4f}") + + print(f"[*] Calculating Target Pose -> Yaw: {args.yaw}°, Pitch: {args.pitch}°, Radius: {args.radius}") + extrinsic_target = compute_target_extrinsic(args.yaw, args.pitch, args.radius) + extrinsic_source = np.eye(4) + + # Construct viewmats tensor: Shape (1, 2, 4, 4) -> [Target, Source] + viewmats = torch.Tensor(np.stack((extrinsic_target, extrinsic_source), axis=0)).unsqueeze(0) + + # ========================================================================= + # 4. Generation model loading and inference (DiT Pipeline) + # ========================================================================= + print("[*] Loading Qwen-Image-Edit Pipeline...") + + if args.qwen_path: + print(f"[*] Loading Qwen-Image-Edit from {args.qwen_path}") + pipe = MetaViewPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(path=glob.glob(f"{args.qwen_path}/Qwen-Image-Edit/transformer/diffusion_pytorch_model*.safetensors")), + ModelConfig(path=glob.glob(f"{args.qwen_path}/Qwen-Image-Edit/text_encoder/model*.safetensors")), + ModelConfig(path=glob.glob(f"{args.qwen_path}/Qwen-Image-Edit/vae/diffusion_pytorch_model.safetensors")), + ], + tokenizer_config=None, + processor_config=ModelConfig(path=f"{args.qwen_path}/Qwen-Image-Edit/processor/"), + ) + else: # Auto download + pipe = MetaViewPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"), + ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors"), + ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=None, + processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"), + ) + + print(f"[*] Loading MetaView Weights from {args.ckpt_path}...") + state_dict = load_state_dict(args.ckpt_path) + pipe.dit.load_state_dict(state_dict, strict=False) + + print("[*] Starting Generation (40 Steps)...") + with torch.inference_mode(): + generated_image = pipe( + prompt, edit_image=edit_image, edit_image_auto_resize=False, + seed=0, + viewmats=viewmats.to(device=device, dtype=torch.bfloat16), + Ks=Ks.to(device=device, dtype=torch.bfloat16), + prope_dim_arrange=prope_dim_arrange, + add_attn=True, + add_3D=True, + feat_3D=feat_3D, + depth=depth.to(device=device, dtype=torch.bfloat16) if add_depth else None, + merge_3D=merge_3D, + val=True, + num_inference_steps=40, + height=528, width=960, + ) + + # ========================================================================= + # 5. Save result (stitch source and generated images for comparison) + # ========================================================================= + stitched_image = Image.new('RGB', (960 * 2, 528), (255, 255, 255)) + stitched_image.paste(edit_image, (0, 0)) + stitched_image.paste(generated_image, (960, 0)) + + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + stitched_image.save(args.output_path) + print(f"Success! Result saved to {args.output_path}") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/lora.py b/src/lora.py new file mode 100644 index 0000000000000000000000000000000000000000..e17dd168378ae2b5731d5e1e158fd3de568946a5 --- /dev/null +++ b/src/lora.py @@ -0,0 +1,44 @@ +import inspect +import math +from typing import Callable, List, Optional, Tuple, Union +from einops import rearrange +import torch +from torch import nn +import torch.nn.functional as F +from torch import Tensor +from diffusers.models.attention_processor import Attention + +class LoRALinearLayer(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + rank: int = 4, + network_alpha: Optional[float] = None, + ): + super().__init__() + # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. + # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning + self.network_alpha = network_alpha + self.rank = rank + self.out_dim = out_dim + self.in_dim = in_dim + + self.down = nn.Linear(in_dim, rank, bias=False) + self.up = nn.Linear(rank, out_dim, bias=False) + + nn.init.normal_(self.down.weight, std=1 / rank) + nn.init.zeros_(self.up.weight) + + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_dtype = hidden_states.dtype + dtype = self.down.weight.dtype + + down_hidden_states = self.down(hidden_states.to(dtype)) + up_hidden_states = self.up(down_hidden_states) + + if self.network_alpha is not None: + up_hidden_states *= self.network_alpha / self.rank + + return up_hidden_states.to(orig_dtype) \ No newline at end of file diff --git a/src/metric.py b/src/metric.py new file mode 100644 index 0000000000000000000000000000000000000000..d0cdc817c26fcabbc669ece9bb961902ebd7003b --- /dev/null +++ b/src/metric.py @@ -0,0 +1,233 @@ +import cv2 +import flow_vis +import matplotlib.pyplot as plt +import numpy as np +import torch +import os, sys +from tqdm import tqdm +import argparse + +from einops import rearrange + +import lpips +from PIL import Image + +import torch.nn.functional as F + +sys.path.append(os.getcwd()) +sys.path.append("./UFM") + +from uniflowmatch.utils.geometry import get_meshgrid_torch + +from uniflowmatch.models.ufm import UniFlowMatchClassificationRefinement + +from skimage.metrics import peak_signal_noise_ratio, structural_similarity + +def warp_image_with_flow(source_image, source_mask, target_image, flow, thresh) -> np.ndarray: + """ + Warp the target to source image using the given flow vectors. + Flow vectors indicate the displacement from source to target. + + Args: + source_image: np.ndarray of shape (H, W, 3), normalized to [0, 1] + target_image: np.ndarray of shape (H, W, 3), normalized to [0, 1] + flow: np.ndarray of shape (H, W, 2) + source_mask: non_occluded mask represented in source image. + + Returns: + warped_image: target_image warped according to flow into frame of source image + np.ndarray of shape (H, W, 3), normalized to [0, 1] + + """ + # assert source_image.shape[-1] == 3 + # assert target_image.shape[-1] == 3 + + assert flow.shape[-1] == 2 + + # Get the shape of the source image + height, width = source_image.shape[:2] + target_height, target_width = target_image.shape[:2] + + # Create mesh grid + x, y = np.meshgrid(np.arange(width), np.arange(height)) + + # Apply flow displacements + flow_x, flow_y = flow[..., 0], flow[..., 1] + x_new = np.clip(x + flow_x, 0, target_width - 1) + 0.5 + y_new = np.clip(y + flow_y, 0, target_height - 1) + 0.5 + + x_new = (x_new / target_image.shape[1]) * 2 - 1 + y_new = (y_new / target_image.shape[0]) * 2 - 1 + + warped_image = F.grid_sample( + torch.from_numpy(target_image).permute(2, 0, 1)[None, ...].float(), + torch.from_numpy(np.stack([x_new, y_new], axis=-1)).float()[None, ...], + mode="bilinear", + align_corners=False, + ) + + warped_image = warped_image[0].permute(1, 2, 0).numpy() + + if source_mask is not None: + warped_image = warped_image * (source_mask > thresh) + + return warped_image + +def compute_distance_vectorized(flow_output, covisibility_src, covisibility, + src_thresh, thresh, target_image_shape): + height, width = target_image_shape[0], target_image_shape[1] + + src_mask = covisibility_src > src_thresh + valid_mask = src_mask & (covisibility > thresh) + penalty_mask = src_mask & ~valid_mask + + cnt = np.sum(src_mask) + + flow_normalized = flow_output / np.array([width, height])[:, None, None] + distances = np.sqrt(np.sum(flow_normalized**2, axis=0)) + + # total_dist = valid_dist + penalty + total_dist = np.sum(distances[valid_mask]) + np.sum(penalty_mask) * np.sqrt(2) + + return total_dist, int(cnt) + + + +parser = argparse.ArgumentParser() +parser.add_argument('--data_path', type=str, default='None', help='Image file path') +args = parser.parse_args() + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +thresh = 0.2 # generate +src_thresh = 0.3 # src + +refinement_model = True +resolution = 560 + +model = UniFlowMatchClassificationRefinement.from_pretrained( + "infinity1096/UFM-Refine" if resolution == 560 else "infinity1096/UFM-Refine-980" +).to(device) + +model.eval() + +data_path = args.data_path + +save = False +files = os.listdir(data_path) + +LPIPS = lpips.LPIPS(net="vgg").to(device) +total_psnr = 0 +total_ssim = 0 +total_lpips = 0 +total_dist = 0 + +count = len(files) +invalid = 0 +for idx, file in tqdm(enumerate(files), total=count): + print(file) + + combined_img = Image.open(os.path.join(data_path, file)).resize((960 * 3, 528)) + + width, height = combined_img.size + + length = width // 3 + + src_img = np.array(combined_img.crop((length * 0, 0, length * 1, height))) + generated_img = np.array(combined_img.crop((length * 1, 0, length * 2, height))) + gt_img = np.array(combined_img.crop((length * 2, 0, length * 3, height))) + + psnr = peak_signal_noise_ratio(generated_img, gt_img) + ssim = structural_similarity(generated_img, gt_img, channel_axis=-1) + + img1 = rearrange(torch.Tensor(generated_img).unsqueeze(0), 'b h w c -> b c h w').to(device) / 127.5 - 1.0 + img2 = rearrange(torch.Tensor(gt_img).unsqueeze(0), 'b h w c -> b c h w').to(device) / 127.5 - 1.0 + lpips = LPIPS(img1, img2) + + total_psnr += psnr + total_ssim += ssim + total_lpips += lpips + + source_image = gt_img + target_image = generated_img + # === Predict Correspondences === + result = model.predict_correspondences_batched( + source_image=torch.from_numpy(gt_img).to(device), + target_image=torch.from_numpy(generated_img).to(device), + ) + + result2 = model.predict_correspondences_batched( + source_image=torch.from_numpy(gt_img).to(device), + target_image=torch.from_numpy(src_img).to(device), + ) + covisibility_src = result2.covisibility.mask[0].cpu().numpy() + + flow_output = result.flow.flow_output[0].cpu().numpy() # 2 H W + covisibility = result.covisibility.mask[0].cpu().numpy() + + + dist, cnt = compute_distance_vectorized( + flow_output, covisibility_src, covisibility, + src_thresh, thresh, target_image.shape + ) + + if cnt: + dist /= cnt + dist = dist / np.sqrt(2) * 100 # normalized to 0~100 + else: + # non-covisible between src and gt + dist = 0 + invalid += 1 + + total_dist += dist + print('') + print(f"PSNR: {psnr}") + print(f"SSIM: {ssim}") + print(f"LPIPS: {lpips}") + print(f"dist: {dist} {cnt}") + + if not save: + continue + + # === Visualize Results === + fig, axs = plt.subplots(2, 3, figsize=(15, 5)) + + axs[0, 0].imshow(source_image) + axs[0, 0].set_title("Source Image") + + axs[0, 1].imshow(target_image) + axs[0, 1].set_title("Target Image") + + # Warp the image using flow + warped_image = warp_image_with_flow(source_image, None, target_image, flow_output.transpose(1, 2, 0), thresh) + warped_image = covisibility[..., None] * warped_image + (1 - covisibility[..., None]) * 255 * np.ones_like( + warped_image + ) + warped_image /= 255.0 + + axs[0, 2].imshow(warped_image) + axs[0, 2].set_title("Warped Image") + + # Flow visualization + flow_vis_image = flow_vis.flow_to_color(flow_output.transpose(1, 2, 0)) + axs[1, 0].imshow(flow_vis_image) + axs[1, 0].set_title(f"Flow Output (Valid at covisible region) {np.round(dist, decimals=2)}, psnr:{psnr:.2f}, ssim:{ssim:.2f}") + + # Covisibility mask + axs[1, 1].imshow(covisibility > thresh, cmap="gray", vmin=0, vmax=1) + axs[1, 1].set_title(f"Covisibility Mask ({thresh})") + + heatmap = axs[1, 2].imshow(covisibility, cmap="gray", vmin=0, vmax=1) + axs[1, 2].set_title("Covisibility Mask") + plt.colorbar(heatmap, ax=axs[1, 2]) + + if not os.path.exists(f"result_{save}"): + os.mkdir(f"result_{save}") + + plt.tight_layout() + plt.savefig(f"result_{save}/{file}.png") + +print(f"PSNR: {total_psnr / len(files):.2f} dB") +print(f"SSIM: {total_ssim / len(files):.4f}") +print(f"LPIPS: {total_lpips.item() / count:.4f}") +print(f"dist: {total_dist / (len(files) - invalid):.4f}") \ No newline at end of file diff --git a/src/training_module.py b/src/training_module.py new file mode 100644 index 0000000000000000000000000000000000000000..98c33d8d8314622956914db6909b544af081b852 --- /dev/null +++ b/src/training_module.py @@ -0,0 +1,232 @@ +import torch, json +from diffsynth.core import ModelConfig, load_state_dict +from diffsynth.utils.controlnet import ControlNetInput +from peft import LoraConfig, inject_adapter_in_model + +from src.MetaView_dit import MetaViewSelfAttention3D + +class MetaViewTrainingModule(torch.nn.Module): + def __init__(self): + super().__init__() + + + def to(self, *args, **kwargs): + for name, model in self.named_children(): + model.to(*args, **kwargs) + return self + + + def trainable_modules(self): + trainable_modules = filter(lambda p: p.requires_grad, self.parameters()) + return trainable_modules + + + def trainable_param_names(self): + trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters())) + trainable_param_names = set([named_param[0] for named_param in trainable_param_names]) + return trainable_param_names + + + def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None): + if lora_alpha is None: + lora_alpha = lora_rank + if isinstance(target_modules, list) and len(target_modules) == 1: + target_modules = target_modules[0] + lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules) + model = inject_adapter_in_model(lora_config, model) + if upcast_dtype is not None: + for param in model.parameters(): + if param.requires_grad: + param.data = param.to(upcast_dtype) + return model + + + def mapping_lora_state_dict(self, state_dict): + new_state_dict = {} + for key, value in state_dict.items(): + if "lora_A.weight" in key or "lora_B.weight" in key: + new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight") + new_state_dict[new_key] = value + elif "lora_A.default.weight" in key or "lora_B.default.weight" in key: + new_state_dict[key] = value + return new_state_dict + + + def export_trainable_state_dict(self, state_dict, remove_prefix=None): + trainable_param_names = self.trainable_param_names() + state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names} + if remove_prefix is not None: + state_dict_ = {} + for name, param in state_dict.items(): + if name.startswith(remove_prefix): + name = name[len(remove_prefix):] + state_dict_[name] = param + state_dict = state_dict_ + return state_dict + + + def transfer_data_to_device(self, data, device, torch_float_dtype=None): + if data is None: + return data + elif isinstance(data, torch.Tensor): + data = data.to(device) + if torch_float_dtype is not None and data.dtype in [torch.float, torch.float16, torch.bfloat16]: + data = data.to(torch_float_dtype) + return data + elif isinstance(data, tuple): + data = tuple(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data) + return data + elif isinstance(data, list): + data = list(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data) + return data + elif isinstance(data, dict): + data = {i: self.transfer_data_to_device(data[i], device, torch_float_dtype) for i in data} + return data + else: + return data + + def parse_vram_config(self, fp8=False, offload=False, device="cpu"): + if fp8: + return { + "offload_dtype": torch.float8_e4m3fn, + "offload_device": device, + "onload_dtype": torch.float8_e4m3fn, + "onload_device": device, + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": device, + "computation_dtype": torch.bfloat16, + "computation_device": device, + } + elif offload: + return { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": device, + "computation_dtype": torch.bfloat16, + "computation_device": device, + "clear_parameters": True, + } + else: + return {} + + def parse_model_configs(self, model_paths, model_id_with_origin_paths, fp8_models=None, offload_models=None, device="cpu"): + fp8_models = [] if fp8_models is None else fp8_models.split(",") + offload_models = [] if offload_models is None else offload_models.split(",") + model_configs = [] + if model_paths is not None: + model_paths = json.loads(model_paths) + for path in model_paths: + vram_config = self.parse_vram_config( + fp8=path in fp8_models, + offload=path in offload_models, + device=device + ) + model_configs.append(ModelConfig(path=path, **vram_config)) + if model_id_with_origin_paths is not None: + model_id_with_origin_paths = model_id_with_origin_paths.split(",") + for model_id_with_origin_path in model_id_with_origin_paths: + model_id, origin_file_pattern = model_id_with_origin_path.split(":") + vram_config = self.parse_vram_config( + fp8=model_id_with_origin_path in fp8_models, + offload=model_id_with_origin_path in offload_models, + device=device + ) + model_configs.append(ModelConfig(model_id=model_id, origin_file_pattern=origin_file_pattern, **vram_config)) + return model_configs + + + def switch_pipe_to_training_mode( + self, + pipe, + trainable_models=None, + lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None, + preset_lora_path=None, preset_lora_model=None, + task="sft", + ): + # Scheduler + pipe.scheduler.set_timesteps(1000, training=True) + + # Freeze untrainable models + if "prope_attn" in trainable_models: + pipe.eval() + pipe.requires_grad_(False) + cnt = 0 + for name, module in pipe.dit.named_modules(): + if isinstance(module, (MetaViewSelfAttention3D)): + module.train() + module.requires_grad_(True) + cnt += 1 + print(name) + if "3D" in name.split(".")[0]: # _3D_in, proj_3D_out + module.train() + module.requires_grad_(True) + print(name) + + print(f"{cnt} PRoPE Attention layers are trainable. ") + else: + pipe.freeze_except([] if trainable_models is None else trainable_models.split(",")) + + # Preset LoRA + if preset_lora_path is not None: + pipe.load_lora(getattr(pipe, preset_lora_model), preset_lora_path) + + # FP8 + # FP8 relies on a model-specific memory management scheme. + # It is delegated to the subclass. + + # Add LoRA to the base models + if lora_base_model is not None and not task.endswith(":data_process"): + if (not hasattr(pipe, lora_base_model)) or getattr(pipe, lora_base_model) is None: + print(f"No {lora_base_model} models in the pipeline. We cannot patch LoRA on the model. If this occurs during the data processing stage, it is normal.") + return + model = self.add_lora_to_model( + getattr(pipe, lora_base_model), + target_modules=lora_target_modules.split(","), + lora_rank=lora_rank, + upcast_dtype=pipe.torch_dtype, + ) + if lora_checkpoint is not None: + state_dict = load_state_dict(lora_checkpoint) + state_dict = self.mapping_lora_state_dict(state_dict) + load_result = model.load_state_dict(state_dict, strict=False) + print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys") + if len(load_result[1]) > 0: + print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}") + setattr(pipe, lora_base_model, model) + + + def split_pipeline_units(self, task, pipe, trainable_models=None, lora_base_model=None): + models_require_backward = [] + if trainable_models is not None: + models_require_backward += trainable_models.split(",") + if lora_base_model is not None: + models_require_backward += [lora_base_model] + if task.endswith(":data_process"): + print("!!!! task.endswith(:data_process)", ) + _, pipe.units = pipe.split_pipeline_units(models_require_backward) + elif task.endswith(":train"): + print("!!!! task.endswith(:train)", ) + pipe.units, _ = pipe.split_pipeline_units(models_require_backward) + return pipe + + def parse_extra_inputs(self, data, extra_inputs, inputs_shared): + controlnet_keys_map = ( + ("blockwise_controlnet_", "blockwise_controlnet_inputs",), + ("controlnet_", "controlnet_inputs"), + ) + controlnet_inputs = {} + for extra_input in extra_inputs: + for prefix, name in controlnet_keys_map: + if extra_input.startswith(prefix): + if name not in controlnet_inputs: + controlnet_inputs[name] = {} + controlnet_inputs[name][extra_input.replace(prefix, "")] = data[extra_input] + break + else: + inputs_shared[extra_input] = data[extra_input] + for name, params in controlnet_inputs.items(): + inputs_shared[name] = [ControlNetInput(**params)] + return inputs_shared diff --git a/src/unified_dataset.py b/src/unified_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..70b924b89b7572e4bea8c937db42b0871519d732 --- /dev/null +++ b/src/unified_dataset.py @@ -0,0 +1,805 @@ +from diffsynth.core.data.operators import * +import torch, json, pandas, sys, os +import numpy as np +from pathlib import Path +from PIL import Image +import cv2 + +sys.path.append(os.getcwd()) +sys.path.append("./DepthAnything3/src") + +from DepthAnything3.src.depth_anything_3.api import DepthAnything3 + +from openexr_numpy import imread, imwrite +import torch.nn.functional as F + +import pandas as pd + +class MetaViewUnifiedDataset(torch.utils.data.Dataset): + def __init__( + self, + base_path=None, metadata_path=None, + repeat=1, + data_file_keys=tuple(), + main_data_operator=lambda x: x, + special_operator_map=None, + prope=False, + debug=False, + mode="train", + norm_scale=1.0, + path_3D=None, + export_3D_feat_layers=None, + anno_src=None, + add_depth=False, + subset=None, + base_model = "qwen", + ): + self.base_model = base_model + self.base_path = base_path + self.metadata_path = metadata_path + self.repeat = repeat + self.data_file_keys = data_file_keys + self.main_data_operator = main_data_operator + self.cached_data_operator = LoadTorchPickle() + self.special_operator_map = {} if special_operator_map is None else special_operator_map + self.data = [] + self.cached_data = [] + + paths = base_path.split(";") + + if subset is None: + self.videos = os.listdir(base_path) + elif len(paths) > 1: + self.videos = [] + for path in paths: + dirs = os.listdir(path) + for sub in subset: + if sub in dirs: + p = os.path.join(path, sub) + folders = os.listdir(p) + for folder in folders: + self.videos.append(os.path.join(p, folder)) + else: + self.videos = [] + for sub in subset: + p = os.path.join(base_path, sub) + dirs = os.listdir(p) + for d in dirs: + self.videos.append(os.path.join(sub, d)) + self.total_length = len(self.videos) + self.prope = prope + self.mode = mode + self.norm_scale = norm_scale + self.anno_src = anno_src + self.add_depth = add_depth + self.subset = subset + + if prope: + self.load_from_cache = False + + self.model_3D = None + if path_3D is not None: + device = torch.device("cuda") + self.model_3D = DepthAnything3.from_pretrained(path_3D) + self.model_3D = self.model_3D.to(device=device) + export_3D_feat_layers = export_3D_feat_layers.split(",") + self.export_3D_feat_layers = [int(s) for s in export_3D_feat_layers] + + if debug: + self.max_dist() + exit(0) + + + @staticmethod + def default_image_operator( + base_path="", + max_pixels=1920*1080, height=None, width=None, + height_division_factor=16, width_division_factor=16, + ): + return RouteByType(operator_map=[ + (str, ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor)), + (list, SequencialProcess(ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor))), + ]) + + @staticmethod + def default_video_operator( + base_path="", + max_pixels=1920*1080, height=None, width=None, + height_division_factor=16, width_division_factor=16, + num_frames=81, time_division_factor=4, time_division_remainder=1, + ): + return RouteByType(operator_map=[ + (str, ToAbsolutePath(base_path) >> RouteByExtensionName(operator_map=[ + (("jpg", "jpeg", "png", "webp"), LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor) >> ToList()), + (("gif",), LoadGIF( + num_frames, time_division_factor, time_division_remainder, + frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), + )), + (("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), LoadVideo( + num_frames, time_division_factor, time_division_remainder, + frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), + )), + ])), + ]) + + + + def __getitem__(self, data_id): + if self.prope: + if len(self.base_path.split(";")) > 1: + return self.getitem_prope_all(data_id) + elif "DL3DV" in self.base_path and self.metadata_path is None: + return self.getitem_prope_DL3DV(data_id) + elif "DL3DV" in self.base_path and self.metadata_path: + return self.getitem_metadata_DL3DV(data_id) + elif self.metadata_path: + return self.getitem_metadata(data_id) + + return data + + def quick_check(self, extrinsics: torch.Tensor) -> bool: + """check extrinsics""" + if extrinsics.shape[-2:] != (4, 4): + print("Extrinsics wrong shape!") + return False + + t_norm = torch.norm(extrinsics[:, :3, 3], dim=1) + if t_norm.max() > 100: # 阈值根据应用设定 + print(f"Extrinsics shift too large! {t_norm.max()}") + return False + + R = extrinsics[..., :3, :3] + # 检查行列式接近1 + det_R = torch.det(R) + if not torch.allclose(det_R, torch.ones_like(det_R), atol=1e-4): + print("Extrinsics not ortho!") + return False + + # 检查最后一行 + last_row = extrinsics[..., 3, :] + expected = torch.tensor([0.0, 0.0, 0.0, 1.0], device=extrinsics.device) + if not torch.allclose(last_row, expected.expand_as(last_row), atol=1e-4): + print("Extrinsics wrong row!") + return False + + return True + + + def getitem_prope_DL3DV(self, index): + video = self.videos[index % self.total_length] + if not "Evaluation" in self.base_path: + while not os.path.isdir(os.path.join(self.base_path, video)) or not os.path.exists(os.path.join(self.base_path, video, "transforms.json")): + index += 1 + video = self.videos[index % self.total_length] + + blender2opencv = np.array( + [[1, 0, 0, 0], + [0, -1, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, 1]] + ) + + data = {} + # if "Evaluation" in self.base_path: + # video = f"{video}/{video}/nerfstudio" + with open(os.path.join(self.base_path, video, "transforms.json"), 'r', encoding='utf-8') as file: + json_str = file.read() + meta = json.loads(json_str) + # frames = meta["frames"] + # frames = sorted(frames, key=lambda x: x["colmap_im_id"]) + frames = sorted(os.listdir(os.path.join(self.base_path, video, "images_4"))) + import random + + interval = min(40, len(frames)) + extrinsics_check = False + edit_idx = None + while not extrinsics_check: + extrinsics_check = True + + edit_idx = random.randint(0, len(frames) - interval) + if "val" in self.mode: + edit_idx = 10 + # edit_image = os.path.join(video, frames[edit_idx]["file_path"].replace("images", "images_4")) + edit_image = os.path.join(video, "images_4", frames[edit_idx]) + # print(edit_idx, frames[edit_idx]) + # if "Evaluation" in self.base_path and self.anno_src is not None and "vipe" in self.anno_src: + # vipe_pose = np.load(os.path.join(self.base_path, video.split('/')[0], "pose/video.npz")) + # edit_viewmats = vipe_pose["data"][edit_idx] + if self.anno_src is not None and "vipe-DA3" in self.anno_src: + vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) + edit_viewmats = vipe_pose["data"][edit_idx] + elif self.anno_src is not None and self.anno_src == "vipe": + vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) + edit_viewmats = vipe_pose["data"][edit_idx] + else: + edit_viewmats = np.array(frames[edit_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv # c2w! and invert y z axis! + edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) + + max_idx = min(40, len(frames) - edit_idx) + target_idx = random.randint(edit_idx + max_idx // 2, edit_idx + max_idx - 1) + if "val" in self.mode: + target_idx = 30 + # target_image = os.path.join(video, frames[target_idx]["file_path"].replace("images", "images_4")) + target_image = os.path.join(video, "images_4", frames[target_idx]) + # if "Evaluation" in self.base_path and self.anno_src is not None and "vipe" in self.anno_src: + # vipe_pose = np.load(os.path.join(self.base_path, video.split('/')[0], "pose/video.npz")) + # target_viewmats = vipe_pose["data"][target_idx] + if self.anno_src is not None and "vipe-DA3" in self.anno_src: + vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) + target_viewmats = vipe_pose["data"][target_idx] + elif self.anno_src is not None and self.anno_src == "vipe": + vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) + target_viewmats = vipe_pose["data"][target_idx] + else: + target_viewmats = np.array(frames[target_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv + target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) + + edit_c2w = edit_viewmats + target_c2w = target_viewmats + in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) + + # normalize + c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) + c2ws[:, :3, 3] /= self.norm_scale #20.0 # 10.0 # translation normalized + # print(c2ws) + + if not self.quick_check(c2ws[0:1, :, :]): + extrinsics_check = False + + #transform c2w to w2c align with PRoPE implementation + viewmats = torch.linalg.inv(c2ws) + + + s = 0 + ks = [ [meta["fl_x"], s, meta["cx"]], + [ 0, meta["fl_y"], meta["cy"]], + [ 0, 0, 1]] + ks = torch.Tensor(ks).unsqueeze(0) + image_height = meta["h"] + image_width = meta["w"] + ks[..., 0, 0] = ks[..., 0, 0] / image_width + ks[..., 1, 1] = ks[..., 1, 1] / image_height + ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 + ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 + ks[..., 2, 2] = 1.0 + # ks has been normalized!! + + + data["edit_image"] = self.main_data_operator(edit_image).resize((960, 528)) + data["image"] = self.main_data_operator(target_image).resize((960, 528)) + data["viewmats"] = viewmats + data["Ks"] = torch.cat([ks, ks], dim=0) + + # print("original size : ",data["image"].size, data["edit_image"].size) + if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): + print("!!!camera param has NaN!!!") + print(target_image, edit_image) + exit(0) + + if "qwen" in self.base_model: + data["prompt"] = "镜头视角转到指定位置" + elif "flux" in self.base_model: + data["prompt"] = "Turn to the target view" + data["name"] = video + + # if data["edit_image"].size[0] != 960: #7103edc158a862dbfa3c3454e4de584dad59c3c30055919f1dfa7fd7acfdd5c9 + # print(f"{video} has different size!!") + + if self.model_3D is not None and "val" not in self.mode: + feat_3D = self.model_3D.inference( + [data["edit_image"].resize((960, 528))], # (1, 33, 60) + export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 + process_res=840, + ) + feats = [] + for layer in self.export_3D_feat_layers: + feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) + data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C + if torch.isnan(data["feat_3D"]).any(): + print("!!!feat 3D has NaN!!!") + print(video, edit_idx) + exit(0) + + # target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") + # target_z[np.isnan(target_z)] = 1000 + # target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 + # target_depth = torch.Tensor(target_z) + # data["target_depth"] = target_depth + + if self.add_depth: + if "vipe-DA3" in self.anno_src: + z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") + elif "vipe" == self.anno_src: + z_channel = imread(os.path.join(self.base_path, video, f"depth/{(edit_idx):05d}.exr"), "Z") + z_channel[np.isnan(z_channel)] = 0 + z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 + depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] + # print(torch.max(depth_edit), torch.min(depth_edit)) + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w + if torch.isnan(depth).any(): + print("!!!depth has NaN!!!") + print(video, edit_idx) + exit(0) + + # src_depth = np.array(depth_edit[0]) + # mx = np.max(src_depth) + # mn = np.min(src_depth) + # K = [ [meta["fl_x"] / image_width * 960, 0, 960 // 2], + # [ 0, meta["fl_y"] / image_height * 528, 528 // 2], + # [ 0, 0, 1]] + # T = np.array(viewmats[0]) + # tgt_depth = self.transform_depth( + # src_depth, + # K=np.array(K), + # T=T + # ) + # depth_latent = torch.Tensor(tgt_depth).unsqueeze(0) + # depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w + + data["depth"] = depth + + #if "val" in self.mode: + # tgt_depth = (tgt_depth / np.max(tgt_depth) * 255).astype(np.uint8) + # src_depth = (src_depth / np.max(src_depth) * 255).astype(np.uint8) + + # z_channel = imread(os.path.join(self.base_path, video.split('/')[0], f"depth/{target_idx:05d}.exr"), "Z") + # depth_gt = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) + # depth_gt = F.interpolate(depth_gt, size=(528, 960), mode='bilinear', align_corners=False)[0] + # depth_gt = np.array(depth_gt[0]) + # depth_gt = (depth_gt / np.max(depth_gt) * 255).astype(np.uint8) + + # depth_vis = np.concatenate((src_depth, tgt_depth, depth_gt), axis=1) + # im = Image.fromarray(depth_vis) + # im.save(f"depth_vis/{index}_{mx:.2f}_{mn:.2f}.png") + + + return data + + def getitem_metadata_DL3DV(self, index): + + csv = pd.read_csv(self.metadata_path) + row = csv.iloc[index] + video = str(row['video']) + edit_idx = int(row['edit_idx']) + target_idx = int(row['target_idx']) + + data = {} + + with open(os.path.join(self.base_path, video, "transforms.json"), 'r', encoding='utf-8') as file: + json_str = file.read() + meta = json.loads(json_str) + + frames = sorted(os.listdir(os.path.join(self.base_path, video, "images_4"))) + import random + + edit_image = os.path.join(video, "images_4", frames[edit_idx]) + + if self.anno_src is not None and "vipe-DA3" in self.anno_src: + vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) + edit_viewmats = vipe_pose["data"][edit_idx] + elif self.anno_src is not None and self.anno_src == "vipe": + vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) + edit_viewmats = vipe_pose["data"][edit_idx] + else: + edit_viewmats = np.array(frames[edit_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv # c2w! and invert y z axis! + edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) + + + target_image = os.path.join(video, "images_4", frames[target_idx]) + if self.anno_src is not None and "vipe-DA3" in self.anno_src: + vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) + target_viewmats = vipe_pose["data"][target_idx] + elif self.anno_src is not None and self.anno_src == "vipe": + vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) + target_viewmats = vipe_pose["data"][target_idx] + else: + target_viewmats = np.array(frames[target_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv + target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) + + edit_c2w = edit_viewmats + target_c2w = target_viewmats + in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) + + # normalize + c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) + c2ws[:, :3, 3] /= self.norm_scale #20.0 # 10.0 # translation normalized + + #transform c2w to w2c align with PRoPE implementation + viewmats = torch.linalg.inv(c2ws) + + s = 0 + ks = [ [meta["fl_x"], s, meta["cx"]], + [ 0, meta["fl_y"], meta["cy"]], + [ 0, 0, 1]] + ks = torch.Tensor(ks).unsqueeze(0) + image_height = meta["h"] + image_width = meta["w"] + ks[..., 0, 0] = ks[..., 0, 0] / image_width + ks[..., 1, 1] = ks[..., 1, 1] / image_height + ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 + ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 + ks[..., 2, 2] = 1.0 + # ks has been normalized!! + + + data["edit_image"] = self.main_data_operator(edit_image).resize((960, 528)) + data["image"] = self.main_data_operator(target_image).resize((960, 528)) + data["viewmats"] = viewmats + data["Ks"] = torch.cat([ks, ks], dim=0) + + # print("original size : ",data["image"].size, data["edit_image"].size) + if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): + print("!!!camera param has NaN!!!") + print(target_image, edit_image) + exit(0) + + if "qwen" in self.base_model: + data["prompt"] = "镜头视角转到指定位置" + elif "flux" in self.base_model: + data["prompt"] = "Turn to the target view" + data["name"] = video + + if self.model_3D is not None and "val" not in self.mode: + feat_3D = self.model_3D.inference( + [data["edit_image"].resize((960, 528))], # (1, 33, 60) + export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 + process_res=840, + ) + feats = [] + for layer in self.export_3D_feat_layers: + feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) + data["feat_3D"] = torch.cat(feats, dim=-1) # (1, 20, 36, 1536) B, H, W, C + if torch.isnan(data["feat_3D"]).any(): + print("!!!feat 3D has NaN!!!") + print(video, edit_idx) + exit(0) + + target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") + target_z[np.isnan(target_z)] = 1000 + target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 + target_depth = torch.Tensor(target_z) + data["target_depth"] = target_depth + + if self.add_depth: + if "vipe-DA3" in self.anno_src: + z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") + elif "vipe" == self.anno_src: + z_channel = imread(os.path.join(self.base_path, video, f"depth/{(edit_idx):05d}.exr"), "Z") + z_channel[np.isnan(z_channel)] = 0 + z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 + depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] + # print(torch.max(depth_edit), torch.min(depth_edit)) + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w + if torch.isnan(depth).any(): + print("!!!depth has NaN!!!") + print(video, edit_idx) + exit(0) + + data["depth"] = depth + + return data + + def getitem_metadata(self, index): + csv = pd.read_csv(self.metadata_path) + row = csv.iloc[index] + video = str(row['video']) + edit_idx = int(row['edit_idx']) + target_idx = int(row['target_idx']) + + vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) + vipe_intr = np.load(os.path.join(self.base_path, video, "vipe-DA3/intrinsics/video.npz")) + + data = {} + + cap = cv2.VideoCapture(os.path.join(self.base_path, video, "video.mp4")) + len_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 12 # Re10K shot change + + import random + + interval = min(40, len_frames - 1) + + cap.set(cv2.CAP_PROP_POS_FRAMES, edit_idx) + ret, f = cap.read() + edit_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) + edit_viewmats = vipe_pose["data"][edit_idx] + edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) + + cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx) + ret, f = cap.read() + target_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) + target_viewmats = vipe_pose["data"][target_idx] + target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) + + edit_c2w = edit_viewmats + target_c2w = target_viewmats + print(edit_idx, target_idx, edit_c2w.shape, target_c2w.shape) + in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) + + # normalize + c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) + c2ws[:, :3, 3] /= self.norm_scale # translation normalized + + cap.release() # release video + + #transform c2w to w2c align with PRoPE implementation + viewmats = torch.linalg.inv(c2ws) + + intri = vipe_intr["data"][edit_idx] + s = 0 + ks = [ [intri[0], s, intri[2]], + [ 0, intri[1], intri[3]], + [ 0, 0, 1.]] + image_width = torch.tensor(intri[2]) * 2 + image_height = torch.tensor(intri[3]) * 2 + ks = torch.Tensor(ks).unsqueeze(0) + ks[..., 0, 0] = ks[..., 0, 0] / image_width + ks[..., 1, 1] = ks[..., 1, 1] / image_height + ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 + ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 + + + data["edit_image"] = edit_image.resize((960, 528)) + data["image"] = target_image.resize((960, 528)) + data["viewmats"] = viewmats + data["Ks"] = torch.cat([ks, ks], dim=0) + # print(data["Ks"]) + + if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): + print("!!!camera param has NaN!!!") + print(video, edit_idx, target_idx) + exit(0) + + if "qwen" in self.base_model: + data["prompt"] = "镜头视角转到指定位置" + elif "flux" in self.base_model: + data["prompt"] = "Turn to the target view" + data["name"] = video + + if self.model_3D is not None and "val" not in self.mode: + feat_3D = self.model_3D.inference( + [data["edit_image"].resize((960, 528))], # (1, 33, 60) + export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 + process_res=840, + ) + feats = [] + for layer in self.export_3D_feat_layers: + feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) + data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C + if torch.isnan(data["feat_3D"]).any(): + print("!!!feat 3D has NaN!!!") + print(video, edit_idx) + exit(0) + + if self.add_depth: + z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") + z_channel[np.isnan(z_channel)] = 0 + z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 + depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] + # print(torch.max(depth_edit), torch.min(depth_edit)) + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w + if torch.isnan(depth).any(): + print("!!!depth has NaN!!!") + print(video, edit_idx) + exit(0) + + data["depth"] = depth + + + return data + + def get_frame(self, video_path, frame_num): + cap = cv2.VideoCapture(video_path) + cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num) + ret, frame = cap.read() + cap.release() + if ret: + return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + else: + print(f'error video {video_path} {frame_num}') + exit(0) + + def getitem_prope_all(self, index): + video = self.videos[index % self.total_length] + if not "Evaluation" in self.base_path: + while not os.path.isdir(video): + index += 1 + video = self.videos[index % self.total_length] + + vipe_pose = np.load(os.path.join(video, "vipe-DA3/pose/video.npz")) + vipe_intr = np.load(os.path.join(video, "vipe-DA3/intrinsics/video.npz")) + + data = {} + if 'DL3DV' in video: + frames = sorted(os.listdir(os.path.join(video, "images_4"))) + len_frames = len(frames) + else: + cap = cv2.VideoCapture(os.path.join(video, "video.mp4")) + len_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 12 # Re10K shot change + + + import random + + interval = min(40, len_frames - 1) + extrinsics_check = False + edit_idx = None + while not extrinsics_check: + if "val" in self.mode: + edit_idx = 10 + target_idx = 30 + else: + edit_idx = random.randint(0, len_frames - interval) + max_idx = min(40, len_frames - edit_idx) # 0 -12 + target_idx = random.randint(edit_idx + max_idx // 2, edit_idx + max_idx - 1) + # print(len_frames, edit_idx, target_idx) + + if 'DL3DV' in video: + edit_image = os.path.join(video, "images_4", frames[edit_idx]) + edit_image = Image.open(edit_image) + else: + cap.set(cv2.CAP_PROP_POS_FRAMES, edit_idx) + ret, f = cap.read() + edit_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) + edit_viewmats = vipe_pose["data"][edit_idx] + edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) + + if 'DL3DV' in video: + target_image = os.path.join(video, "images_4", frames[target_idx]) + target_image = Image.open(target_image) + else: + cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx) + ret, f = cap.read() + target_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) + target_viewmats = vipe_pose["data"][target_idx] + target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) + + edit_c2w = edit_viewmats + target_c2w = target_viewmats + in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) + + # normalize + c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) + c2ws[:, :3, 3] /= self.norm_scale # translation normalized + + if self.quick_check(c2ws[0:1, :, :]): + extrinsics_check = True + if 'DL3DV' not in video: + cap.release() # release video + break + + #transform c2w to w2c align with PRoPE implementation + viewmats = torch.linalg.inv(c2ws) + + intri = vipe_intr["data"][edit_idx] + s = 0 + ks = [ [intri[0], s, intri[2]], + [ 0, intri[1], intri[3]], + [ 0, 0, 1.]] + image_width = torch.tensor(intri[2]) * 2 + image_height = torch.tensor(intri[3]) * 2 + ks = torch.Tensor(ks).unsqueeze(0) + ks[..., 0, 0] = ks[..., 0, 0] / image_width + ks[..., 1, 1] = ks[..., 1, 1] / image_height + ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 + ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 + + # ks = [ [meta["fl_x"], s, meta["cx"]], + # [ 0, meta["fl_y"], meta["cy"]], + # [ 0, 0, 1]] + # ks = torch.Tensor(ks).unsqueeze(0) + # image_height = meta["h"] + # image_width = meta["w"] + # ks[..., 0, 0] = ks[..., 0, 0] / image_width + # ks[..., 1, 1] = ks[..., 1, 1] / image_height + # ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 + # ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 + # ks[..., 2, 2] = 1.0 + # ks has been normalized!! + + + data["edit_image"] = edit_image.resize((960, 528)) + data["image"] = target_image.resize((960, 528)) + data["viewmats"] = viewmats + data["Ks"] = torch.cat([ks, ks], dim=0) + # print(data["Ks"]) + + if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): + print("!!!camera param has NaN!!!") + print(video, edit_idx, target_idx) + exit(0) + + if "qwen" in self.base_model: + data["prompt"] = "镜头视角转到指定位置" + elif "flux" in self.base_model: + data["prompt"] = "Turn to the target view" + data["name"] = video + + if self.model_3D is not None and "val" not in self.mode: + feat_3D = self.model_3D.inference( + [data["edit_image"].resize((960, 528))], # (1, 33, 60) + export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 + process_res=840, + ) + feats = [] + for layer in self.export_3D_feat_layers: + feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) + data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C + if torch.isnan(data["feat_3D"]).any(): + print("!!!feat 3D has NaN!!!") + print(video, edit_idx) + exit(0) + + # target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") + # target_z[np.isnan(target_z)] = 1000 + # target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 + # target_depth = torch.Tensor(target_z) + # data["target_depth"] = target_depth + + if self.add_depth: + z_channel = imread(os.path.join(video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") + z_channel[np.isnan(z_channel)] = 0 + z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 + depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) + depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] + # print(torch.max(depth_edit), torch.min(depth_edit)) + depth_latent = torch.zeros_like(depth_edit) + depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w + if torch.isnan(depth).any(): + print("!!!depth has NaN!!!") + print(video, edit_idx) + exit(0) + + data["depth"] = depth + + + return data + + def __len__(self): + if self.prope: + return self.total_length * self.repeat + + +if __name__ == '__main__': + + metadata_path = '../RealEstate10K/meta_view/hard.csv' + base_path = "../RealEstate10K/JiaHWang/Re10K/test" + # 48aaed5a44005bccd51d529ab90335b144fe5e7f3c8a22ba399f4ee3b3fb6728 + dataset = MetaViewUnifiedDataset( + base_path=base_path, + metadata_path=metadata_path, + # subset=['1K', '2K', '3K', '4K', '5K', '6K', '7K', '9K', '10K', '11K', 'train', 'Sekai-Real-Walking-HQ-split'], + repeat=1, + data_file_keys="image,edit_image".split(","), + prope=True, + debug=False, + mode="train", + norm_scale=1.0, + # path_3D="../Depth-Anything-3/model/DA3-GIANT-1.1", + # export_3D_feat_layers="19,39", + anno_src="vipe-DA3", + add_depth=True, + main_data_operator=UnifiedDataset.default_image_operator( + base_path=base_path, + max_pixels=1048576, + height=None, + width=None, + height_division_factor=16, + width_division_factor=16, + ) + ) + + dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=1) + print(len(dataset)) + cnt = 0 + for data in dataloader: + # print(data["name"]) + cnt += 1 + print(data['name']) + #if cnt > 11: + # exit(0) + print(cnt) \ No newline at end of file