Upload folder using huggingface_hub

.ruff_cache/.gitignore

@@ -0,0 +1,2 @@
+# Automatically created by ruff.
+*

.ruff_cache/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

before_denoise.py

@@ -37,17 +37,11 @@ class ChronoEditSetTimestepsStep(ModularPipelineBlocks):
 
     @property
     def expected_components(self) -> List[ComponentSpec]:
-        return [
-            ComponentSpec("scheduler", UniPCMultistepScheduler)
-        ]
+        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
 
     @property
     def inputs(self) -> List[InputParam]:
-        return [
-            InputParam("num_inference_steps", default=50),
-            InputParam("timesteps"),
-            InputParam("sigmas")
-        ]
+        return [InputParam("num_inference_steps", default=50), InputParam("timesteps"), InputParam("sigmas")]
 
     @property
     def intermediate_outputs(self) -> List[OutputParam]:
@@ -72,7 +66,7 @@ class ChronoEditSetTimestepsStep(ModularPipelineBlocks):
             block_state.timesteps,
             block_state.sigmas,
         )
-        
+
         self.set_block_state(state, block_state)
         return components, state
 
@@ -186,18 +180,18 @@ class ChronoEditPrepareLatentStep(ModularPipelineBlocks):
         )
         mask_lat_size = mask_lat_size.transpose(1, 2)
         mask_lat_size = mask_lat_size.to(latent_condition.device)
-        
+
         return latents, torch.concat([mask_lat_size, latent_condition], dim=1)
 
     @torch.no_grad()
     def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
-        
+
         self.check_inputs(block_state.height, block_state.width)
 
         block_state.device = components._execution_device
         block_state.num_channels_latents = components.num_channels_latents
-        
+
         batch_size = block_state.batch_size * block_state.num_videos_per_prompt
         block_state.latents, block_state.condition = self.prepare_latents(
             components,
@@ -214,5 +208,5 @@ class ChronoEditPrepareLatentStep(ModularPipelineBlocks):
         )
 
         self.set_block_state(state, block_state)
-        
+
         return components, state

decoders.py

@@ -14,14 +14,12 @@
 # limitations under the License.
 
 from diffusers.configuration_utils import FrozenDict
-from diffusers.modular_pipelines import (
-    ModularPipelineBlocks, ComponentSpec, InputParam, OutputParam, PipelineState
-)
+from diffusers.modular_pipelines import ModularPipelineBlocks, ComponentSpec, InputParam, OutputParam, PipelineState
 from typing import List, Union
 from diffusers import AutoencoderKLWan
 from diffusers.video_processor import VideoProcessor
 import torch
-import PIL 
+import PIL
 import numpy as np
 
 
@@ -93,4 +91,4 @@ class ChronoEditDecodeStep(ModularPipelineBlocks):
 
         self.set_block_state(state, block_state)
 
-        return components, state
+        return components, state

denoise.py

@@ -190,7 +190,7 @@ class ChronoEditDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
     @torch.no_grad()
     def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
-        
+
         block_state.num_warmup_steps = max(
             len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
         )

encoders.py

@@ -88,7 +88,7 @@ class ChronoEditProcessImageStep(ModularPipelineBlocks):
     def intermediate_outputs(self) -> List[OutputParam]:
         return [
             OutputParam("processed_image", type_hint=PipelineImageInput),
-            OutputParam("image_embeds", type_hint=torch.Tensor)
+            OutputParam("image_embeds", type_hint=torch.Tensor),
         ]
 
     @property
@@ -118,7 +118,7 @@ class ChronoEditProcessImageStep(ModularPipelineBlocks):
             block_state.image_embeds = image_embeds.repeat(batch_size, 1, 1).to(torch.bfloat16)
 
         self.set_block_state(state, block_state)
-        
+
         return components, state
 
 

example.py

@@ -2,24 +2,22 @@
 Mimicked from https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py
 """
 
-from diffusers.modular_pipelines import ModularPipelineBlocks, WanModularPipeline
+from diffusers.modular_pipelines import WanModularPipeline, ModularPipelineBlocks
 from diffusers.utils import load_image
 from diffusers import UniPCMultistepScheduler
-from modular_blocks import ChronoEditBlocks
 import torch
 from PIL import Image
 
-repo_id = "nvidia/ChronoEdit-14B-Diffusers"
-blocks = ChronoEditBlocks()
+repo_id = "diffusers-internal-dev/chronoedit-modular"
+blocks = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True)
 pipe = WanModularPipeline(blocks, repo_id)
-
 pipe.load_components(
     trust_remote_code=True,
     device_map="cuda",
     torch_dtype={"default": torch.bfloat16, "image_encoder": torch.float32},
 )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)
-pipe.load_lora_weights(repo_id, weight_name="lora/chronoedit_distill_lora.safetensors")
+pipe.load_lora_weights("nvidia/ChronoEdit-14B-Diffusers", weight_name="lora/chronoedit_distill_lora.safetensors")
 pipe.fuse_lora(lora_scale=1.0)
 
 image = load_image("https://huggingface.co/spaces/nvidia/ChronoEdit/resolve/main/examples/3.png")
@@ -31,8 +29,9 @@ out = pipe(
     image=image,
     prompt=prompt,  # todo: enhance prompt
     num_inference_steps=8,  # todo: implement temporal reasoning
-    num_frames=5, # https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py#L152
+    num_frames=5,  # https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py#L152
     output_type="np",
+    generator=torch.manual_seed(0),
 )
 frames = out.values["videos"][0]
 Image.fromarray((frames[-1] * 255).clip(0, 255).astype("uint8")).save("demo.png")

inputs.py

@@ -86,11 +86,11 @@ class ChronoEditImageInputStep(ModularPipelineBlocks):
         block_state = self.get_block_state(state)
         image = block_state.image
         mod_value = components.vae_scale_factor_spatial * components.transformer.config.patch_size[1]
-        
+
         width, height = calculate_dimensions(image, mod_value)
         block_state.image = image.resize((width, height))
-        block_state.height = height 
+        block_state.height = height
         block_state.width = width
-        
+
         self.set_block_state(state, block_state)
         return components, state

upload_block.py

@@ -0,0 +1,11 @@
+from diffusers.modular_pipelines import WanModularPipeline
+
+from modular_blocks import ChronoEditBlocks
+
+repo_id = "nvidia/ChronoEdit-14B-Diffusers"
+blocks = ChronoEditBlocks()
+
+blocks.push_to_hub("diffusers-internal-dev/chronoedit-modular")
+
+pipe = WanModularPipeline(blocks, repo_id)
+pipe.push_to_hub("diffusers-internal-dev/chronoedit-modular")