smthemex-ComfyUI_ChronoEdit_SM/Chrono/scripts/run_inference_diffusers.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

"""
ChronoEdit Diffusers Inference Script

This script provides a command-line interface for running video editing inference
using the ChronoEdit model with the Diffusers backend.

Example usage:

# Basic usage
PYTHONPATH=$(pwd) python scripts/run_inference_diffusers.py \
    --input assets/images/input_2.png --offload_model --use-prompt-enhancer \
    --prompt "Add a sunglasses to the cat's face" \
    --output output2.mp4 --seed 42  \
    --model-path ./checkpoints/ChronoEdit-14B-Diffusers

# Basic usage with temporal reasoning
PYTHONPATH=$(pwd) python scripts/run_inference_diffusers.py \
    --enable-temporal-reasoning assets/images/input_2.png --offload_model \
    --prompt "Add a sunglasses to the cat's face"  \
    --output output_reasoning_offload.mp4 --seed 42 \
    --model-path ./checkpoints/ChronoEdit-14B-Diffusers

# Advanced usage with lora settings
PYTHONPATH=$(pwd) python scripts/run_inference_diffusers.py \
    --input assets/images/input_2.png \
    --prompt "Add a sunglasses to the cat's face"  \
    --output output_lora.mp4 \
    --num-inference-steps 8 \
    --guidance-scale 1.0 \
    --flow-shift 2.0 \
    --lora-scale 1.0 \
    --seed 42 \
    --lora-path ./checkpoints/ChronoEdit-14B/nvidia/lora/chronoedit_distill_lora_clean.safetensors \
    --model-path ./checkpoints/ChronoEdit-14B-Diffusers
"""


import os
import sys
import numpy as np
import torch
from diffusers.schedulers import UniPCMultistepScheduler
from omegaconf import OmegaConf
from diffusers import  GGUFQuantizationConfig
from diffusers.hooks import apply_group_offloading
#from diffusers import WanTransformer3DModel
from ..chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
from contextlib import contextmanager
from ..chronoedit_diffusers.transformer_chronoedit_ import WanTransformer3DModel
from .prompt_enhancer import load_model as load_prompt_enhancer
from .prompt_enhancer import enhance_prompt
from transformers import Qwen2_5_VLForConditionalGeneration,AutoConfig

from ..util import load_checkpoint_and_dispatch_
try:
    from transformers import Qwen3VLMoeForConditionalGeneration
    transfomer_vrsion_low=False
except:
    transfomer_vrsion_low=True
from accelerate import init_empty_weights

# try:
#     diffusers_module = sys.modules.get('diffusers')
#     if diffusers_module:
#         setattr(diffusers_module, 'WanTransformer3DModel', WanTransformer3DModel)
# except Exception as e:
#     print(f"Warning: Could not register WanTransformer3DModel with diffusers module: {e}")

# Resolution presets
RESOLUTION_PRESETS = {
    "480p": 480 * 832,
    "720p": 720 * 1280,
    "1080p": 1080 * 1920,
}

@contextmanager
def temp_patch_module_attr(module_name: str, attr_name: str, new_obj):
    mod = sys.modules.get(module_name)
    if mod is None:
        yield
        return
    had = hasattr(mod, attr_name)
    orig = getattr(mod, attr_name, None)
    setattr(mod, attr_name, new_obj)
    try:
        yield
    finally:
        if had:
            setattr(mod, attr_name, orig)
        else:
            try:
                delattr(mod, attr_name)
            except Exception:
                pass

def calculate_dimensions(image,  mod_value):
    """
    Calculate output dimensions based on resolution settings.
    
    Args:
        image: PIL Image
        mod_value: Modulo value for dimension alignment # mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]  *2
        
    Returns:
        Tuple of (width, height)
    """
    
    # Get max area from preset or override 
    target_area = 720 * 1280
    
    # Calculate dimensions maintaining aspect ratio
    if isinstance(image, torch.Tensor):
        aspect_ratio=image.shape[1] / image.shape[2]
    else:
        aspect_ratio = image.height / image.width
    calculated_height = round(np.sqrt(target_area * aspect_ratio)) // mod_value * mod_value
    calculated_width = round(np.sqrt(target_area / aspect_ratio)) // mod_value * mod_value
    
    return calculated_width, calculated_height


def load_dit_model(ckpt_path, gguf_path,dir_path):

    if gguf_path is not None:
        with temp_patch_module_attr("diffusers", "WanTransformer3DModel", WanTransformer3DModel):
            transformer = WanTransformer3DModel.from_single_file(
                gguf_path,
                config=os.path.join(dir_path, "Chrono/ChronoEdit-14B-Diffusers/transformer"),
                quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
                torch_dtype=torch.bfloat16,
                )
    elif ckpt_path is not None:
        from safetensors.torch import load_file
        sd=load_file(ckpt_path,device="cpu")
        sd=replace_key(sd)
        with init_empty_weights():
            with temp_patch_module_attr("diffusers", "WanTransformer3DModel", WanTransformer3DModel):
                config=WanTransformer3DModel.load_config(os.path.join(dir_path, "Chrono/ChronoEdit-14B-Diffusers/transformer"))
                transformer=WanTransformer3DModel.from_config(config)
        model_state_dict = transformer.state_dict()
        expected_keys = set(model_state_dict.keys())
        del transformer
        # 过滤掉模型不需要的键值对
        filtered_sd = {k: v for k, v in sd.items() if k in expected_keys}

        # 检查是否有缺失的键
        missing_keys = expected_keys - set(sd.keys())
        if missing_keys:
            print(f"Warning: Missing keys in checkpoint: {missing_keys}")
        del sd  
        with temp_patch_module_attr("diffusers", "WanTransformer3DModel", WanTransformer3DModel): 
            transformer = WanTransformer3DModel.from_single_file(filtered_sd,config=os.path.join(dir_path, "Chrono/ChronoEdit-14B-Diffusers/transformer"),torch_dtype=torch.bfloat16)
        del filtered_sd
    vae_config = OmegaConf.load(os.path.join(dir_path, "Chrono/ChronoEdit-14B-Diffusers/vae/config.json")) 
    pipe = ChronoEditPipeline.from_pretrained(
            os.path.join(dir_path, "Chrono/ChronoEdit-14B-Diffusers"),
            image_encoder=None,
            text_encoder=None,
            transformer=transformer,
            vae=None,
            torch_dtype=torch.bfloat16,
            vae_config=vae_config,
        )
    return pipe


def load_prompt_model(prompt_enhancer_model):
    model, processor = load_prompt_enhancer(prompt_enhancer_model)
    return model, processor


def prompt_enhance(input_image,prompt,prompt_model,processor,):

    cot_prompt = enhance_prompt(
            input_image,
            prompt,
            prompt_model,
            processor,
        )

    # try:
    #     if isinstance(prompt_model, torch.nn.Module) and hasattr(prompt_model, "_hf_hook"):
    #         accelerate.hooks.remove_hook_from_module(prompt_model, recurse=True)
    #         print("Offloaded prompt model to CPU")
    # except:pass
        
    torch.cuda.empty_cache()
    return cot_prompt

def load_prompt_enhancer_cf(clip_path,repo):
    if "qwen3" in clip_path.lower() and transfomer_vrsion_low:
        raise ValueError("Qwen3 is not supported in low version transfomer")
    print(f"Loading clip text encoder from {clip_path}...")
    config=AutoConfig.from_pretrained(repo)
    with init_empty_weights():
        if "qwen3" in clip_path.lower() and  not transfomer_vrsion_low:  
            model = Qwen3VLMoeForConditionalGeneration._from_config(config)
        elif "qwen2" in clip_path.lower() or "2.5" in clip_path.lower():    
            model = Qwen2_5_VLForConditionalGeneration._from_config(config)
    if clip_path.endswith('.safetensors'):
        from safetensors.torch import load_file
        state_dict = load_file(clip_path)
    else:
        state_dict = torch.load(clip_path,weights_only=False, map_location='cpu')   
    state_dict = remap_state_dict_keys(state_dict)
    model=load_checkpoint_and_dispatch_(model, state_dict, device_map="auto",dtype=torch.bfloat16)
    del state_dict
    model.eval() 
    print("Text encoder loaded successfully")
    return model


def replace_key(t_state_dict):
    new_dict = {}
    unused_keys = {'model_sampling.sigmas'}
    for k, v in t_state_dict.items():
        if k in unused_keys:
            continue
        new_k = k
        if k.startswith("text_embedding.0."):
            new_k = k.replace("text_embedding.0.", "condition_embedder.text_embedder.linear_1.")
        elif k.startswith("text_embedding.2."):
            new_k = k.replace("text_embedding.2.", "condition_embedder.text_embedder.linear_2.")
        elif k.startswith("time_embedding.0."):
            new_k = k.replace("time_embedding.0.", "condition_embedder.time_embedder.linear_1.")
        elif k.startswith("time_embedding.2."):
            new_k = k.replace("time_embedding.2.", "condition_embedder.time_embedder.linear_2.")
        elif k.startswith("time_projection.1."):
            new_k = k.replace("time_projection.1.", "condition_embedder.time_proj.")
        elif k.startswith("img_emb.proj.1."):
            new_k = k.replace("img_emb.proj.1.", "condition_embedder.image_embedder.ff.net.0.proj.")
        elif k.startswith("img_emb.proj.3."):
            new_k = k.replace("img_emb.proj.3.", "condition_embedder.image_embedder.ff.net.2.")
        elif k.startswith("img_emb.proj.0."):
            new_k = k.replace("img_emb.proj.0.", "condition_embedder.image_embedder.norm1.")
        elif k.startswith("img_emb.proj.4."):
            new_k = k.replace("img_emb.proj.4.", "condition_embedder.image_embedder.norm2.")
        elif k.startswith("head.modulation"):
            new_k = k.replace("head.modulation", "scale_shift_table")
        elif k.startswith("head.head."):
            new_k = k.replace("head.head.", "proj_out.")
        elif k.startswith("blocks."):
            if ".ffn.0." in k:
                new_k = k.replace(".ffn.0.", ".ffn.net.0.proj.")
            elif '.ffn.2.' in k:
                new_k = k.replace('.ffn.2.', '.ffn.net.2.')
            elif '.modulation' in k:
                new_k = k.replace('.modulation', '.scale_shift_table')
            elif ".norm3." in k:
                new_k = k.replace('.norm3.', '.norm2.')
            elif "self_attn.o." in k:
                new_k = k.replace("self_attn.o.", "attn1.to_out.0.")
            elif "self_attn.q." in k:
                new_k = k.replace("self_attn.q.", "attn1.to_q.")
            elif "self_attn.k." in k:
                new_k = k.replace("self_attn.k.", "attn1.to_k.")
            elif "self_attn.v." in k:
                new_k = k.replace("self_attn.v.", "attn1.to_v.")
            elif ".self_attn.norm_k." in k:
                new_k = k.replace(".self_attn.norm_k.", ".attn1.norm_k.")
            elif ".self_attn.norm_q." in k:
                new_k = k.replace(".self_attn.norm_q.", ".attn1.norm_q.")
            elif ".cross_attn.k." in k:
                new_k = k.replace(".cross_attn.k.", ".attn2.to_k.")
            elif ".cross_attn.q." in k:
                new_k = k.replace(".cross_attn.q.", ".attn2.to_q.")
            elif ".cross_attn.v." in k:
                new_k = k.replace(".cross_attn.v.", ".attn2.to_v.")
            elif ".cross_attn.o." in k:
                new_k = k.replace(".cross_attn.o.", ".attn2.to_out.0.")
            elif ".cross_attn.norm_k." in k:
                new_k = k.replace(".cross_attn.norm_k.", ".attn2.norm_k.")
            elif ".cross_attn.norm_q." in k:
                new_k = k.replace(".cross_attn.norm_q.", ".attn2.norm_q.")
            elif ".cross_attn.norm_k_img." in k:
                new_k = k.replace(".cross_attn.norm_k_img.", ".attn2.norm_added_k.")
            elif ".cross_attn.k_img." in k:
                new_k = k.replace(".cross_attn.k_img.", ".attn2.add_k_proj.")
            elif ".cross_attn.v_img." in k:
                new_k = k.replace(".cross_attn.v_img.", ".attn2.add_v_proj.")
            else:
                new_k = k   
        else:
            new_k = k
        
        new_dict[new_k] = v
    return new_dict



def remap_state_dict_keys(state_dict):
    """Remap state dict keys to match model architecture"""
    new_state_dict = {}
    
    for key, value in state_dict.items():
        new_key = key
        # 处理视觉编码器前缀
        if key.startswith('visual.') and not key.startswith('model.visual.'):
            new_key = 'model.' + key
        # 处理语言模型前缀
        elif key.startswith('model.layers.') and not key.startswith('model.language_model.'):
            new_key = 'model.language_model.' + key[len('model.'):]
        # 处理嵌入层
        elif key == 'model.embed_tokens.weight':
            new_key = 'model.language_model.embed_tokens.weight'
        elif key == 'model.norm.weight':
            new_key = 'model.language_model.norm.weight'
            
        new_state_dict[new_key] = value
    
    return new_state_dict

def inference_chrono(pipe,image_latent,positive,image_embeds,negative,flow_shift,seed,num_inference_steps,height,width,
              num_frames,num_temporal_reasoning_steps,guidance_scale=5.0,offload_model=True,block_num=1,device="cuda"):
    
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config,flow_shift=flow_shift)
    apply_group_offloading(pipe.transformer, onload_device=torch.device("cuda"), offload_type="block_level", num_blocks_per_group=block_num)

    generator = torch.Generator(device=device).manual_seed(seed)
    #num_frames = 29 if enable_temporal_reasoning else 5
    enable_temporal_reasoning=True if num_frames == 29 else False
    print(f"Generating {num_frames} frames with {num_inference_steps} steps...")
   
    output = pipe(
        image=None,
        prompt=None,
        negative_prompt=None,
        image_embeds=image_embeds, #torch.Size([1, 257, 1280])
        prompt_embeds=positive,
        negative_prompt_embeds=negative,
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        enable_temporal_reasoning=enable_temporal_reasoning,
        num_temporal_reasoning_steps=num_temporal_reasoning_steps,
        generator=generator,
        offload_model=offload_model,
        image_latent=image_latent,
    ).frames

    return output