bernini/pipeline.py

# Copyright (c) 2026 Bytedance Ltd. and/or its affiliate
#
# 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.

"""End-to-end Bernini Renderer inference pipeline: preprocess -> sample -> decode -> save."""

import html
import json
import logging
import math
import os
import random
import re
from functools import partial
from types import SimpleNamespace
from typing import Optional

import ftfy
import numpy as np
import torch
from diffusers.models import AutoencoderKLWan
from diffusers.video_processor import VideoProcessor
from tqdm import tqdm
from transformers import AutoConfig, AutoTokenizer, AutoProcessor, Qwen2_5_VLModel

from .data.bernini_process import bernini_process_sample
from .data.bernini_template import BerniniTemplate
from .data.utils.video_utils import PathVideoReader, smart_video_nframes
from .data_utils import make_divisible, preprocess_image, preprocess_video, tensor_to_bytes, get_vit_features, get_vae_features, FakeVideoReader, create_fake_image, VAEVideoTransform
from .io_utils import save_output
from .models import BerniniConfig, BerniniModel
from .models import BerniniRendererConfig, BerniniRendererModel
from .models.transformer_wan import WanRotaryPosEmbed
from .weights import load_weights

logger = logging.getLogger("bernini.pipeline")


def _resolve_cached_hf_path(path: Optional[str]) -> Optional[str]:
    """Resolve a HF repo id (or repo id + subpath) to a local cached path.

    Examples:
    - ``ByteDance/Bernini-Diffusers`` -> ``~/.cache/.../snapshots/<rev>``
    - ``ByteDance/Bernini-Diffusers/vae/config.json`` ->
      ``~/.cache/.../snapshots/<rev>/vae/config.json``

    Returns the original ``path`` if it already exists locally or if it cannot
    be resolved from the local HF cache.
    """
    if path is None or os.path.exists(path):
        return path

    parts = path.split("/")
    if len(parts) < 2:
        return path

    repo_id = "/".join(parts[:2])
    subpath = os.path.join(*parts[2:]) if len(parts) > 2 else ""

    try:
        from huggingface_hub import snapshot_download

        repo_dir = snapshot_download(repo_id, local_files_only=True)
    except Exception:
        return path

    candidate = os.path.join(repo_dir, subpath) if subpath else repo_dir
    if os.path.exists(candidate):
        if candidate != path:
            logger.info("resolved hub path '%s' to cached path '%s'", path, candidate)
        return candidate
    return path


def _prefer_local_dir(current, config_dir, *required):
    """Resolve a component base path from config.json against the directory
    passed to from_pretrained.

    Released configs may carry a hub repo id or a path relative to the repo
    root; when that path does not exist locally but `config_dir` contains the
    `required` entries, load from `config_dir` instead so an
    already-downloaded directory is not re-fetched from the Hub.
    """
    config_dir = _resolve_cached_hf_path(config_dir)
    if current is not None and os.path.exists(current):
        return current
    if os.path.isdir(config_dir) and all(
        os.path.exists(os.path.join(config_dir, r)) for r in required if r
    ):
        if current is not None and current != config_dir:
            logger.info(
                "component path '%s' not found locally; loading from '%s' instead",
                current,
                config_dir,
            )
        return config_dir
    current = _resolve_cached_hf_path(current)
    if current is not None and os.path.exists(current):
        return current
    return current


def _prefer_local_file(current, config_dir, *parts):
    """Like _prefer_local_dir, but the local candidate is an entry inside
    `config_dir` rather than the directory itself."""
    config_dir = _resolve_cached_hf_path(config_dir)
    if current is not None and os.path.exists(current):
        return current
    candidate = os.path.join(config_dir, *parts)
    if os.path.exists(candidate):
        return candidate
    current = _resolve_cached_hf_path(current)
    if current is not None and os.path.exists(current):
        return current
    return current


def _localize_bernini_config(config, config_dir):
    """Point the component paths of a BerniniConfig at `config_dir` when the
    paths baked into config.json (repo-root-relative in the released
    Bernini-Diffusers layout) do not exist locally."""
    config_dir = _resolve_cached_hf_path(config_dir)
    config.base_dir = _prefer_local_dir(config.base_dir, config_dir)
    config.diff_dec_config_path = _prefer_local_dir(config.diff_dec_config_path, config.base_dir or config_dir)
    config.mllm_config_path = _prefer_local_dir(
        config.mllm_config_path, config.base_dir or config_dir, config.mllm_subfolder
    )
    config.processor_config_path = _prefer_local_dir(
        config.processor_config_path, config.base_dir or config_dir, config.processor_subfolder
    )
    config.t5_text_encoder_path = _prefer_local_dir(
        config.t5_text_encoder_path, config.base_dir or config_dir, config.t5_text_encoder_subfolder
    )
    config.t5_tokenizer_path = _prefer_local_dir(
        config.t5_tokenizer_path, config.base_dir or config_dir, config.t5_tokenizer_subfolder
    )
    config.vae_model_path = _prefer_local_dir(
        config.vae_model_path, config.base_dir or config_dir, config.vae_subfolder or "vae"
    )
    config.vae_config_path = _prefer_local_file(
        config.vae_config_path, config.base_dir or config_dir, config.vae_subfolder or "vae", "config.json"
    )
    config.transformer_config_path = _prefer_local_file(
        config.transformer_config_path, config.base_dir or config_dir, "transformer_config.json"
    )
    config.transformer_2_config_path = _prefer_local_file(
        config.transformer_2_config_path, config.base_dir or config_dir, "transformer_2_config.json"
    )
    config.scheduler_config_path = _prefer_local_file(
        config.scheduler_config_path, config.base_dir or config_dir, "scheduler"
    )


def _prompt_clean(text: str) -> str:
    text = ftfy.fix_text(text)
    text = html.unescape(html.unescape(text))
    return re.sub(r"\s+", " ", text).strip()


def _vae_encode(vae, x: torch.Tensor) -> torch.Tensor:
    """Encode `[1,C,T,H,W]` pixels into normalized VAE latents."""
    latents = vae.encode(x).latent_dist.mode()
    z = vae.config.z_dim
    mean = torch.tensor(vae.config.latents_mean, dtype=latents.dtype, device=latents.device).view(1, z, 1, 1, 1)
    std = torch.tensor(vae.config.latents_std, dtype=latents.dtype, device=latents.device).view(1, z, 1, 1, 1)
    return (latents - mean) / std


def _get_t5_text_ids(text, tokenizer, max_length: int = 512):
    """Tokenize text for the T5 encoder, returning input_ids and attention_mask."""
    text = _prompt_clean(text)
    out = tokenizer(
        text,
        padding="max_length",
        max_length=max_length,
        truncation=True,
        add_special_tokens=True,
        return_attention_mask=True,
        return_tensors="pt",
    )
    return out.input_ids, out.attention_mask


def _vae_decode(vae, latents: torch.Tensor):
    """Decode VAE latents into a numpy clip `[T, H, W, C]` in [0, 1]."""
    latents = latents.to(vae.dtype)
    z = vae.config.z_dim
    mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype).view(1, z, 1, 1, 1)
    std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype).view(1, z, 1, 1, 1)
    latents = latents * std + mean
    video = vae.decode(latents, return_dict=False)[0]
    processor = VideoProcessor(vae_scale_factor=2 ** len(vae.temperal_downsample))
    return processor.postprocess_video(video, output_type="np")[0]


class BerniniRendererPipeline:
    """Loads the model once; each call generates one video / image."""

    def __init__(self, model, vae, tokenizer, device):
        self.model = model
        self.vae = vae
        self.tokenizer = tokenizer
        self.device = device
        self.weight_dtype = torch.bfloat16

    @classmethod
    def from_pretrained(
        cls,
        config_dir: str,
        high_noise_ckpt: str = None,
        low_noise_ckpt: str = None,
        device="cuda",
        load_ckpt_weights: bool = True,
        **config_overrides,
    ) -> "BerniniRendererPipeline":
        config = BerniniRendererConfig.from_pretrained(config_dir, **config_overrides)
        config.wan22_base = _prefer_local_dir(
            config.wan22_base, config_dir, "tokenizer", "text_encoder", "vae"
        )
        tokenizer = AutoTokenizer.from_pretrained(
            config.wan22_base, subfolder="tokenizer", trust_remote_code=True
        )
        vae = AutoencoderKLWan.from_pretrained(config.wan22_base, subfolder="vae", torch_dtype=torch.float32)
        vae.eval()
        vae.requires_grad_(False)
        model = BerniniRendererModel(config)
        if load_ckpt_weights:
            load_weights(model, high_noise_ckpt, low_noise_ckpt)
        model.eval()
        return cls(model, vae, tokenizer, device)

    def _tokenize(self, prompt: str):
        out = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=512,
            truncation=True,
            add_special_tokens=True,
            return_attention_mask=True,
            return_tensors="pt",
        )
        return out.input_ids, out.attention_mask

    @torch.no_grad()
    def __call__(
        self,
        prompt: str,
        *,
        neg_prompt: str = "",
        num_frames: int = 81,
        max_image_size: int = 624,
        height: int = 480,
        width: int = 832,
        video=None,
        image=None,
        images=None,
        num_inference_steps: int = 40,
        guidance_mode: str = "rv2v",
        omega_vid: float = 3.0,
        omega_img: float = 3.0,
        omega_txt: float = 4.0,
        omega_scale: float = 0.75,
        flow_shift: float = 5.0,
        seed: int = 42,
        fps: int = 16,
        vae_fps: int = None,
        vit_fps: int = None,
        eta: float = 0.5,
        norm_threshold=(50.0, 50.0),
        momentum: float = -0.5,
        system_prompt: str = "",
        output_path: str = "output.mp4",
        write_output: bool = True,
        **kwargs
    ):
        """Generate one clip and write it to `output_path`.

        `video` drives video editing, `image` a single-image edit, `images` a
        list of reference images; the output size follows the source video or
        single image, otherwise `height` / `width`.

        With `write_output=False` the decode/save step is skipped (used by the
        redundant ranks of an Ulysses group) and ``None`` is returned.
        """
        device = self.device
        prompt = system_prompt + _prompt_clean(prompt)
        logger.info("prompt: %s", prompt)
        prompt_ids, prompt_mask = self._tokenize(prompt)
        neg_ids, neg_mask = self._tokenize(neg_prompt)

        # ---- encode visual conditions on the VAE ----
        self.vae.to(device)
        t, h, w = num_frames, None, None

        multi_video_vae_latents = None
        if video is not None:
            paths = video if isinstance(video, list) else [video]
            multi_video_vae_latents = []
            first_shape = None
            for vp in paths:
                pv = preprocess_video(
                    vp, fps=fps, max_image_size=max_image_size, max_image_num=num_frames, device=device
                )
                if first_shape is None:
                    first_shape = pv.shape
                multi_video_vae_latents.append(_vae_encode(self.vae, pv))
            t, h, w = first_shape[-3], first_shape[-2], first_shape[-1]

        image_vae_latents = None
        if image is not None:
            pi = preprocess_image(image, max_image_size=max_image_size, device=device)
            if h is None:
                h, w = pi.shape[-2], pi.shape[-1]
            image_vae_latents = _vae_encode(self.vae, pi)

        multi_image_vae_latents = None
        if images:
            multi_image_vae_latents = [
                _vae_encode(self.vae, preprocess_image(img, max_image_size=max_image_size, device=device))
                for img in images
            ]

        self.vae.to("cpu")
        torch.cuda.empty_cache()

        if h is None:
            h, w = height, width
        h, w = make_divisible(h, 16), make_divisible(w, 16)

        # ---- diffusion sampling ----
        latents = self.model.sample(
            input_ids=prompt_ids.to(device),
            attention_mask=prompt_mask.to(device),
            uncond_input_ids=neg_ids.to(device),
            uncond_attention_mask=neg_mask.to(device),
            image_vae_latents=image_vae_latents,
            multi_video_vae_latents=multi_video_vae_latents,
            multi_image_vae_latents=multi_image_vae_latents,
            num_frames=t,
            width=w,
            height=h,
            num_inference_steps=num_inference_steps,
            guidance_mode=guidance_mode,
            omega_vid=omega_vid,
            omega_img=omega_img,
            omega_txt=omega_txt,
            omega_scale=omega_scale,
            flow_shift=flow_shift,
            seed=seed,
            device=device,
            eta=eta,
            norm_threshold=norm_threshold,
            momentum=momentum,
        )
        self.model.to("cpu")
        torch.cuda.empty_cache()

        if not write_output:
            return None

        # ---- decode + save ----
        self.vae.to(device)
        output = _vae_decode(self.vae, latents)
        self.vae.to("cpu")
        torch.cuda.empty_cache()

        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        save_output(output, output_path, fps=fps)
        logger.info("saved -> %s  (%d frames, %dx%d)", output_path, output.shape[0], h, w)
        return output_path

class BerniniPipeline:
    """Loads the model once; each call generates one video / image."""

    def __init__(self, config, model, vae, t5_tokenizer, vit_processor, device):
        self.config = config
        self.model = model
        self.vae = vae
        self.t5_tokenizer = t5_tokenizer
        self.vit_processor = vit_processor
        self.device = device
        self.weight_dtype = torch.bfloat16
        self.text_encoder = model.mllm
        self.connector = getattr(model, "connector", None)

    @classmethod
    def from_pretrained(
        cls,
        config_dir: str,
        ckpt: str = None,
        device="cuda",
        **config_overrides,
    ) -> "BerniniPipeline":
        config = BerniniConfig.from_pretrained(config_dir, **config_overrides)
        _localize_bernini_config(config, config_dir)
        if ckpt is None: ckpt = config_dir
        model = BerniniModel.from_pretrained(
            ckpt,
            subfolder=config.bernini_ckpt_subfolder,
            config=config,
        )
        # transformer_1 is loaded in diff_dec, while transformer_2 is loaded in diff_dec_low and then
        # attached back to diff_dec before sampling.
        setattr(model.diff_dec, "transformer_2", model.diff_dec_low.transformer_2)
        model.eval()

        t5_tokenizer = AutoTokenizer.from_pretrained(
            config.t5_tokenizer_path,
            subfolder=config.t5_tokenizer_subfolder,
            trust_remote_code=True,
        )
        vit_processor = AutoProcessor.from_pretrained(
            config.processor_config_path,
            subfolder=config.processor_subfolder,
            padding_side="right",
            trust_remote_code=True,
        )

        vae = AutoencoderKLWan.from_pretrained(
            config.vae_model_path,
            subfolder=config.vae_subfolder,
            torch_dtype=torch.float32,
        )
        vae.eval()
        vae.requires_grad_(False)
        return cls(config, model, vae, t5_tokenizer, vit_processor, device)

    @torch.no_grad()
    def sample_vit_decoder(
        self,
        vit_embed,
        uncond_vit_embed,
        imgcond_vit_embed,
        vit_txt_cfg,
        sample_steps,
        vit_img_cfg=None,
        verbose=True,
    ):
        dtype = vit_embed.dtype
        if vit_img_cfg is not None and vit_txt_cfg > 1.0:
            vit_embed = torch.cat([vit_embed, uncond_vit_embed, imgcond_vit_embed], dim=1)
        elif vit_txt_cfg > 1.0:
            vit_embed = torch.cat([vit_embed, uncond_vit_embed], dim=1)

        vit_embed = (
            self.model.vit_decoder.sample(
                z=vit_embed[0],
                cfg=vit_txt_cfg,
                img_cfg=vit_img_cfg,
                num_inference_steps=sample_steps,
                verbose=verbose,
            )
            .unsqueeze(0)
            .to(dtype)
        )

        if vit_img_cfg is not None and vit_txt_cfg > 1.0:
            vit_embed = vit_embed[:, : vit_embed.shape[1] // 3, :]
        elif vit_txt_cfg > 1.0:
            vit_embed = vit_embed[:, : vit_embed.shape[1] // 2, :]

        return vit_embed

    @torch.no_grad()
    def preprocess_inputs(
        self,
        prompt,
        mllm_model,
        vae_model,
        vae_transform,
        row=None,
        vit_min_pixels: int = 3136,
        vit_max_pixels: int = 50176,
        vae_fps: int = 16,
        vit_fps: int = 2,
        num_frames: int = 81,
        height: int = 480,
        width: int = 832,
        video=None,
        image=None,
        images=None,
        max_duration: int = None,
    ):
        from bernini.data_utils import generate_unified_inputs
        # Build image/video path lists, filtering out None values.
        if images is not None:
            input_image_paths = [img for img in images if img is not None]
        elif image is not None:
            input_image_paths = [image]
        else:
            input_image_paths = []
        if video is None:
            input_video_paths = []
        elif isinstance(video, str):
            input_video_paths = [video]
        else:
            input_video_paths = [vid for vid in video if vid is not None]
        inputs_structure = generate_unified_inputs(
            prompt,
            input_image_paths=input_image_paths,
            input_video_paths=input_video_paths,
            has_video_input=bool(input_video_paths),
            output_t=num_frames,
            output_h=height,
            output_w=width,
        )
        row_data = {}
        row_data['inputs'] = inputs_structure
        if (images is not None and len(images) > 0) or image is not None or num_frames == 1:
            images = [image] if image is not None else list(images or [])
            if num_frames == 1: images.append("output_img_flag")
            images = [create_fake_image(height, width) if img == "output_img_flag" else img for img in images]
            image_inputs = self.vit_processor.image_processor(
                images=images, return_tensors="pt",
                min_pixels=vit_min_pixels, max_pixels=vit_max_pixels
            )
            pixel_values = image_inputs['pixel_values']
            image_grid_thw = image_inputs['image_grid_thw']
            image_embeds = get_vit_features(mllm_model, pixel_values, image_grid_thw)
            row_data['image_embeds'] = [tensor_to_bytes(embed.detach().cpu()) for embed in image_embeds]
            row_data['image_grid_thw'] = image_grid_thw.numpy().tolist()
            # VAE
            image_tensors = [vae_transform(img) for img in images]
            image_vae_latents = []
            for img_tensor in image_tensors:
                latent = get_vae_features(vae_model, img_tensor)
                image_vae_latents.append(latent)
            row_data['image_vae_latents'] = image_vae_latents
            del image_tensors
            torch.cuda.empty_cache()
        
        if video is not None or num_frames > 1:
            row_data["video_embeds"] = []
            row_data["video_grid_thw"] = []
            row_data["video_vae_latents"] = []
            video_meta = []
            if video is not None: 
                if isinstance(video, str):
                    video_meta.append(video)
                elif isinstance(video, list):
                    video_meta.extend(video)
            if num_frames > 1:
                if video is not None:
                    video_meta.append(video_meta[0])
                else:
                    video_meta.append('output_vid_flag')
            for video_path in video_meta:
                if video_path != "output_vid_flag":
                    duration = None
                    if row is not None and 'videos' in row:
                        for v_meta in row['videos']:
                            if v_meta.get('video_path') == video_path:
                                duration = v_meta.get('duration', None)
                                break
                    if duration is not None and max_duration is not None and duration > max_duration:
                        duration = max_duration
                    video_reader = PathVideoReader(video_path, duration=duration, crop_method='left')
                else:
                    video_reader = FakeVideoReader(
                        num_frames=num_frames,
                        height=height,
                        width=width,
                        fps=vae_fps
                    )
                vit_idx = smart_video_nframes(
                    total_frames=video_reader.length, video_fps=video_reader.fps,
                    fps=vit_fps, frame_factor=2,
                    max_frames=num_frames, add_one=False
                )
                video_for_vit = video_reader.sample(vit_idx)
                video_inputs = self.vit_processor.video_processor(
                    videos=video_for_vit, return_tensors="pt",
                    size={'shortest_edge': vit_min_pixels, 'longest_edge': vit_max_pixels},
                )
                vid_pixel_values = video_inputs['pixel_values_videos']
                vid_grid_thw = video_inputs['video_grid_thw']
                video_embeds = get_vit_features(mllm_model, vid_pixel_values, vid_grid_thw)
                row_data['video_embeds'].extend([tensor_to_bytes(embed.detach().cpu()) for embed in video_embeds])
                row_data['video_grid_thw'].extend(vid_grid_thw.numpy().tolist())
                del video_inputs
                vae_idx = smart_video_nframes(
                    total_frames=video_reader.length, video_fps=video_reader.fps,
                    fps=vae_fps, frame_factor=4,
                    max_frames=num_frames, add_one=True
                )
                video_for_vae = video_reader.sample(vae_idx)
                video_tensor = torch.stack([vae_transform(frame) for frame in video_for_vae], dim=1)
                video_vae_latent = get_vae_features(vae_model, video_tensor)
                row_data['video_vae_latents'].append(video_vae_latent)
                del video_tensor

                torch.cuda.empty_cache()
        
        return row_data

    def transform_inputs(
        self,
        sample,
        max_vae_frames: int = 81,
        task_name: str = "t2v",
        neg_prompt: Optional[str] = None,
        t5_neg_prompt: str = "",
        use_qwen_neg_prompt: bool = True,
    ):
        if neg_prompt is not None:
            t5_neg_prompt = neg_prompt
        rope = WanRotaryPosEmbed(
            128,
            (1, 2, 2),
            1024,
            use_src_id_rotary_emb=True,
        )
        mllm_config_path = self.config.mllm_config_path
        mllm_config_subfolder = getattr(self.config, "mllm_subfolder", None)
        vae_model_path = getattr(self.config, "vae_model_path", None)
        vae_subfolder = getattr(self.config, "vae_subfolder", None)
        vae_config_path = self.config.vae_config_path
        
        mllm_config = AutoConfig.from_pretrained(mllm_config_path, subfolder=mllm_config_subfolder)
        fake_model = SimpleNamespace(
            config=mllm_config,
            image_token_id=mllm_config.image_token_id,
            video_token_id=mllm_config.video_token_id,
        )
        position_id_func = partial(Qwen2_5_VLModel.get_rope_index, fake_model)
        
        processor = self.vit_processor
        chat_template = BerniniTemplate(
            processor.tokenizer,
            t5_tokenizer=self.t5_tokenizer
        )

        with open(vae_config_path, 'r') as f:
            vae_config = json.load(f)
        vae_latent_mean = torch.tensor(vae_config['latents_mean'], device="cpu")
        vae_latent_std = torch.tensor(vae_config['latents_std'], device="cpu")
        vae_latent_mean = vae_latent_mean.view(vae_config['z_dim'], 1, 1, 1)
        vae_latent_std = vae_latent_std.view(vae_config['z_dim'], 1, 1, 1)

        src_id_kwargs = dict(
            interpolate_src_id=getattr(self.config, "interpolate_src_id", True),
            max_trained_src_id=getattr(self.config, "max_trained_src_id", 5),
        )
        transform = partial(
            bernini_process_sample,
            processor=processor,
            chat_template=chat_template,
            position_id_func=position_id_func,
            vae_rope_func=rope,
            vae_latent_mean=vae_latent_mean,
            vae_latent_std=vae_latent_std,
            text_dropout_rate=0.0,
            img_dropout_rate=0.0,
            video_dropout_rate=0.0,
            max_vae_frames=max_vae_frames,
            source_name=task_name,
            **src_id_kwargs,
        )
        uncond_transform = partial(
            bernini_process_sample,
            processor=processor,
            chat_template=chat_template,
            position_id_func=position_id_func,
            vae_rope_func=rope,
            vae_latent_mean=vae_latent_mean,
            vae_latent_std=vae_latent_std,
            text_dropout_rate=1.0,
            img_dropout_rate=1.0,
            video_dropout_rate=1.0,
            max_vae_frames=max_vae_frames,
            source_name=task_name,
            **src_id_kwargs,
        )
        imgcond_transform = partial(
            bernini_process_sample,
            processor=processor,
            chat_template=chat_template,
            position_id_func=position_id_func,
            vae_rope_func=rope,
            vae_latent_mean=vae_latent_mean,
            vae_latent_std=vae_latent_std,
            text_dropout_rate=0.0,
            img_dropout_rate=1.0,
            video_dropout_rate=1.0,
            max_vae_frames=max_vae_frames,
            source_name=task_name,
            **src_id_kwargs,
        )
        def process_sample(sample, sample_idx):
            neg_prompt = sample.get('neg_prompt', t5_neg_prompt)
            tokenized_example = transform(sample)[0]
            imgcond_tokenized_example = imgcond_transform(sample)[0]
            uncond_tokenized_example = uncond_transform(sample, neg_prompt=neg_prompt if use_qwen_neg_prompt else "")[0]

            for k in tokenized_example:
                if isinstance(tokenized_example[k], torch.Tensor):
                    tokenized_example[k] = tokenized_example[k]
                    uncond_tokenized_example[k] = uncond_tokenized_example[k]
                    imgcond_tokenized_example[k] = imgcond_tokenized_example[k]
            
            sample.pop('inputs')
            return dict(
                uid=sample.pop('uid', f'{sample_idx:03d}').split('/')[-1].split('.')[0],
                edit_type=sample.pop('edit_type', 'unknown'),
                inputs=tokenized_example,
                uncond_inputs=uncond_tokenized_example,
                imgcond_inputs=imgcond_tokenized_example,
                **sample
            )
        
        return process_sample(sample, 0)

    @torch.no_grad()
    def sample_vit_embed(
        self,
        input_embeds: Optional[torch.FloatTensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        attention_mask_4d: Optional[torch.Tensor] = None,
        visual_output_token_mask=None,
        uncond_input_embeds: Optional[torch.FloatTensor] = None,
        uncond_position_ids: Optional[torch.Tensor] = None,
        uncond_attention_mask_4d: Optional[torch.Tensor] = None,
        uncond_visual_output_token_mask=None,
        imgcond_input_embeds: Optional[torch.FloatTensor] = None,
        imgcond_position_ids: Optional[torch.Tensor] = None,
        imgcond_attention_mask_4d: Optional[torch.Tensor] = None,
        imgcond_visual_output_token_mask=None,
        planning_step=3,
        vit_denoising_step=1,
        vit_txt_cfg=1.4,
        vit_img_cfg=1.2,
    ):  
        device = input_embeds.device
        mask_ratio_generator_infer = lambda s, totals: np.cos(math.pi / 2.0 * (s + 1) / totals)
        # Init and sample generation orders
        n_query_tokens = visual_output_token_mask.sum().detach().cpu().numpy()
        order = np.array(list(range(n_query_tokens)))
        np.random.shuffle(order)
        order = torch.Tensor(np.array(order)).to(device).long()
        mask = torch.ones(n_query_tokens).to(device)

        if position_ids.shape[1] == 3:
            position_ids = position_ids.transpose(
                0, 1).contiguous()  # bs, dim, l -> dim, bs, l
        if uncond_position_ids.shape[1] == 3:
            uncond_position_ids = uncond_position_ids.transpose(
                0, 1).contiguous()  # bs, dim, l -> dim, bs, l
        if imgcond_position_ids.shape[1] == 3:
            imgcond_position_ids = imgcond_position_ids.transpose(
                0, 1).contiguous()  # bs, dim, l -> dim, bs, l

        if self.model.vit_decoder is not None:
            for step in tqdm(range(planning_step), desc=f"Sample FM+MAR clip in {planning_step} steps"):
                if self.connector is not None:
                    connector_param = next(self.connector.parameters())
                    if connector_param.device != input_embeds.device or connector_param.dtype != input_embeds.dtype:
                        self.connector.to(device=input_embeds.device, dtype=input_embeds.dtype)
                hidden_state = self.text_encoder(
                    inputs_embeds=input_embeds.clone(),
                    position_ids=position_ids.clone(),
                    attention_mask=attention_mask_4d.clone(),
                    output_hidden_states=True,
                ).hidden_states[-2]
                uncond_hidden_state = self.text_encoder(
                    inputs_embeds=uncond_input_embeds.clone(),
                    position_ids=uncond_position_ids.clone(),
                    attention_mask=uncond_attention_mask_4d.clone(),
                    output_hidden_states=True,
                ).hidden_states[-2]
                imgcond_hidden_state = self.text_encoder(
                    inputs_embeds=imgcond_input_embeds.clone(),
                    position_ids=imgcond_position_ids.clone(),
                    attention_mask=imgcond_attention_mask_4d.clone(),
                    output_hidden_states=True,
                ).hidden_states[-2]
                # extract feat from stageone to feed stagetwo
                cond_vit_embed = hidden_state[:, visual_output_token_mask, :]
                uncond_vit_embed = uncond_hidden_state[:, uncond_visual_output_token_mask, :]
                imgcond_vit_embed = imgcond_hidden_state[:, imgcond_visual_output_token_mask, :]
                pred_vit_embed_mllm = self.connector.for_vit(cond_vit_embed)
                uncond_pred_vit_embed_mllm = self.connector.for_vit(uncond_vit_embed)
                imgcond_pred_vit_embed_mllm = self.connector.for_vit(imgcond_vit_embed)

                # mask ratio for the next round, following MaskGIT and MAGE.
                mask_ratio = mask_ratio_generator_infer(step, planning_step)
                mask_len = torch.Tensor([np.floor(n_query_tokens * mask_ratio)]).to(device)
                # masks out at least one for the next iteration
                mask_len = torch.maximum(
                    torch.Tensor([1]).cuda(), torch.minimum(torch.sum(mask, dim=-1, keepdims=True) - 1, mask_len)
                )
                # get masking for next iteration
                mask_next = torch.zeros_like(mask)  # zero init
                mask_next = torch.scatter(
                    mask_next,
                    dim=-1,
                    index=order[: mask_len.long()],  
                    src=torch.ones_like(mask),
                ).bool()
                if step >= planning_step - 1:
                    mask_to_pred = mask.bool()  # Predict the left mask tokens
                else:
                    mask_to_pred = torch.logical_xor(mask.bool(), mask_next)
                mask = mask_next

                if mask_to_pred.nonzero(as_tuple=True)[0].sum() == 0:
                    continue 
                cond_pred_vit_embed = pred_vit_embed_mllm[:, mask_to_pred.nonzero(as_tuple=True)[0]]
                uncond_pred_vit_embed = uncond_pred_vit_embed_mllm[:, mask_to_pred.nonzero(as_tuple=True)[0]]
                imgcond_pred_vit_embed = imgcond_pred_vit_embed_mllm[:, mask_to_pred.nonzero(as_tuple=True)[0]]
                cur_pred_vit_embed = self.sample_vit_decoder(
                    vit_embed=cond_pred_vit_embed,
                    uncond_vit_embed=uncond_pred_vit_embed,
                    imgcond_vit_embed=imgcond_pred_vit_embed,
                    vit_txt_cfg=vit_txt_cfg,
                    vit_img_cfg=vit_img_cfg,
                    sample_steps=vit_denoising_step,
                    verbose=False,
                )

                all_target_vit_embed = input_embeds[:, visual_output_token_mask, :]
                all_target_vit_embed[:, mask_to_pred.nonzero(as_tuple=True)[0]] = cur_pred_vit_embed
                input_embeds[:, visual_output_token_mask] = all_target_vit_embed
                uncond_input_embeds[:, uncond_visual_output_token_mask] = all_target_vit_embed
                imgcond_input_embeds[:, imgcond_visual_output_token_mask] = all_target_vit_embed

        pred_vit_embed_diff = input_embeds[:, visual_output_token_mask, :]

        outputs = self.text_encoder(
            inputs_embeds=input_embeds.clone(),
            position_ids=position_ids.clone(),
            attention_mask=attention_mask_4d.clone(),
            output_hidden_states=True,
        )
        uncond_outputs = self.text_encoder(
            inputs_embeds=uncond_input_embeds.clone(),
            position_ids=uncond_position_ids.clone(),
            attention_mask=uncond_attention_mask_4d.clone(),
            output_hidden_states=True,
        )
        cond_outputs = self.model.feat_from_planner_to_renderer(
            hidden_states=outputs.hidden_states[-2],
            visual_output_mask=visual_output_token_mask, 
            tgt_vit_mask=None,
            inference=True
        )
        uncond_outputs = self.model.feat_from_planner_to_renderer(
            hidden_states=uncond_outputs.hidden_states[-2],
            visual_output_mask=uncond_visual_output_token_mask, 
            tgt_vit_mask=None,
            inference=True
        )

        if self.model.feature_type_from_stage_one in ["masked_tgt_embed_with_qwen_txt_tokens"]:
            cond_embeds_wotxt_wovit = uncond_outputs['diff_mllm_contexts']
            cond_embeds_wtxt_wvit = cond_outputs['diff_mllm_contexts']
            cond_embeds_wtxt_wovit = None
            cond_embeds_wotxt_wvit = None
        else:
            uncond_cond_embeds = uncond_outputs['diff_mllm_contexts']
            diff_mllm_context_txt_mask = uncond_outputs['diff_mllm_context_txt_mask']
            cond_embeds_wotxt_wovit = uncond_cond_embeds[:, diff_mllm_context_txt_mask] 
            diff_mllm_context_txt_mask = cond_outputs['diff_mllm_context_txt_mask']
            diff_mllm_context_vit_mask = cond_outputs['diff_mllm_context_vit_mask']
            cond_embeds_wtxt_wvit = cond_outputs['diff_mllm_contexts']
            cond_embeds_wtxt_wovit = cond_embeds_wtxt_wvit[:, diff_mllm_context_txt_mask]
            cond_embeds_wotxt_wvit = cond_embeds_wtxt_wvit[:, diff_mllm_context_vit_mask]

        return dict(
            cond_embeds_wtxt_wvit=cond_embeds_wtxt_wvit,
            cond_embeds_wtxt_wovit=cond_embeds_wtxt_wovit,
            cond_embeds_wotxt_wvit=cond_embeds_wotxt_wvit,
            cond_embeds_wotxt_wovit=cond_embeds_wotxt_wovit,
            pred_vit_embed=pred_vit_embed_diff
        )

    @torch.no_grad()
    def __call__(
        self,
        task_name: str,
        prompt: str,
        *,
        neg_prompt: str = "",
        num_frames: int = 81,
        max_image_size: int = 624,
        height: int = 480,
        width: int = 832,
        video=None,
        image=None,
        images=None,
        num_inference_steps: int = 40,
        guidance_mode: str = "rv2v",
        omega_vid: float = 3.0, 
        omega_img: float = 3.0, 
        omega_txt: float = 4.0, 
        omega_tgt: float = 4.0, 
        omega_scale: float = 0.75,
        planning_step: int = 25,
        vit_txt_cfg: float = 1.4,
        vit_img_cfg: float = 1.2,
        vit_denoising_step: int = 3,
        flow_shift: float = 5.0,
        seed: int = 42,
        fps: int = 16,
        eta: float = 0.5,
        norm_threshold=(50.0, 50.0),
        momentum: float = -0.5,
        system_prompt: str = "",
        output_path: str = "output.mp4",
        write_output: bool = True,
        use_truncate: bool = False,
        max_sequence_length: int = 512,
    ):
        """Generate one clip and write it to `output_path`.

        `video` drives video editing, `image` a single-image edit, `images` a
        list of reference images; the output size follows the source video or
        single image, otherwise `height` / `width`.

        With `write_output=False` the decode/save step is skipped (used by the
        redundant ranks of an Ulysses group) and ``None`` is returned.
        """
        device = self.device
        # Resets the torch RNG from the request seed before feature extraction and sampling.
        random.seed(int(seed))
        np.random.seed(int(seed))
        torch.manual_seed(int(seed))
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(int(seed))
        # Resolve fps aliases: fps sets default for both vae_fps and vit_fps
        vae_fps = fps
        vit_fps = fps // 8  # default vit_fps is 1/8 of vae_fps
        raw_prompt = _prompt_clean(prompt)
        t5_prompt = _prompt_clean(system_prompt + raw_prompt)
        logger.info("prompt: %s", t5_prompt)
        # ---- encode visual conditions on the VAE ----
        self.vae.to(device)
        self.model.mllm.to(device)
        self.model.mllm.to(self.weight_dtype)
        if self.connector is not None:
            self.connector.to(device=device, dtype=self.weight_dtype)
        if getattr(self.model, "vit_decoder", None) is not None:
            self.model.vit_decoder.to(device=device, dtype=self.weight_dtype)
        vae_transform = VAEVideoTransform(
            max_image_size=max_image_size,
            min_image_size=240,
            image_stride=16,
        )
        sample = self.preprocess_inputs(
            raw_prompt,
            mllm_model=self.model.mllm,
            vae_model=self.vae,
            vae_transform=vae_transform,
            num_frames=num_frames,
            height=height,
            width=width,
            video=video,
            image=image,
            images=images,
            vit_fps=vit_fps,
            vae_fps=vae_fps,
        )
        self.vae.to("cpu")
        torch.cuda.empty_cache()
        input_dict = self.transform_inputs(
            sample,
            num_frames,
            task_name=task_name,
            neg_prompt=neg_prompt,
        )

        def _move_to_device(obj):
            if isinstance(obj, torch.Tensor):
                return obj.to(device)
            if isinstance(obj, dict):
                return {k: _move_to_device(v) for k, v in obj.items()}
            if isinstance(obj, list):
                return [_move_to_device(v) for v in obj]
            if isinstance(obj, tuple):
                return tuple(_move_to_device(v) for v in obj)
            return obj

        input_dict = _move_to_device(input_dict)
        inputs = input_dict['inputs']
        uncond_inputs = input_dict['uncond_inputs']
        imgcond_inputs = input_dict['imgcond_inputs']

        input_embeds = self.model.format_mllm_inputs_embeds(
            input_ids=inputs['input_ids'],
            visual_embeds=inputs['visual_embeds'],
            visual_input_mask=inputs['visual_input_token_mask'],
            visual_output_mask=inputs['visual_output_token_mask'],
        ).to(self.weight_dtype)
        uncond_input_embeds = self.model.format_mllm_inputs_embeds(
            input_ids=uncond_inputs['input_ids'],
            visual_embeds=uncond_inputs['visual_embeds'],
            visual_input_mask=uncond_inputs['visual_input_token_mask'],
            visual_output_mask=uncond_inputs['visual_output_token_mask'],
        ).to(self.weight_dtype)
        imgcond_input_embeds = self.model.format_mllm_inputs_embeds(
            input_ids=imgcond_inputs['input_ids'],
            visual_embeds=imgcond_inputs['visual_embeds'],
            visual_input_mask=imgcond_inputs['visual_input_token_mask'],
            visual_output_mask=imgcond_inputs['visual_output_token_mask'],
        ).to(self.weight_dtype)

        post_process_out = self.model.post_process_input_embeds(
                input_embeds.unsqueeze(0), 
                inputs['visual_output_token_mask'], 
                tgt_vit_mask=None,
                inference=True
            )
        inputs_embed = post_process_out['input_embeds']
        uncond_post_process_out = self.model.post_process_input_embeds(
                uncond_input_embeds.unsqueeze(0), 
                uncond_inputs['visual_output_token_mask'], 
                tgt_vit_mask=None,
                inference=True
            )
        uncond_inputs_embed = uncond_post_process_out['input_embeds']
        imgcond_post_process_out = self.model.post_process_input_embeds(
                imgcond_input_embeds.unsqueeze(0), 
                imgcond_inputs['visual_output_token_mask'], 
                tgt_vit_mask=None,
                inference=True
            )
        imgcond_inputs_embed = imgcond_post_process_out['input_embeds']
        ret = self.sample_vit_embed(
            input_embeds=inputs_embed,
            attention_mask_4d=inputs['attention_mask_4d'].unsqueeze(0),
            position_ids=inputs['position_ids'].unsqueeze(0),
            visual_output_token_mask=inputs['visual_output_token_mask'],
            uncond_input_embeds=uncond_inputs_embed,
            uncond_position_ids=uncond_inputs['position_ids'].unsqueeze(0),
            uncond_attention_mask_4d=uncond_inputs['attention_mask_4d'].unsqueeze(0),
            uncond_visual_output_token_mask=uncond_inputs['visual_output_token_mask'],
            imgcond_input_embeds=imgcond_inputs_embed,
            imgcond_position_ids=imgcond_inputs['position_ids'].unsqueeze(0),
            imgcond_attention_mask_4d=imgcond_inputs['attention_mask_4d'].unsqueeze(0),
            imgcond_visual_output_token_mask=imgcond_inputs['visual_output_token_mask'],
            planning_step=planning_step,
            vit_txt_cfg=vit_txt_cfg,
            vit_img_cfg=vit_img_cfg,
            vit_denoising_step=vit_denoising_step,
        )
        
        cond_embeds_wtxt_wvit = ret['cond_embeds_wtxt_wvit']
        cond_embeds_wtxt_wovit = ret['cond_embeds_wtxt_wovit']
        cond_embeds_wotxt_wvit = ret['cond_embeds_wotxt_wvit']
        cond_embeds_wotxt_wovit = ret['cond_embeds_wotxt_wovit']

        self.model.mllm.to('cpu')
        if self.connector is not None:
            self.connector.to('cpu')
        if getattr(self.model, "vit_decoder", None) is not None:
            self.model.vit_decoder.to('cpu')
        torch.cuda.empty_cache()

        if getattr(self.model, "t5_text_encoder", None) is not None:
            self.model.t5_text_encoder.to(device)
        t5_input_ids, t5_attention_mask = _get_t5_text_ids(
            t5_prompt, self.t5_tokenizer,
        )
        t5_embeds = self.model.get_t5_text_embeddings_sample(
            t5_input_ids.to(device), t5_attention_mask.to(device)
        )

        neg_prompt_ids, neg_prompt_attention_mask = _get_t5_text_ids(
            _prompt_clean(neg_prompt),
            self.t5_tokenizer,
        )
        neg_t5_embeds = self.model.get_t5_text_embeddings_sample(neg_prompt_ids.to(device), neg_prompt_attention_mask.to(device))
        
        cond_embeds_wtxt_wvit = torch.cat([t5_embeds, cond_embeds_wtxt_wvit], dim=1)
        if cond_embeds_wtxt_wovit is not None:
            cond_embeds_wtxt_wovit = torch.cat([t5_embeds, cond_embeds_wtxt_wovit], dim=1)
        if cond_embeds_wotxt_wvit is not None:
            cond_embeds_wotxt_wvit = torch.cat([neg_t5_embeds, cond_embeds_wotxt_wvit], dim=1)
        cond_embeds_wotxt_wovit = torch.cat([neg_t5_embeds, cond_embeds_wotxt_wovit], dim=1)
        if getattr(self.model, "t5_text_encoder", None) is not None:
            self.model.t5_text_encoder.to('cpu')
        torch.cuda.empty_cache()

        def is_image_vae_shape(shape):
            shape = shape.tolist() if isinstance(shape, torch.Tensor) else shape
            return shape[0] == 1 or (len(shape) > 1 and shape[1] == 1)

        def pad_and_truncate_feat(feat, max_sequence_length=max_sequence_length, truncate=use_truncate):
            if feat is None:
                return None
            if feat.shape[1] < max_sequence_length:
                feat = torch.cat(
                    [feat, feat.new_zeros((1, max_sequence_length-feat.shape[1], feat.shape[-1]))],
                    dim=1
                )
            if truncate and feat.shape[1] > max_sequence_length:
                feat = feat[:, :max_sequence_length, :]
            return feat

        if max_sequence_length > 0:
            cond_embeds_wtxt_wvit = pad_and_truncate_feat(cond_embeds_wtxt_wvit)
            cond_embeds_wotxt_wovit = pad_and_truncate_feat(cond_embeds_wotxt_wovit)
            if cond_embeds_wtxt_wovit is not None:
                cond_embeds_wtxt_wovit = pad_and_truncate_feat(cond_embeds_wtxt_wovit)
            if cond_embeds_wotxt_wvit is not None:
                cond_embeds_wotxt_wvit = pad_and_truncate_feat(cond_embeds_wotxt_wvit)
        
        pos = 0
        all_vae_latents = inputs['input_vae_latents']
        all_vae_rope = inputs['input_vae_rope']
        src_image_vae_latents, src_image_vae_shapes, src_image_vae_rope = [], [], []
        src_video_vae_latents, src_video_vae_shapes, src_video_vae_rope = [], [], []
        for i, shape in enumerate(inputs['input_vae_shape'][:-1]):
            vae_len = shape[1] * shape[2] // 4 * shape[0]
            cur_vae_rope = all_vae_rope[pos:pos+vae_len]
            cur_vae_latent = all_vae_latents[pos:pos+vae_len]
            if is_image_vae_shape(shape):
                src_image_vae_latents.append(cur_vae_latent)
                src_image_vae_shapes.append(shape)
                src_image_vae_rope.append(cur_vae_rope)
            else:
                src_video_vae_latents.append(cur_vae_latent)
                src_video_vae_shapes.append(shape)
                src_video_vae_rope.append(cur_vae_rope)
            pos += vae_len

        # Diffusion dimensions are derived from the transformed target VAE shape
        target_vae_shape = inputs['input_vae_shape'][-1]
        target_t = int(target_vae_shape[0])
        num_frames = min(num_frames, 1 + (target_t - 1) * 4)
        if height is None or height <= 0:
            height = int(target_vae_shape[1]) * 8
        if width is None or width <= 0:
            width = int(target_vae_shape[2]) * 8

        torch.cuda.empty_cache()
        latents = self.model.diff_dec.sample_bernini_wvitcfg(
            prompt_embeds_wtxt_wvit=cond_embeds_wtxt_wvit.to(self.weight_dtype),
            prompt_embeds_wtxt_wovit=cond_embeds_wtxt_wovit.to(self.weight_dtype) if cond_embeds_wtxt_wovit is not None else None,
            prompt_embeds_wotxt_wvit=cond_embeds_wotxt_wvit.to(self.weight_dtype) if cond_embeds_wotxt_wvit is not None else None,
            prompt_embeds_wotxt_wovit=cond_embeds_wotxt_wovit.to(self.weight_dtype),
            source_image_vae_latents=None if len(src_image_vae_latents) == 0 else torch.cat(src_image_vae_latents, dim=0),
            source_image_vae_rope=None if len(src_image_vae_rope) == 0 else torch.cat(src_image_vae_rope, dim=0),
            source_video_vae_latents=None if len(src_video_vae_latents) == 0 else torch.cat(src_video_vae_latents, dim=0),
            source_video_vae_rope=None if len(src_video_vae_rope) == 0 else torch.cat(src_video_vae_rope, dim=0),
            num_frames=num_frames,
            width=width,
            height=height,
            omega_txt=omega_txt,
            omega_img=omega_img,
            omega_vid=omega_vid,
            omega_tgt=omega_tgt,
            omega_scale=omega_scale,
            num_inference_steps=num_inference_steps,
            guidance_mode=guidance_mode,
            flow_shift=flow_shift,
            seed=seed,
            device=device,
        )

        if not write_output:
            return None

        self.vae.to(device)
        output = _vae_decode(self.vae, latents)
        self.vae.to("cpu")
        torch.cuda.empty_cache()

        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        save_output(output, output_path, fps=vae_fps)
        logger.info("saved -> %s  (%d frames, %dx%d)", output_path, output.shape[0], height, width)
        return output_path