diffusers-internal-dev/chronoedit-modular

1"""2Mimicked from https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py3"""4 5from diffusers.modular_pipelines import WanModularPipeline, ModularPipelineBlocks6from diffusers.utils import load_image7from diffusers import UniPCMultistepScheduler8import torch9from PIL import Image10 11repo_id = "diffusers-internal-dev/chronoedit-modular"12blocks = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True)13pipe = WanModularPipeline(blocks, repo_id)14pipe.load_components(15    trust_remote_code=True,16    device_map="cuda",17    torch_dtype={"default": torch.bfloat16, "image_encoder": torch.float32},18)19pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)20pipe.load_lora_weights("nvidia/ChronoEdit-14B-Diffusers", weight_name="lora/chronoedit_distill_lora.safetensors")21pipe.fuse_lora(lora_scale=1.0)22 23image = load_image("https://huggingface.co/spaces/nvidia/ChronoEdit/resolve/main/examples/3.png")24prompt = "Transform the image so that inside the floral teacup of steaming tea, a small, cute mouse is sitting and taking a bath; the mouse should look relaxed and cheerful, with a tiny white bath towel draped over its head as if enjoying a spa moment, while the steam rises gently around it, blending seamlessly with the warm and cozy atmosphere."25 26# image is resized within the pipeline unlike https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py#L15127# refer to `ChronoEditImageInputStep`.28out = pipe(29    image=image,30    prompt=prompt,  # todo: enhance prompt31    num_inference_steps=8,  # todo: implement temporal reasoning32    num_frames=5,  # https://huggingface.co/spaces/nvidia/ChronoEdit/blob/main/app.py#L15233    output_type="np",34    generator=torch.manual_seed(0),35)36frames = out.values["videos"][0]37Image.fromarray((frames[-1] * 255).clip(0, 255).astype("uint8")).save("demo.png")38