diffusers-internal-dev/chronoedit-modular

1# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.2# SPDX-License-Identifier: Apache-2.03#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16from diffusers.modular_pipelines import ModularPipelineBlocks, InputParam, OutputParam, ModularPipeline, PipelineState17import numpy as np18import torch19import PIL20from typing import List21from diffusers.modular_pipelines.wan.before_denoise import WanInputStep22 23 24def calculate_dimensions(image, mod_value):25    """26    Calculate output dimensions based on resolution settings.27 28    Args:29        image: PIL Image30        mod_value: Modulo value for dimension alignment31 32    Returns:33        Tuple of (width, height)34    """35 36    # Get max area from preset or override37    target_area = 720 * 128038 39    # Calculate dimensions maintaining aspect ratio40    aspect_ratio = image.height / image.width41    calculated_height = round(np.sqrt(target_area * aspect_ratio)) // mod_value * mod_value42    calculated_width = round(np.sqrt(target_area / aspect_ratio)) // mod_value * mod_value43 44    return calculated_width, calculated_height45 46 47# Make the input step aware of `negative_prompt_embeds`.48# ChronoEdit uses a `guidance_scale` of 1.49class ChronoEditInputStep(WanInputStep):50    model_name = "chronoedit"51 52    @property53    def inputs(self) -> List[InputParam]:54        return [55            InputParam("num_videos_per_prompt", default=1),56            InputParam(57                "prompt_embeds",58                required=True,59                type_hint=torch.Tensor,60                description="Pre-generated text embeddings. Can be generated from text_encoder step.",61            ),62            InputParam(63                "negative_prompt_embeds",64                type_hint=torch.Tensor,65                description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",66            ),67        ]68 69 70class ChronoEditImageInputStep(ModularPipelineBlocks):71    model_name = "chronoedit"72 73    @property74    def inputs(self) -> List[InputParam]:75        return [InputParam(name="image")]76 77    @property78    def intermediate_outputs(self) -> List[OutputParam]:79        return [80            OutputParam(name="image", type_hint=PIL.Image.Image),81            OutputParam(name="height", type_hint=int, description="The height set w.r.t input image and specs"),82            OutputParam(name="width", type_hint=int, description="The width set w.r.t input image and specs"),83        ]84 85    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:86        block_state = self.get_block_state(state)87        image = block_state.image88        mod_value = components.vae_scale_factor_spatial * components.transformer.config.patch_size[1]89 90        width, height = calculate_dimensions(image, mod_value)91        block_state.image = image.resize((width, height))92        block_state.height = height93        block_state.width = width94 95        self.set_block_state(state, block_state)96        return components, state97