diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 82a15e66d5bd..5c6ac126dac7 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -349,6 +349,8 @@ title: DiTTransformer2DModel - local: api/models/easyanimate_transformer3d title: EasyAnimateTransformer3DModel + - local: api/models/edenlive_transformer_3d + title: EdenLiveTransformer3DModel - local: api/models/ernie_image_transformer2d title: ErnieImageTransformer2DModel - local: api/models/flux2_transformer @@ -729,6 +731,8 @@ title: Cosmos - local: api/pipelines/cosmos3 title: Cosmos3 + - local: api/pipelines/edenlive + title: EdenLive - local: api/pipelines/framepack title: Framepack - local: api/pipelines/helios @@ -798,6 +802,8 @@ title: DPMSolverSDEScheduler - local: api/schedulers/singlestep_dpm_solver title: DPMSolverSinglestepScheduler + - local: api/schedulers/edenlive + title: EdenLiveFlowMatchEulerDiscreteScheduler - local: api/schedulers/edm_multistep_dpm_solver title: EDMDPMSolverMultistepScheduler - local: api/schedulers/edm_euler diff --git a/docs/source/en/api/models/edenlive_transformer_3d.md b/docs/source/en/api/models/edenlive_transformer_3d.md new file mode 100644 index 000000000000..d0bac2de858a --- /dev/null +++ b/docs/source/en/api/models/edenlive_transformer_3d.md @@ -0,0 +1,47 @@ + + +# EdenLiveTransformer3DModel + +[EdenLive](https://github.com/VAST-AI-Research/EdenLive) uses a video diffusion transformer conditioned on +text, a reference image, and DINO proxy features projected inside each transformer forward pass. The same model +class supports bidirectional and causal attention. See the [pipeline guide](../pipelines/edenlive) for video generation. + +```python +import torch +from diffusers import EdenLiveTransformer3DModel + +transformer = EdenLiveTransformer3DModel.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 +) +``` + +For causal weights, load `VAST-AI/EdenLive-5B-Causal-Diffusers` or `VAST-AI/EdenLive-14B-Causal-Diffusers` with the same +`subfolder="transformer"`. Keep model timesteps in FP32 when calling the transformer directly, including when its +weights are BF16. Pass raw DINO patch features as `proxy_features` with shape `(batch, proxy_input_channels, frames, +height, width)`; `proxy_channels` sets the number of projected channels. The internal projection remains in FP32 +when loaded with `from_pretrained`. + +## EdenLiveTransformer3DModel + +[[autodoc]] EdenLiveTransformer3DModel + - forward + - create_causal_state + +## EdenLiveCausalState + +The causal pipeline manages this state automatically. For custom loops, create a separate state for each CFG branch +with `transformer.create_causal_state()`. Denoise with `cache_update=False`, then commit each clean chunk at timestep +zero with `cache_update=True`. Call `reset()` before an independent video. + +[[autodoc]] EdenLiveCausalState + - reset diff --git a/docs/source/en/api/pipelines/edenlive.md b/docs/source/en/api/pipelines/edenlive.md new file mode 100644 index 000000000000..f4791bb2ef04 --- /dev/null +++ b/docs/source/en/api/pipelines/edenlive.md @@ -0,0 +1,136 @@ + + +# EdenLive + +[EdenLive](https://github.com/VAST-AI-Research/EdenLive) generates video from a proxy video, a reference +image, and a text prompt. [`EdenLivePipeline`] supports bidirectional generation; [`EdenLiveCausalPipeline`] +generates a complete video in causal chunks and supports changing prompts between chunks. + +## Basic usage + +Prepare an RGB reference image and a proxy video with at least 49 frames sampled at 24 fps. The image provides +appearance, while the proxy provides motion and scene structure. + +```python +import torch +from diffusers import EdenLivePipeline +from diffusers.utils import export_to_video, load_image, load_video + +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, +).to("cuda") + +image = load_image("reference.png") +video = load_video("proxy-24fps.mp4") + +frames = pipe( + image=image, + video=video, + prompt="A character explores a snowy landscape.", + num_frames=49, +).frames[0] +export_to_video(frames, "edenlive.mp4", fps=24) +``` + +Keep the proxy encoder in FP32 when loading the other models in BF16, as shown above. The transformer keeps its +internal proxy projection in FP32 when loaded with `from_pretrained`. + +## Causal generation + +Load the causal pipeline from its dedicated repository. Both model sizes use four fixed denoising steps per chunk. + +```python +from diffusers import EdenLiveCausalPipeline + +causal_pipe = EdenLiveCausalPipeline.from_pretrained( + "VAST-AI/EdenLive-5B-Causal-Diffusers", + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, +).to("cuda") +``` + +Use `prompt` for fixed text, or `chunk_prompts` to change it during a single video. At the default chunk size, +49 frames require five prompts. Reusing the image and video above: + +```python +frames = causal_pipe( + image=image, + video=video, + chunk_prompts=[ + "A character walks through snow.", + "A character walks through snow.", + "The character approaches a campfire.", + "The character warms their hands beside the campfire.", + "The character rests beside the campfire.", + ], + num_frames=49, +).frames[0] +export_to_video(frames, "chunk-prompts.mp4", fps=24) +``` + +`chunk_prompts` overrides a single `prompt` and cannot be combined with `prompt_embeds`. The negative prompt remains +fixed. A list passed as `prompt` represents a batch of independent videos. Each pipeline call starts a new video. + +## LightTAE decoder + +The full Wan VAE is the default. For faster decoding, pass the optional [`EdenLiveLightTAEDecoder`] to either pipeline: + +```python +from diffusers import EdenLiveLightTAEDecoder + +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", + decoder=EdenLiveLightTAEDecoder.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", subfolder="decoder", torch_dtype=torch.bfloat16 + ), + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, +).to("cuda") +``` + +## Notes + +| Checkpoint | Pipeline | Default height × width | Example fps | Steps | +|---|---|---|---|---| +| [5B bidirectional](https://huggingface.co/VAST-AI/EdenLive-5B-Diffusers) | `EdenLivePipeline` | 704 × 1280 | 24 | 48 | +| [14B bidirectional](https://huggingface.co/VAST-AI/EdenLive-14B-Diffusers) | `EdenLivePipeline` | 480 × 832 | 16 | 30 | +| [5B causal](https://huggingface.co/VAST-AI/EdenLive-5B-Causal-Diffusers) | `EdenLiveCausalPipeline` | 704 × 1280 | 24 | 4 per chunk | +| [14B causal](https://huggingface.co/VAST-AI/EdenLive-14B-Causal-Diffusers) | `EdenLiveCausalPipeline` | 480 × 832 | 16 | 4 per chunk | + +- Sample proxy frames at the desired output frame rate. `num_frames` must be $4k+1$ and fit within the input video; + omitting it uses the largest valid prefix. Custom dimensions must be multiples of 32 for 5B or 16 for 14B. +- Guidance defaults to 3 for bidirectional generation and 1 for causal generation. Set `guidance_scale` above one + to use a `negative_prompt`. +- For CPU offload, load without `.to("cuda")` and call `enable_model_cpu_offload()`. +- See the [example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive) for installation, + video resampling, CPU offload, and checkpoint conversion. + +## EdenLivePipeline + +[[autodoc]] EdenLivePipeline + - all + - __call__ + +## EdenLiveCausalPipeline + +[[autodoc]] EdenLiveCausalPipeline + - all + - __call__ + +## EdenLiveLightTAEDecoder + +[[autodoc]] EdenLiveLightTAEDecoder + - decode + +## EdenLivePipelineOutput + +[[autodoc]] pipelines.edenlive.pipeline_output.EdenLivePipelineOutput diff --git a/docs/source/en/api/schedulers/edenlive.md b/docs/source/en/api/schedulers/edenlive.md new file mode 100644 index 000000000000..94bbe1124b11 --- /dev/null +++ b/docs/source/en/api/schedulers/edenlive.md @@ -0,0 +1,27 @@ + + +# EdenLiveFlowMatchEulerDiscreteScheduler + +[`EdenLiveFlowMatchEulerDiscreteScheduler`] implements deterministic flow Euler sampling for EdenLive 14B +bidirectional generation and stochastic sampling for the distilled causal checkpoints. The 5B bidirectional model +uses [`UniPCMultistepScheduler`]. + +Load the `scheduler` subfolder from the matching checkpoint repository. For causal generation, use +`VAST-AI/EdenLive-5B-Causal-Diffusers` or `VAST-AI/EdenLive-14B-Causal-Diffusers`; the released causal schedules have +four fixed steps per chunk. See the [pipeline guide](../pipelines/edenlive#causal-generation) for loading the pipeline. + +## EdenLiveFlowMatchEulerDiscreteScheduler + +[[autodoc]] EdenLiveFlowMatchEulerDiscreteScheduler + - set_timesteps + - step diff --git a/examples/edenlive/README.md b/examples/edenlive/README.md new file mode 100644 index 000000000000..f69c109be028 --- /dev/null +++ b/examples/edenlive/README.md @@ -0,0 +1,72 @@ +# EdenLive inference + +Generate video from a reference image, proxy video, and prompt. See the [pipeline guide](../../docs/source/en/api/pipelines/edenlive.md) +for the Python API and model defaults. + +## Setup + +Install a CUDA-compatible [PyTorch](https://pytorch.org/get-started/locally/) and torchvision, then run from the +Diffusers repository root: + +```bash +pip install -e . "transformers>=4.56.0" accelerate sentencepiece ftfy av imageio imageio-ffmpeg +``` + +The [requirements file](requirements.txt) also lists the example dependencies. + +## Single GPU + +```bash +python examples/edenlive/inference_edenlive.py \ + --model VAST-AI/EdenLive-5B-Causal-Diffusers --causal \ + --video proxy.mp4 --image reference.png \ + --prompt "A character explores a snowy landscape." \ + --num-frames 49 +``` + +Use `--causal` with a causal repository, or omit it for a bidirectional repository. `--model` also accepts a local +Diffusers directory. The script resamples the proxy video to the model's default frame rate or `--fps`. +`--num-frames` must be $4k+1$ and fit within the resampled video; omit it to use the largest valid prefix. + +The default output is `edenlive.mp4`, with timing and input details in `edenlive.json`. Set `--output` to change it. +Use `--lighttae` for faster decoding, `--vae-tiling` to reduce VAE memory use, or `--cpu-offload` to move models +to the GPU only when needed. Run +`python examples/edenlive/inference_edenlive.py --help` for all options. + +## Model repositories + +| Size | Bidirectional | Causal | +|---|---|---| +| 5B | [EdenLive-5B-Diffusers](https://huggingface.co/VAST-AI/EdenLive-5B-Diffusers) | [EdenLive-5B-Causal-Diffusers](https://huggingface.co/VAST-AI/EdenLive-5B-Causal-Diffusers) | +| 14B | [EdenLive-14B-Diffusers](https://huggingface.co/VAST-AI/EdenLive-14B-Diffusers) | [EdenLive-14B-Causal-Diffusers](https://huggingface.co/VAST-AI/EdenLive-14B-Causal-Diffusers) | + +Each repository contains one complete pipeline. Select the matching repository and pipeline variant. + +## Changing prompts + +With `--causal`, add `--chunk-prompts prompts.json`. For 49 frames, supply five prompts: + +```json +[ + "A character walks through snow.", + "A character walks through snow.", + "The character approaches a campfire.", + "The character warms their hands beside the campfire.", + "The character rests beside the campfire." +] +``` + +The list replaces `--prompt`, with one prompt per chunk. The negative prompt remains fixed. + +## Convert original checkpoints + +To convert a local original release instead of loading the published Diffusers weights: + +```bash +python scripts/convert_edenlive_to_diffusers.py \ + --source /path/to/EdenLive --size 5b --variant bidirectional \ + --output /path/to/EdenLive-5B-Diffusers +``` + +Use `--size 14b` for 14B or `--variant causal` for causal generation, with a separate output directory for each +combination. Each package contains one sampling variant and the optional LightTAE decoder. diff --git a/examples/edenlive/inference_edenlive.py b/examples/edenlive/inference_edenlive.py new file mode 100644 index 000000000000..015b7ee6a6c9 --- /dev/null +++ b/examples/edenlive/inference_edenlive.py @@ -0,0 +1,174 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +"""Offline EdenLive inference on one GPU, with optional CPU offload.""" + +import argparse +import json +import math +import time +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from diffusers import ( + EdenLiveCausalPipeline, + EdenLiveLightTAEDecoder, + EdenLivePipeline, +) +from diffusers.utils import export_to_video + + +def load_video(path, fps, height, width, num_frames=None): + import av + + if fps <= 0: + raise ValueError("fps must be positive.") + with av.open(str(path)) as container: + stream = container.streams.video[0] + duration = float(stream.duration * stream.time_base) if stream.duration is not None else None + timestamps = np.asarray([float(frame.time) for frame in container.decode(video=0)]) + if len(timestamps) == 0: + raise ValueError("The input video has no decoded frames.") + timestamps -= timestamps[0] + if duration is None: + duration = timestamps[-1] + (np.median(np.diff(timestamps)) if len(timestamps) > 1 else 1 / fps) + available = math.floor(duration * fps) + count = 4 * ((available - 1) // 4) + 1 if num_frames is None else num_frames + if count < 1 or (count - 1) % 4 or count > available: + raise ValueError(f"num_frames must be 4k+1 and no larger than {available} frames at {fps} fps.") + indices = np.searchsorted((timestamps[:-1] + timestamps[1:]) / 2, np.arange(count) / fps) + frames = [] + with av.open(str(path)) as container: + for index, frame in enumerate(container.decode(video=0)): + if len(frames) == count: + break + if index != indices[len(frames)]: + continue + picture = frame.to_image().convert("RGB") + crop_width = min(picture.width, round(picture.height * width / height)) + crop_height = min(picture.height, round(picture.width * height / width)) + left, top = (picture.width - crop_width) // 2, (picture.height - crop_height) // 2 + picture = picture.crop((left, top, left + crop_width, top + crop_height)) + picture = picture.resize((width, height), Image.Resampling.BILINEAR) + while len(frames) < count and indices[len(frames)] == index: + frames.append(picture) + if len(frames) != count: + raise ValueError(f"Decoded {len(frames)} of {count} requested frames.") + return frames, {"source_frame_indices": indices.tolist(), "fps": fps, "num_frames": count} + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Standard Diffusers model directory or Hub repository.") + parser.add_argument("--causal", action="store_true", help="Load a causal model with EdenLiveCausalPipeline.") + parser.add_argument( + "--lighttae", action="store_true", help="Use the optional LightTAE decoder instead of the full Wan VAE." + ) + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--image", type=Path, required=True) + parser.add_argument("--prompt") + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--chunk-prompts", type=Path, help="JSON list of prompts for chunk-causal inference.") + parser.add_argument("--output", type=Path, default=Path("edenlive.mp4")) + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--num-frames", type=int) + parser.add_argument("--fps", type=float) + parser.add_argument("--num-inference-steps", type=int) + parser.add_argument("--guidance-scale", type=float) + parser.add_argument("--seed", type=int, help="Noise seed; defaults to 42 for 14b-causal and 0 otherwise.") + parser.add_argument("--attention-backend") + parser.add_argument("--cpu-offload", action="store_true", help="Move models to the GPU only when needed.") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--vae-tiling", action="store_true") + args = parser.parse_args() + if args.chunk_prompts is not None and not args.causal: + parser.error("--chunk-prompts requires --causal.") + if args.prompt is None and args.chunk_prompts is None: + parser.error("Provide --prompt or --chunk-prompts.") + return args + + +def main(): + args = parse_args() + device = torch.device("cuda") + loading_kwargs = {"local_files_only": args.local_files_only} + components = {"decoder": None} + pipeline_class = EdenLiveCausalPipeline if args.causal else EdenLivePipeline + if args.lighttae: + components["decoder"] = EdenLiveLightTAEDecoder.from_pretrained( + args.model, subfolder="decoder", torch_dtype=torch.bfloat16, **loading_kwargs + ) + pipeline = pipeline_class.from_pretrained( + args.model, + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, + **components, + **loading_kwargs, + ) + if args.attention_backend is not None: + pipeline.transformer.set_attention_backend(args.attention_backend) + if args.cpu_offload: + pipeline.enable_model_cpu_offload() + else: + pipeline.to(device) + if args.vae_tiling: + pipeline.vae.enable_tiling() + is_5b = pipeline.transformer.config.image_dim is None + args.height = args.height if args.height is not None else (704 if is_5b else 480) + args.width = args.width if args.width is not None else (1280 if is_5b else 832) + args.fps = args.fps if args.fps is not None else (24 if is_5b else 16) + args.seed = args.seed if args.seed is not None else (42 if args.causal and not is_5b else 0) + video, video_info = load_video(args.video, args.fps, args.height, args.width, args.num_frames) + inputs = { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "image": Image.open(args.image).convert("RGB"), + "video": video, + "height": args.height, + "width": args.width, + "num_frames": len(video), + "generator": torch.Generator(device=device).manual_seed(args.seed), + "output_type": "np", + } + for name in ("num_inference_steps", "guidance_scale"): + if getattr(args, name) is not None: + inputs[name] = getattr(args, name) + if args.chunk_prompts is not None: + inputs["chunk_prompts"] = json.loads(args.chunk_prompts.read_text()) + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + started = time.perf_counter() + frames = pipeline(**inputs).frames[0] + torch.cuda.synchronize(device) + metrics = { + "seconds": time.perf_counter() - started, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + export_to_video(frames, str(args.output), fps=args.fps) + report = { + "arguments": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()}, + "pipeline": pipeline.__class__.__name__, + "video": video_info, + "metrics": metrics, + } + args.output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/examples/edenlive/requirements.txt b/examples/edenlive/requirements.txt new file mode 100644 index 000000000000..f68559fc7f2b --- /dev/null +++ b/examples/edenlive/requirements.txt @@ -0,0 +1,9 @@ +torch +torchvision +transformers>=4.56.0 +accelerate +av +imageio +imageio-ffmpeg +sentencepiece +ftfy diff --git a/scripts/convert_edenlive_to_diffusers.py b/scripts/convert_edenlive_to_diffusers.py new file mode 100644 index 000000000000..090420b21790 --- /dev/null +++ b/scripts/convert_edenlive_to_diffusers.py @@ -0,0 +1,276 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +"""Convert a local original EdenLive release into standard Diffusers components.""" + +import argparse +import gc +import hashlib +import json +import shutil +from pathlib import Path + +import torch +from transformers import ( + AutoTokenizer, + CLIPVisionConfig, + CLIPVisionModel, + DINOv3ViTConfig, + DINOv3ViTModel, + UMT5EncoderModel, +) + +from diffusers import ( + AutoencoderKLWan, + EdenLiveFlowMatchEulerDiscreteScheduler, + EdenLiveLightTAEDecoder, + EdenLiveTransformer3DModel, + UniPCMultistepScheduler, +) +from diffusers.loaders.single_file_utils import convert_wan_vae_to_diffusers +from diffusers.models.model_loading_utils import load_state_dict + + +def read_config(path): + return json.loads(path.read_text()) + + +def tensor_record(value): + raw = value.detach().cpu().contiguous().view(torch.uint8).numpy() + return {"shape": list(value.shape), "dtype": str(value.dtype), "sha256": hashlib.sha256(raw).hexdigest()} + + +def save_component(model, state, destination, manifest): + model.load_state_dict(state, strict=True, assign=True) + model.eval() + model.save_pretrained(destination, safe_serialization=True, max_shard_size="5GB") + dtype = next(iter(state.values())).dtype + restored = type(model).from_pretrained(destination, torch_dtype=dtype, local_files_only=True) + actual = restored.state_dict() + if actual.keys() != state.keys(): + raise ValueError(f"Round-trip keys differ for {destination.name}") + records = {} + for name, original in state.items(): + value = actual[name] + if original.dtype != value.dtype or not torch.equal(original, value): + raise ValueError(f"Round-trip changed {destination.name}/{name}") + records[name] = tensor_record(original) + manifest[destination.name] = records + del restored + gc.collect() + print(f"Verified {destination}: {len(records)} tensors", flush=True) + + +def convert_lighttae(state, model): + converted = { + key.replace("decoder.blocks.", "decoder.", 1): value + for key, value in state.items() + if not key.startswith("encoder.") + } + for key, expected in model.state_dict().items(): + if key.endswith(".conv.weight") and key in converted and converted[key].shape[0] > expected.shape[0]: + converted[key] = converted[key][-expected.shape[0] :] + return converted + + +def _proxy_state(state): + normalized = {} + for key, value in state.items(): + for prefix in ("model.module.", "model.", "module."): + if key.startswith(prefix): + key = key.removeprefix(prefix) + break + key = key.removeprefix("encoder.") + if key in normalized: + raise ValueError(f"Duplicate proxy checkpoint key: {key}.") + if not key.startswith(("backbone.", "patch_projection.")) and key not in ( + "global_projection.weight", + "global_projection.bias", + ): + raise ValueError(f"Unrecognized proxy checkpoint key: {key}.") + normalized[key] = value + return normalized + + +def _remap_wan22_vae_key(key): + """Normalize only the nested Wan 2.2 blocks missing from the native mapper.""" + parts = key.split(".") + if len(parts) < 5 or parts[:2] not in (["encoder", "downsamples"], ["decoder", "upsamples"]): + return key + block = "down_blocks" if parts[0] == "encoder" else "up_blocks" + prefix = f"{parts[0]}.{block}.{parts[2]}" + if parts[3] == "avg_shortcut": + return f"{prefix}.avg_shortcut.{'.'.join(parts[4:])}" + if len(parts) < 6 or parts[3] != parts[1]: + return key + slot, rest = parts[4], parts[5:] + if rest[0] == "residual" and len(rest) > 2: + field = {"0": "norm1", "2": "conv1", "3": "norm2", "6": "conv2"}.get(rest[1]) + if field is not None: + return f"{prefix}.resnets.{slot}.{field}.{'.'.join(rest[2:])}" + elif rest[0] == "shortcut": + return f"{prefix}.resnets.{slot}.conv_shortcut.{'.'.join(rest[1:])}" + elif parts[0] == "encoder" and slot == "2": + return f"{prefix}.downsampler.{'.'.join(rest)}" + elif parts[0] == "decoder" and slot == "3": + return f"{prefix}.upsampler.{'.'.join(rest)}" + return key + + +def convert(source, output, size, variant): + if output.exists() and any(output.iterdir()): + raise ValueError("The output directory must be empty.") + output.mkdir(parents=True, exist_ok=True) + manifest = {} + proxy_state = _proxy_state(load_state_dict(str(source / "vision-encoder/projection.safetensors"))) + projection = { + key.replace("patch_projection.", "proxy_projection.", 1): value + for key, value in proxy_state.items() + if key.startswith("patch_projection.") + } + del proxy_state + folder = f"dit-{size}-bidir" if variant == "bidirectional" else f"dit-{size}" + config = read_config(source / folder / "config.json") + config.pop("cross_attn_norm", None) + state = load_state_dict(str(source / folder / "diffusion_pytorch_model.safetensors")) + config["proxy_input_channels"] = projection["proxy_projection.weight"].shape[1] + config["proxy_channels"] = state["patch_embedding.weight"].shape[1] - config["in_channels"] + config["condition_rmsnorm"] = "patch_embedding.condition_norm.weight" in state + state.update(projection) + with torch.device("meta"): + model = EdenLiveTransformer3DModel.from_config(config) + for attribute in ("rope", "cache_rope"): + rope = getattr(model, attribute) + setattr(model, attribute, type(rope)(rope.attention_head_dim, tuple(rope.patch_size), rope.max_seq_len)) + save_component(model, state, output / "transformer", manifest) + del model, state, projection + gc.collect() + + config = read_config(source / f"vae-{size}" / "config.json") + version = "2" if size == "5b" else "1" + state = load_state_dict(str(source / f"vae-{size}" / f"Wan2.{version}_VAE.pth")) + state = convert_wan_vae_to_diffusers({_remap_wan22_vae_key(key): value for key, value in state.items()}) + config["in_channels"] = state["encoder.conv_in.weight"].shape[1] + config["out_channels"] = state["decoder.conv_out.weight"].shape[0] + with torch.device("meta"): + model = AutoencoderKLWan.from_config(config) + save_component(model, state, output / "vae", manifest) + del model, state + gc.collect() + + state = load_state_dict(str(source / f"decoder-{size}" / f"lighttaew2_{version}.pth")) + with torch.device("meta"): + model = EdenLiveLightTAEDecoder( + latent_channels=48 if size == "5b" else 16, patch_size=2 if size == "5b" else 1 + ) + save_component(model, convert_lighttae(state, model), output / "decoder", manifest) + del model, state + gc.collect() + + state = _proxy_state(load_state_dict(str(source / "vision-encoder/projection.safetensors"))) + config = DINOv3ViTConfig.from_dict(read_config(source / "vision-encoder/backbone/config.json")) + with torch.device("meta"): + model = DINOv3ViTModel(config) + model.rope_embeddings = type(model.rope_embeddings)(config) + expected = set(model.state_dict()) + backbone = {} + for key, value in state.items(): + if key.startswith("backbone."): + key = key.removeprefix("backbone.") + if key not in expected and f"model.{key}" in expected: + key = f"model.{key}" + backbone[key] = value + save_component(model, backbone, output / "proxy_encoder", manifest) + del model, state, backbone + gc.collect() + + for name in ("text_encoder", "tokenizer"): + shutil.copytree(source / "umt5" / name, output / name, dirs_exist_ok=True) + tokenizer = AutoTokenizer.from_pretrained(output / "tokenizer", local_files_only=True) + original = load_state_dict(str(source / "umt5/text_encoder/model.safetensors")) + model = UMT5EncoderModel.from_pretrained( + output / "text_encoder", dtype=next(iter(original.values())).dtype, local_files_only=True + ) + restored = model.state_dict() + for key, value in original.items(): + if restored[key].dtype != value.dtype or not torch.equal(restored[key], value): + raise ValueError(f"Text encoder round-trip changed {key}") + manifest["text_encoder"] = {key: tensor_record(value) for key, value in original.items()} + del original, restored, model + gc.collect() + + if size == "14b": + config = CLIPVisionConfig.from_dict(read_config(source / "clip/image_encoder/config.json")) + state = load_state_dict(str(source / "clip/image_encoder/model.safetensors")) + state.pop("visual_projection.weight") + with torch.device("meta"): + model = CLIPVisionModel(config) + expected_keys = set(model.state_dict()) + state = { + key if key in expected_keys else key.removeprefix("vision_model."): value for key, value in state.items() + } + save_component(model, state, output / "image_encoder", manifest) + shutil.copytree(source / "clip/image_processor", output / "image_processor", dirs_exist_ok=True) + del model, state + gc.collect() + + if size == "5b": + scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5.0) + timesteps = [1000, 937.5, 833.3333333333, 625] + else: + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(shift=3.0) + timesteps = [1000, 900, 750, 500] + if variant == "causal": + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True, default_timesteps=timesteps) + scheduler.save_pretrained(output / "scheduler") + index = { + "_class_name": "EdenLiveCausalPipeline" if variant == "causal" else "EdenLivePipeline", + "transformer": ["diffusers", "EdenLiveTransformer3DModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + "decoder": [None, None], + "text_encoder": ["transformers", "UMT5EncoderModel"], + "tokenizer": ["transformers", type(tokenizer).__name__], + "proxy_encoder": ["transformers", "DINOv3ViTModel"], + "image_encoder": ["transformers", "CLIPVisionModel"] if size == "14b" else [None, None], + "image_processor": ["transformers", "CLIPImageProcessor"] if size == "14b" else [None, None], + "scheduler": ["diffusers", type(scheduler).__name__], + "proxy_image_height": 480, + "proxy_image_width": 832, + } + if variant == "causal": + index.update(chunk_size=3, max_context_frames=18, sink_frames=6) + (output / "model_index.json").write_text(json.dumps(index, indent=2) + "\n") + for name in ("LICENSE", "NOTICE"): + if (source / name).is_file(): + shutil.copy2(source / name, output / name) + for component, destinations in ( + ("vision-encoder", ("proxy_encoder", "transformer")), + ("clip", ("image_encoder", "image_processor") if size == "14b" else ()), + ): + for destination in destinations: + for name in ("LICENSE.md", "NOTICE.md"): + shutil.copy2(source / component / name, output / destination / name) + (output / "conversion_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True, help="Local original EdenLive release directory.") + parser.add_argument( + "--output", type=Path, required=True, help="Destination for one standard Diffusers model root." + ) + parser.add_argument("--size", choices=("5b", "14b"), required=True) + parser.add_argument("--variant", choices=("bidirectional", "causal"), required=True) + args = parser.parse_args() + convert(args.source, args.output, args.size, args.variant) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 4f6790e5e4fa..f859454abb06 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -298,6 +298,9 @@ "DreamLiteTransformer2DModel", "DreamLiteUNetModel", "EasyAnimateTransformer3DModel", + "EdenLiveCausalState", + "EdenLiveLightTAEDecoder", + "EdenLiveTransformer3DModel", "ErnieImageTransformer2DModel", "Flux2Transformer2DModel", "Flux3Transformer2DModel", @@ -462,6 +465,7 @@ "DPMSolverMultistepInverseScheduler", "DPMSolverMultistepScheduler", "DPMSolverSinglestepScheduler", + "EdenLiveFlowMatchEulerDiscreteScheduler", "EDMDPMSolverMultistepScheduler", "EDMEulerScheduler", "EntropyBoundScheduler", @@ -680,6 +684,9 @@ "EasyAnimateControlPipeline", "EasyAnimateInpaintPipeline", "EasyAnimatePipeline", + "EdenLiveCausalPipeline", + "EdenLivePipeline", + "EdenLivePipelineOutput", "ErnieImagePipeline", "Flux2KleinInpaintPipeline", "Flux2KleinKVPipeline", @@ -1213,6 +1220,9 @@ DreamLiteTransformer2DModel, DreamLiteUNetModel, EasyAnimateTransformer3DModel, + EdenLiveCausalState, + EdenLiveLightTAEDecoder, + EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, Flux3Transformer2DModel, @@ -1373,6 +1383,7 @@ DPMSolverMultistepInverseScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, + EdenLiveFlowMatchEulerDiscreteScheduler, EDMDPMSolverMultistepScheduler, EDMEulerScheduler, EntropyBoundScheduler, @@ -1570,6 +1581,9 @@ EasyAnimateControlPipeline, EasyAnimateInpaintPipeline, EasyAnimatePipeline, + EdenLiveCausalPipeline, + EdenLivePipeline, + EdenLivePipelineOutput, ErnieImagePipeline, Flux2KleinInpaintPipeline, Flux2KleinKVPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 9e27e27e39ac..61b238152846 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -31,6 +31,7 @@ _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"] _import_structure["autoencoders.autoencoder_cosmos3_audio"] = ["Cosmos3AVAEAudioTokenizer"] _import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"] + _import_structure["autoencoders.autoencoder_edenlive"] = ["EdenLiveLightTAEDecoder"] _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"] _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"] @@ -120,6 +121,10 @@ _import_structure["transformers.transformer_cosmos"] = ["CosmosTransformer3DModel"] _import_structure["transformers.transformer_cosmos3"] = ["Cosmos3OmniTransformer"] _import_structure["transformers.transformer_easyanimate"] = ["EasyAnimateTransformer3DModel"] + _import_structure["transformers.transformer_edenlive"] = [ + "EdenLiveCausalState", + "EdenLiveTransformer3DModel", + ] _import_structure["transformers.transformer_ernie_image"] = ["ErnieImageTransformer2DModel"] _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"] _import_structure["transformers.transformer_flux2"] = ["Flux2Transformer2DModel"] @@ -216,6 +221,7 @@ AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, + EdenLiveLightTAEDecoder, Kandinsky6SRVAE, LongCatAudioDiTVae, LTX2VideoDiffusionDecoderModel, @@ -269,6 +275,8 @@ DreamLiteTransformer2DModel, DualTransformer2DModel, EasyAnimateTransformer3DModel, + EdenLiveCausalState, + EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, Flux3Transformer2DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index bb4abb06928b..a9664df7b8c0 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -1,6 +1,7 @@ from .autoencoder_asym_kl import AsymmetricAutoencoderKL from .autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from .autoencoder_dc import AutoencoderDC +from .autoencoder_edenlive import EdenLiveLightTAEDecoder from .autoencoder_kandinsky6_sr import Kandinsky6SRVAE from .autoencoder_kl import AutoencoderKL from .autoencoder_kl_allegro import AutoencoderKLAllegro diff --git a/src/diffusers/models/autoencoders/autoencoder_edenlive.py b/src/diffusers/models/autoencoders/autoencoder_edenlive.py new file mode 100644 index 000000000000..a91ac9f64b99 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_edenlive.py @@ -0,0 +1,176 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +from dataclasses import dataclass, field + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput + + +@dataclass +class _EdenLiveTAEDecodeState: + memory: dict[int, torch.Tensor] = field(default_factory=dict) + + +class EdenLiveTAEClamp(nn.Module): + def forward(self, hidden_states) -> torch.Tensor: + return torch.tanh(hidden_states / 3) * 3 + + +class EdenLiveTAEMemoryBlock(nn.Module): + def __init__(self, channels): + super().__init__() + self.conv = nn.Sequential( + nn.Conv2d(channels * 2, channels, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(channels, channels, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(channels, channels, 3, padding=1), + ) + self.skip = nn.Identity() + self.act = nn.ReLU(inplace=True) + + def forward(self, hidden_states, state, batch_size) -> torch.Tensor: + batch_frames, channels, height, width = hidden_states.shape + frames = hidden_states.reshape(batch_size, -1, channels, height, width) + previous = state.memory.get(id(self)) + if previous is None: + previous = torch.zeros_like(frames[:, :1]) + past = torch.cat([previous, frames[:, :-1]], dim=1).reshape(batch_frames, channels, height, width) + # Copy the final frame so a short tail does not retain the whole previous chunk. + state.memory[id(self)] = frames[:, -1:].clone() + return self.act(self.conv(torch.cat([hidden_states, past], dim=1)) + self.skip(hidden_states)) + + +class EdenLiveTAETemporalUpsample(nn.Module): + def __init__(self, channels, stride): + super().__init__() + self.conv = nn.Conv2d(channels, channels * stride, 1, bias=False) + + def forward(self, hidden_states) -> torch.Tensor: + _, channels, height, width = hidden_states.shape + return self.conv(hidden_states).reshape(-1, channels, height, width) + + +class EdenLiveLightTAEDecoder(ModelMixin, ConfigMixin): + """Decoder-only LightTAE for EdenLive's Wan latents. + + Input latents use the VAE's raw (unnormalized) scale and `[batch, channels, frames, height, width]` layout. Decoded + RGB samples use the same layout, range `[-1, 1]`, and `(frames - 1) * 4 + 1` frames. Temporal state belongs to each + `decode` call and is never retained on the model. + + Parameters: + latent_channels (`int`, *optional*, defaults to 16): + Latent channels: 16 for Wan 2.1 / EdenLive 14B, or 48 for Wan 2.2 / EdenLive 5B. + patch_size (`int`, *optional*, defaults to 1): + Output pixel shuffle factor: 1 for 14B or 2 for 5B. + block_out_channels (`tuple[int]`, *optional*, defaults to `(256, 128, 64, 64)`): + Decoder widths at its four spatial resolutions. + """ + + _no_split_modules = ["EdenLiveTAEMemoryBlock", "EdenLiveTAETemporalUpsample"] + + @register_to_config + def __init__(self, latent_channels=16, patch_size=1, block_out_channels=(256, 128, 64, 64)): + super().__init__() + if latent_channels < 1 or patch_size not in (1, 2): + raise ValueError("LightTAE requires positive latent channels and patch_size of 1 or 2.") + if len(block_out_channels) != 4 or any(channels < 1 for channels in block_out_channels): + raise ValueError("`block_out_channels` must contain four positive widths.") + widths = block_out_channels + layers = [EdenLiveTAEClamp(), nn.Conv2d(latent_channels, widths[0], 3, padding=1), nn.ReLU(inplace=True)] + for stage in range(3): + layers.extend(EdenLiveTAEMemoryBlock(widths[stage]) for _ in range(3)) + layers.extend( + [ + nn.Upsample(scale_factor=2), + EdenLiveTAETemporalUpsample(widths[stage], 1 if stage == 0 else 2), + nn.Conv2d(widths[stage], widths[stage + 1], 3, padding=1, bias=False), + ] + ) + layers.extend([nn.ReLU(inplace=True), nn.Conv2d(widths[3], 3 * patch_size**2, 3, padding=1)]) + self.decoder = nn.Sequential(*layers) + + @property + def scale_factor_temporal(self): + return 4 + + @property + def scale_factor_spatial(self): + return 8 * self.config.patch_size + + def _decode_chunk(self, latents, state): + first_chunk = not state.memory + batch_size, _, num_frames, height, width = latents.shape + hidden_states = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, -1, height, width) + for layer in self.decoder: + if isinstance(layer, EdenLiveTAEMemoryBlock): + hidden_states = layer(hidden_states, state, batch_size) + else: + hidden_states = layer(hidden_states) + hidden_states = hidden_states.clamp(0, 1) + if self.config.patch_size > 1: + hidden_states = F.pixel_shuffle(hidden_states, self.config.patch_size) + frames = hidden_states.reshape(batch_size, -1, *hidden_states.shape[1:]) + if first_chunk: + frames = frames[:, self.scale_factor_temporal - 1 :] + return frames.permute(0, 2, 1, 3, 4) * 2 - 1 + + @apply_forward_hook + def decode(self, latents: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: + """Decode a complete video from raw VAE latents. + + Args: + latents (`torch.Tensor`): + Unnormalized VAE latents with shape `(B, C, T_latent, H_latent, W_latent)` on the decoder device and in + its dtype. `C` must equal `config.latent_channels`, and at least one latent frame is required. Pipeline + outputs with `output_type="latent"` must first be unnormalized with the VAE mean and standard + deviation. The pipeline performs this conversion automatically when decoding frames. + return_dict (`bool`, defaults to `True`): + Return a [`DecoderOutput`] when true; otherwise return a one-element tuple containing the sample. + + Returns: + [`DecoderOutput`] or `tuple[torch.Tensor]`: + RGB sample in `(B, 3, T, H, W)` layout and [-1, 1] range. The output has `(T_latent - 1) * + scale_factor_temporal + 1` frames and spatial dimensions multiplied by `scale_factor_spatial`. Internal + chunks share temporal state, including a short final chunk. + """ + if latents.ndim != 5 or latents.shape[1] != self.config.latent_channels or latents.shape[2] < 1: + raise ValueError("Expected nonempty NCTHW latents with the configured latent channel count.") + # Keep one state object across hooks that rebuild dictionary and list arguments when moving tensors. + state = _EdenLiveTAEDecodeState() + sample = torch.cat([self._decode_chunk(chunk, state) for chunk in latents.split(3, dim=2)], dim=2) + if not return_dict: + return (sample,) + return DecoderOutput(sample=sample) + + def forward(self, latents: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: + """ + Args: + latents (`torch.Tensor`): + Raw video latents in `[batch, channels, frames, height, width]` layout. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`DecoderOutput`] instead of a tuple containing the decoded sample. + + Returns: + [`DecoderOutput`] or `tuple[torch.Tensor]`: + The decoded RGB video in NCTHW layout and `[-1, 1]` range. + """ + return self.decode(latents, return_dict=return_dict) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 1076b2b33fae..876353d7c842 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -31,6 +31,7 @@ from .transformer_cosmos import CosmosTransformer3DModel from .transformer_cosmos3 import Cosmos3OmniTransformer from .transformer_easyanimate import EasyAnimateTransformer3DModel + from .transformer_edenlive import EdenLiveCausalState, EdenLiveTransformer3DModel from .transformer_ernie_image import ErnieImageTransformer2DModel from .transformer_flux import FluxTransformer2DModel from .transformer_flux2 import Flux2Transformer2DModel diff --git a/src/diffusers/models/transformers/transformer_edenlive.py b/src/diffusers/models/transformers/transformer_edenlive.py new file mode 100644 index 000000000000..10b091d763f1 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_edenlive.py @@ -0,0 +1,908 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import math +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import apply_lora_scale +from ...utils.torch_utils import get_module_execution_device, maybe_adjust_dtype_for_device, maybe_allow_in_graph +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +def _apply_rotary_emb(hidden_states, rotary_emb): + freqs_cos, freqs_sin = rotary_emb + first, second = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cos[..., 0::2], freqs_sin[..., 1::2] + output = torch.empty_like(hidden_states) + output[..., 0::2] = first * cos - second * sin + output[..., 1::2] = first * sin + second * cos + return output + + +def _get_qkv_projections(attn, hidden_states, encoder_hidden_states=None): + context = hidden_states if encoder_hidden_states is None else encoder_hidden_states + if attn.fused_projections: + if not attn.is_cross_attention: + return attn.to_qkv(hidden_states).chunk(3, dim=-1) + query = attn.to_q(hidden_states) + key, value = attn.to_kv(context).chunk(2, dim=-1) + return query, key, value + return attn.to_q(hidden_states), attn.to_k(context), attn.to_v(context) + + +def _get_added_kv_projections(attn, image_context): + if attn.fused_projections: + return attn.to_added_kv(image_context).chunk(2, dim=-1) + return attn.add_k_proj(image_context), attn.add_v_proj(image_context) + + +@dataclass +class _EdenLiveLayerCache: + key: torch.Tensor + value: torch.Tensor + num_tokens: int = 0 + + +@dataclass +class EdenLiveCausalState: + """KV storage for one fixed-resolution rollout and guidance branch. Positions count latent frames.""" + + layers: list[_EdenLiveLayerCache] + chunk_size: int + max_context_frames: int + sink_frames: int + next_frame: int = 0 + valid: bool = True + + def reset(self) -> None: + """Clear committed token counts and reset the frame cursor and validity, retaining allocated cache storage.""" + for layer in self.layers: + layer.num_tokens = 0 + self.next_frame = 0 + self.valid = True + + +class EdenLiveAttnProcessor: + _attention_backend = None + + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + *, + text_context_length=512, + ): + image_context = None + if attn.add_k_proj is not None: + if encoder_hidden_states is None: + raise ValueError("Image cross-attention requires image and text embeddings.") + image_length = encoder_hidden_states.shape[1] - text_context_length + if image_length <= 0: + raise ValueError("The 14B model requires CLIP image embeddings before the text embeddings.") + image_context = encoder_hidden_states[:, :image_length] + encoder_hidden_states = encoder_hidden_states[:, image_length:] + + query, key, value = _get_qkv_projections(attn, hidden_states, encoder_hidden_states) + query = attn.norm_q(query).unflatten(2, (attn.heads, -1)) + key = attn.norm_k(key).unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + image_output = None + if image_context is not None: + image_key, image_value = _get_added_kv_projections(attn, image_context) + image_key = attn.norm_added_k(image_key).unflatten(2, (attn.heads, -1)) + image_value = image_value.unflatten(2, (attn.heads, -1)) + image_output = ( + dispatch_attention_fn( + query, + image_key, + image_value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) + .flatten(2, 3) + .type_as(query) + ) + + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) + output = output.flatten(2, 3).type_as(query) + if image_output is not None: + output = output + image_output + return attn.to_out[1](attn.to_out[0](output)) + + +class EdenLiveCausalAttnProcessor: + _attention_backend = None + + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + rotary_emb=None, + *, + layer_cache=None, + key_rotary_emb=None, + cache_update=False, + cache_window_tokens=None, + cache_sink_tokens=0, + ): + if encoder_hidden_states is not None or attention_mask is not None: + raise ValueError("The causal processor only supports unmasked self-attention.") + + query, key, value = _get_qkv_projections(attn, hidden_states) + query = attn.norm_q(query).unflatten(2, (attn.heads, -1)) + key = attn.norm_k(key).unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + if rotary_emb is not None: + query = _apply_rotary_emb(query, rotary_emb) + if layer_cache is None: + if rotary_emb is not None: + key = _apply_rotary_emb(key, rotary_emb) + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) + output = output.flatten(2, 3).type_as(query) + return attn.to_out[1](attn.to_out[0](output)) + new_tokens = key.shape[1] + capacity = layer_cache.key.shape[1] + stored_end = layer_cache.num_tokens + if key.dtype != layer_cache.key.dtype or key.device != layer_cache.key.device: + raise ValueError("KV cache and attention must have the same device and dtype; recreate the causal state.") + if new_tokens > capacity: + raise ValueError("The current chunk exceeds the allocated KV cache.") + + if cache_update: + if stored_end + new_tokens > capacity: + sink = min(cache_sink_tokens, capacity - new_tokens) + keep = capacity - sink - new_tokens + recent_start = max(sink, stored_end - keep) + retained = stored_end - recent_start + if retained: + layer_cache.key[:, sink : sink + retained].copy_( + layer_cache.key[:, recent_start:stored_end].clone() + ) + layer_cache.value[:, sink : sink + retained].copy_( + layer_cache.value[:, recent_start:stored_end].clone() + ) + stored_end = sink + retained + layer_cache.key[:, stored_end : stored_end + new_tokens].copy_(key.detach()) + layer_cache.value[:, stored_end : stored_end + new_tokens].copy_(value.detach()) + stored_end += new_tokens + layer_cache.num_tokens = stored_end + sink = min(cache_sink_tokens, stored_end) + recent_start = max(sink, stored_end - (cache_window_tokens - sink)) + key = torch.cat((layer_cache.key[:, :sink], layer_cache.key[:, recent_start:stored_end]), dim=1) + value = torch.cat((layer_cache.value[:, :sink], layer_cache.value[:, recent_start:stored_end]), dim=1) + else: + sink = min(cache_sink_tokens, cache_window_tokens - new_tokens, stored_end) + recent_start = max(sink, stored_end - (cache_window_tokens - sink - new_tokens)) + key = torch.cat((layer_cache.key[:, :sink], layer_cache.key[:, recent_start:stored_end], key), dim=1) + value = torch.cat( + (layer_cache.value[:, :sink], layer_cache.value[:, recent_start:stored_end], value), dim=1 + ) + + key = _apply_rotary_emb(key, key_rotary_emb) + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) + output = output.flatten(2, 3).type_as(query) + return attn.to_out[1](attn.to_out[0](output)) + + +class EdenLivePatchEmbedding(nn.Conv3d): + def __init__(self, in_channels, out_channels, patch_size, proxy_channels, condition_rmsnorm): + super().__init__(in_channels + proxy_channels, out_channels, kernel_size=patch_size, stride=patch_size) + self.proxy_channels = proxy_channels + self.condition_norm = nn.RMSNorm(proxy_channels, eps=1e-6) if condition_rmsnorm else None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.condition_norm is not None: + base, condition = hidden_states[:, : -self.proxy_channels], hidden_states[:, -self.proxy_channels :] + condition = condition.permute(0, 2, 3, 4, 1) + condition = self.condition_norm(condition.float()).to(hidden_states.dtype) + hidden_states = torch.cat((base, condition.permute(0, 4, 1, 2, 3)), dim=1) + return super().forward(hidden_states) + + +class EdenLiveAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = EdenLiveAttnProcessor + _available_processors = [EdenLiveAttnProcessor, EdenLiveCausalAttnProcessor] + + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + eps: float = 1e-5, + dropout: float = 0.0, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, + processor=None, + ): + super().__init__() + + self.inner_dim = dim_head * heads + self.heads = heads + self.added_kv_proj_dim = added_kv_proj_dim + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + + self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = torch.nn.ModuleList( + [ + torch.nn.Linear(self.inner_dim, dim, bias=True), + torch.nn.Dropout(dropout), + ] + ) + self.norm_q = torch.nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = torch.nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + + self.add_k_proj = self.add_v_proj = None + if added_kv_proj_dim is not None: + self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps) + + self.is_cross_attention = cross_attention_dim_head is not None + + self.set_processor(processor or EdenLiveAttnProcessor()) + + # Copied from diffusers.models.transformers.transformer_wan.WanAttention.fuse_projections + def fuse_projections(self): + if getattr(self, "fused_projections", False): + return + + if not self.is_cross_attention: + concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_qkv = nn.Linear(in_features, out_features, bias=True) + self.to_qkv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + else: + concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_kv = nn.Linear(in_features, out_features, bias=True) + self.to_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + if self.added_kv_proj_dim is not None: + concatenated_weights = torch.cat([self.add_k_proj.weight.data, self.add_v_proj.weight.data]) + concatenated_bias = torch.cat([self.add_k_proj.bias.data, self.add_v_proj.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_added_kv = nn.Linear(in_features, out_features, bias=True) + self.to_added_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + self.fused_projections = True + + @torch.no_grad() + # Copied from diffusers.models.transformers.transformer_wan.WanAttention.unfuse_projections + def unfuse_projections(self): + if not getattr(self, "fused_projections", False): + return + + if hasattr(self, "to_qkv"): + delattr(self, "to_qkv") + if hasattr(self, "to_kv"): + delattr(self, "to_kv") + if hasattr(self, "to_added_kv"): + delattr(self, "to_added_kv") + + self.fused_projections = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, **kwargs) + + +class EdenLiveImageEmbedding(torch.nn.Module): + def __init__(self, in_features: int, out_features: int): + super().__init__() + + self.norm1 = FP32LayerNorm(in_features) + self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu") + self.norm2 = FP32LayerNorm(out_features) + + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm1(encoder_hidden_states_image) + hidden_states = self.ff(hidden_states) + hidden_states = self.norm2(hidden_states) + return hidden_states + + +class EdenLiveTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + image_embed_dim: int | None = None, + ): + super().__init__() + + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = EdenLiveImageEmbedding(image_embed_dim, dim) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + timestep = self.timesteps_proj(timestep) + if timestep_seq_len is not None: + timestep = timestep.unflatten(0, (-1, timestep_seq_len)) + + timestep = timestep.to(dtype=encoder_hidden_states.dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image is not None: + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.cat((encoder_hidden_states_image, encoder_hidden_states), dim=1) + return temb, timestep_proj, encoder_hidden_states + + +class EdenLiveRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + + self.t_dim = t_dim + self.h_dim = h_dim + self.w_dim = w_dim + + freqs_dtype = maybe_adjust_dtype_for_device(torch.float64, torch.empty(0).device) + + freqs_cos = [] + freqs_sin = [] + + for dim in [t_dim, h_dim, w_dim]: + freq_cos, freq_sin = get_1d_rotary_pos_embed( + dim, + max_seq_len, + theta, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + freqs_cos.append(freq_cos) + freqs_sin.append(freq_sin) + + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward( + self, hidden_states: torch.Tensor, start_frame: int = 0, num_frames: int | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, input_frames, height, width = hidden_states.shape + num_frames = input_frames if num_frames is None else num_frames + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w + + if start_frame + ppf > self.max_seq_len or max(pph, ppw) > self.max_seq_len: + raise ValueError("The requested rotary position exceeds rope_max_seq_len.") + split_sizes = [self.t_dim, self.h_dim, self.w_dim] + + freqs_cos = self.freqs_cos.to(device=hidden_states.device, dtype=hidden_states.dtype).split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.to(device=hidden_states.device, dtype=hidden_states.dtype).split(split_sizes, dim=1) + + freqs_cos_f = freqs_cos[0][start_frame : start_frame + ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_sin_f = freqs_sin[0][start_frame : start_frame + ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + + return freqs_cos, freqs_sin + + +@maybe_allow_in_graph +class EdenLiveTransformerBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, eps=1e-6, added_kv_proj_dim=None): + super().__init__() + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = EdenLiveAttention( + dim=dim, heads=num_heads, dim_head=dim // num_heads, eps=eps, processor=EdenLiveCausalAttnProcessor() + ) + self.attn2 = EdenLiveAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + added_kv_proj_dim=added_kv_proj_dim, + cross_attention_dim_head=dim // num_heads, + processor=EdenLiveAttnProcessor(), + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states, + encoder_hidden_states, + temb, + rotary_emb, + *, + layer_cache=None, + key_rotary_emb=None, + cache_update=False, + cache_window_tokens=None, + cache_sink_tokens=0, + text_context_length=512, + ) -> torch.Tensor: + if temb.ndim == 4: + modulation = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + tensor.squeeze(2) for tensor in modulation + ) + else: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb.float() + ).chunk(6, dim=1) + + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.attn1( + norm_hidden_states, + rotary_emb=rotary_emb, + layer_cache=layer_cache, + key_rotary_emb=key_rotary_emb, + cache_update=cache_update, + cache_window_tokens=cache_window_tokens, + cache_sink_tokens=cache_sink_tokens, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states, + text_context_length=text_context_length, + ) + hidden_states = hidden_states + attn_output + norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + ff_output = self.ffn(norm_hidden_states) + return (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + +class EdenLiveTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): + r""" + Proxy-conditioned video transformer with bidirectional and chunk-causal inference. + + Args: + patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): + Temporal and spatial patch dimensions. The temporal patch size must be one. + num_attention_heads (`int`, defaults to `40`): + Number of attention heads. + attention_head_dim (`int`, defaults to `128`): + Channels per attention head. + in_channels (`int`, defaults to `36`): + Input latent channels before appending proxy features. This includes mask and image latent channels for + 14B. + out_channels (`int`, defaults to `16`): + Output latent channels. + text_dim (`int`, defaults to `4096`): + Input text embedding dimension. + freq_dim (`int`, defaults to `256`): + Sinusoidal timestep embedding dimension. + ffn_dim (`int`, defaults to `13824`): + Feed-forward intermediate dimension. + num_layers (`int`, defaults to `40`): + Number of transformer blocks. + qk_norm (`str`, defaults to `"rms_norm_across_heads"`): + Query and key normalization. Only the checkpoint's normalization across heads is supported. + eps (`float`, defaults to `1e-6`): + Normalization epsilon. + image_dim (`int`, *optional*): + Input CLIP image embedding dimension for the 14B model. + added_kv_proj_dim (`int`, *optional*): + Dimension of projected image embeddings consumed by image cross-attention. + rope_max_seq_len (`int`, defaults to `1024`): + Maximum rotary position on each temporal or spatial axis. + proxy_channels (`int`, defaults to `1024`): + Projected DINO feature channels appended before patch embedding. + proxy_input_channels (`int`, defaults to `1024`): + DINO feature channels before the learned proxy projection. + condition_rmsnorm (`bool`, defaults to `False`): + Whether to RMS-normalize proxy features before patch embedding, as required by the 14B checkpoint. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["EdenLiveTransformerBlock", "EdenLiveTimeTextImageEmbedding", "EdenLivePatchEmbedding"] + _repeated_blocks = ["EdenLiveTransformerBlock"] + _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] + _keep_in_fp32_modules = ["proxy_projection"] + _skip_keys = ["causal_state"] + + @register_to_config + def __init__( + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 36, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + qk_norm: str = "rms_norm_across_heads", + eps: float = 1e-6, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, + rope_max_seq_len: int = 1024, + proxy_channels: int = 1024, + proxy_input_channels: int = 1024, + condition_rmsnorm: bool = False, + ): + super().__init__() + if len(patch_size) != 3 or patch_size[0] != 1 or min(patch_size) < 1: + raise ValueError("EdenLive requires temporal patch size 1 and positive spatial patch sizes.") + if qk_norm != "rms_norm_across_heads": + raise ValueError("EdenLive checkpoints use rms_norm_across_heads.") + if min(proxy_input_channels, proxy_channels, num_layers) < 1: + raise ValueError("proxy_input_channels, proxy_channels and num_layers must be positive.") + if (image_dim is None) != (added_kv_proj_dim is None): + raise ValueError("image_dim and added_kv_proj_dim must both be set for the 14B image-conditioned model.") + inner_dim = num_attention_heads * attention_head_dim + self.rope = EdenLiveRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.cache_rope = EdenLiveRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.proxy_projection = nn.Linear(proxy_input_channels, proxy_channels) + self.patch_embedding = EdenLivePatchEmbedding( + in_channels, inner_dim, patch_size, proxy_channels, condition_rmsnorm + ) + self.condition_embedder = EdenLiveTimeTextImageEmbedding( + inner_dim, freq_dim, inner_dim * 6, text_dim, image_dim + ) + self.blocks = nn.ModuleList( + [ + EdenLiveTransformerBlock(inner_dim, ffn_dim, num_attention_heads, eps, added_kv_proj_dim) + for _ in range(num_layers) + ] + ) + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + self.gradient_checkpointing = False + + def create_causal_state( + self, + batch_size: int, + height: int, + width: int, + chunk_size: int = 3, + max_context_frames: int = 18, + sink_frames: int = 6, + device=None, + dtype=None, + ) -> EdenLiveCausalState: + r""" + Allocate independent KV storage for one guidance branch and rollout with fixed latent spatial dimensions. + + Args: + batch_size (`int`): Batch size of the model inputs. + height (`int`): Height of the latent input, before patch embedding. + width (`int`): Width of the latent input, before patch embedding. + chunk_size (`int`, defaults to `3`): Number of latent frames per full chunk. + max_context_frames (`int`, defaults to `18`): Total context, including sink and current chunk. + sink_frames (`int`, defaults to `6`): Initial latent frames retained when the context fills. + device (`torch.device` or `str`, *optional*): + Cache device for a model without `device_map`; defaults to the model execution device. With + `device_map`, each layer's execution device takes precedence. + dtype (`torch.dtype`, *optional*): Cache dtype; defaults to the model dtype. + + Returns: + `EdenLiveCausalState`: Empty KV storage for one rollout and guidance branch. + """ + _, patch_height, patch_width = self.config.patch_size + if min(batch_size, height, width, chunk_size, max_context_frames) < 1: + raise ValueError("Batch size, latent dimensions, chunk size and context size must be positive.") + if height % patch_height or width % patch_width: + raise ValueError("Latent dimensions must be divisible by the patch size.") + if not 0 <= sink_frames < max_context_frames or max_context_frames < chunk_size: + raise ValueError("The context must contain the current chunk and be larger than its nonnegative sink.") + if max_context_frames % chunk_size or sink_frames % chunk_size: + raise ValueError("The fixed context and sink must contain whole chunks.") + tokens_per_frame = (height // patch_height) * (width // patch_width) + shape = ( + batch_size, + max_context_frames * tokens_per_frame, + self.config.num_attention_heads, + self.config.attention_head_dim, + ) + layers = [] + for block in self.blocks: + cache_device = ( + get_module_execution_device(block) + if getattr(self, "hf_device_map", None) + else device or get_module_execution_device(self) + ) + layers.append( + _EdenLiveLayerCache( + torch.zeros(shape, device=cache_device, dtype=dtype or self.dtype), + torch.zeros(shape, device=cache_device, dtype=dtype or self.dtype), + ) + ) + return EdenLiveCausalState( + layers, + chunk_size, + max_context_frames, + sink_frames, + ) + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + proxy_features: torch.Tensor | None = None, + causal_state: EdenLiveCausalState | None = None, + cache_update: bool = False, + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: + r""" + Denoise one clip or chunk. A clean-chunk call with `cache_update=True` advances the causal state. + + Args: + hidden_states (`torch.Tensor`): + Input latents of shape `(batch_size, in_channels, num_frames, height, width)`, including the image + conditioning channels for 14B. + timestep (`torch.Tensor`): + Diffusion timestep, as a scalar, `(batch_size,)`, `(batch_size, num_frames)`, or `(batch_size, + num_patch_tokens)`. Preserve float32 timestep precision when using reduced-precision model weights. + encoder_hidden_states (`torch.Tensor`): + Text embeddings of shape `(batch_size, sequence_length, text_dim)`. + encoder_hidden_states_image (`torch.Tensor`, *optional*): + CLIP image embeddings required by the 14B model. + proxy_features (`torch.Tensor`, *optional*): + Required DINO features of shape + `(batch_size, proxy_input_channels, num_frames, proxy_height, proxy_width)`. Features are projected and + normalized before their spatial dimensions are interpolated to match the latent input. + causal_state (`EdenLiveCausalState`, *optional*): + KV storage created by `create_causal_state`. Use separate states for different guidance branches. + cache_update (`bool`, defaults to `False`): + Whether to commit a clean chunk to the KV state and advance its frame cursor. Denoising calls leave the + cache unchanged. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output instead of a tuple. + attention_kwargs (`dict[str, Any]`, *optional*): + Attention options, including the LoRA `scale` when using PEFT adapters. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + The predicted latent tensor, with shape `(batch_size, out_channels, num_frames, height, width)`. + """ + if hidden_states.ndim != 5 or hidden_states.shape[1] != self.config.in_channels: + raise ValueError(f"hidden_states must have shape [B, {self.config.in_channels}, T, H, W].") + batch_size, _, num_frames, height, width = hidden_states.shape + patch_frames, patch_height, patch_width = self.config.patch_size + if num_frames < 1 or height % patch_height or width % patch_width: + raise ValueError( + "Input spatial dimensions must be divisible by the patch size and the clip must be nonempty." + ) + if proxy_features is None or proxy_features.ndim != 5: + raise ValueError("proxy_features must be supplied as a [B, C, T, H, W] tensor.") + if proxy_features.shape[:2] != (batch_size, self.config.proxy_input_channels): + raise ValueError("Proxy batch size and channels must match the configured model.") + if proxy_features.shape[2] != num_frames: + raise ValueError("Proxy features must have one control frame per latent frame.") + if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch_size: + raise ValueError("encoder_hidden_states must have shape [B, sequence_length, text_dim].") + if self.config.image_dim is not None and encoder_hidden_states_image is None: + raise ValueError("This image-conditioned model requires CLIP encoder_hidden_states_image.") + if self.config.image_dim is None and encoder_hidden_states_image is not None: + raise ValueError("The 5B model does not use CLIP image embeddings.") + + tokens_per_frame = (height // patch_height) * (width // patch_width) + num_tokens = num_frames * tokens_per_frame + text_context_length = encoder_hidden_states.shape[1] + key_rotary_emb = None + cache_window_tokens = None + cache_sink_tokens = 0 + query_start_frame = 0 + if cache_update and causal_state is None: + raise ValueError("cache_update requires a causal state.") + if causal_state is not None: + state = causal_state + if not state.valid: + raise ValueError("A previous cache commit failed. Reset or recreate the causal state.") + cache_shape = ( + batch_size, + state.max_context_frames * tokens_per_frame, + self.config.num_attention_heads, + self.config.attention_head_dim, + ) + if len(state.layers) != len(self.blocks) or any( + layer.key.shape != cache_shape or layer.value.shape != cache_shape for layer in state.layers + ): + raise ValueError("The causal state does not match the model's batch size or token dimensions.") + if state.next_frame % state.chunk_size or num_frames > state.chunk_size: + raise ValueError("Causal calls must follow whole chunks; a shorter chunk can only finish a rollout.") + if torch.is_grad_enabled() and self.gradient_checkpointing: + raise ValueError("Mutable causal inference state cannot be used with gradient checkpointing.") + query_start_frame = min(state.next_frame, state.max_context_frames - state.chunk_size) + visible_frames = query_start_frame + num_frames + cache_window_tokens = visible_frames * tokens_per_frame + cache_sink_tokens = state.sink_frames * tokens_per_frame + key_rotary_emb = self.cache_rope(hidden_states, num_frames=visible_frames) + + rotary_emb = self.rope(hidden_states, start_frame=query_start_frame) + projected_features = [] + for frame in proxy_features.unbind(2): + frame = frame.permute(0, 2, 3, 1).contiguous().to(device=hidden_states.device, dtype=torch.float32) + with torch.autocast(hidden_states.device.type, dtype=torch.bfloat16, enabled=hidden_states.is_cuda): + projected = self.proxy_projection(frame) + projected = F.normalize(projected, p=2, dim=-1).float() + projected_features.append(projected.permute(0, 3, 1, 2)) + proxy_features = torch.stack(projected_features, dim=2) + proxy_features = proxy_features.to(device=hidden_states.device, dtype=hidden_states.dtype) + if proxy_features.shape[2:] != hidden_states.shape[2:]: + proxy_features = F.interpolate( + proxy_features, size=hidden_states.shape[2:], mode="trilinear", align_corners=False + ) + hidden_states = self.patch_embedding(torch.cat((hidden_states, proxy_features), dim=1)) + hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous() + + if timestep.ndim == 0: + timestep = timestep.expand(batch_size) + if timestep.ndim == 1 and timestep.shape[0] == 1: + timestep = timestep.expand(batch_size) + if timestep.ndim not in (1, 2) or timestep.shape[0] != batch_size: + raise ValueError("timestep must be scalar, [B], [B, latent_frames], or [B, patch_tokens].") + timestep_seq_len = None + if timestep.ndim == 2: + if timestep.shape[1] == num_frames: + timestep = timestep.repeat_interleave(tokens_per_frame, dim=1) + if timestep.shape[1] != num_tokens: + raise ValueError("Per-token timesteps must match the latent frame count or patch token count.") + timestep_seq_len = num_tokens + timestep = timestep.flatten() + temb, timestep_proj, encoder_hidden_states = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=timestep_seq_len + ) + timestep_proj = timestep_proj.unflatten(2 if timestep_seq_len is not None else 1, (6, -1)) + if cache_update: + causal_state.valid = False + for index, block in enumerate(self.blocks): + kwargs = { + "layer_cache": causal_state.layers[index] if causal_state is not None else None, + "key_rotary_emb": key_rotary_emb, + "cache_update": cache_update, + "cache_window_tokens": cache_window_tokens, + "cache_sink_tokens": cache_sink_tokens, + "text_context_length": text_context_length, + } + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + partial(block, **kwargs), hidden_states, encoder_hidden_states, timestep_proj, rotary_emb + ) + else: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, **kwargs) + + if temb.ndim == 3: + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = shift.squeeze(2), scale.squeeze(2) + else: + shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) + hidden_states_dtype = hidden_states.dtype + hidden_states = self.norm_out(hidden_states.float()) + hidden_states = (hidden_states * (1 + scale.to(hidden_states.device)) + shift.to(hidden_states.device)).to( + hidden_states_dtype + ) + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape( + batch_size, + num_frames, + height // patch_height, + width // patch_width, + patch_frames, + patch_height, + patch_width, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + if cache_update: + causal_state.next_frame += num_frames + causal_state.valid = True + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 4f1f600e4faf..d0d8142573c1 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -274,6 +274,7 @@ ] _import_structure["diffusion_gemma"] = ["DiffusionGemmaPipeline", "DiffusionGemmaPipelineOutput"] _import_structure["dreamlite"] = ["DreamLitePipeline", "DreamLiteMobilePipeline", "DreamLitePipelineOutput"] + _import_structure["edenlive"] = ["EdenLiveCausalPipeline", "EdenLivePipeline", "EdenLivePipelineOutput"] _import_structure["easyanimate"] = [ "EasyAnimatePipeline", "EasyAnimateInpaintPipeline", @@ -724,6 +725,7 @@ EasyAnimateInpaintPipeline, EasyAnimatePipeline, ) + from .edenlive import EdenLiveCausalPipeline, EdenLivePipeline, EdenLivePipelineOutput from .ernie_image import ErnieImagePipeline from .flux import ( FluxControlImg2ImgPipeline, diff --git a/src/diffusers/pipelines/edenlive/__init__.py b/src/diffusers/pipelines/edenlive/__init__.py new file mode 100644 index 000000000000..d54aa3d6ea93 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/__init__.py @@ -0,0 +1,45 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_edenlive"] = ["EdenLivePipeline"] + _import_structure["pipeline_edenlive_causal"] = ["EdenLiveCausalPipeline"] + _import_structure["pipeline_output"] = ["EdenLivePipelineOutput"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_edenlive import EdenLivePipeline + from .pipeline_edenlive_causal import EdenLiveCausalPipeline + from .pipeline_output import EdenLivePipelineOutput +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py new file mode 100644 index 000000000000..7cb5e727b8cf --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -0,0 +1,500 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import html +import math +import re +from typing import Callable + +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel + +from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveTransformer3DModel +from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler +from ...utils import is_ftfy_available, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import EdenLivePipelineOutput +from .pipeline_utils import prepare_model_input +from .video_processor import EdenLiveVideoProcessor + + +if is_ftfy_available(): + import ftfy + + +# docstyle-ignore +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import EdenLivePipeline + >>> from diffusers.utils import export_to_video, load_image, load_video + >>> pipe = EdenLivePipeline.from_pretrained( + ... "VAST-AI/EdenLive-5B-Diffusers", + ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, + ... ).to("cuda") + >>> frames = pipe( + ... image=load_image("reference.png"), video=load_video("proxy.mp4"), + ... prompt="A character walks through snow.", num_frames=49, + ... ).frames[0] + >>> export_to_video(frames, "edenlive.mp4", fps=24) + ``` +""" + + +class EdenLivePipeline(DiffusionPipeline): + """Bidirectional video generation conditioned on a reference image, text, and an RGB proxy video. + + Args: + vae ([`AutoencoderKLWan`]): + Wan VAE for reference encoding and full-quality video decoding. + text_encoder (`UMT5EncoderModel`): + Text encoder. + tokenizer (`AutoTokenizer`): + Tokenizer corresponding to the text encoder. + transformer ([`EdenLiveTransformer3DModel`]): + EdenLive bidirectional video denoiser. + scheduler ([`EdenLiveFlowMatchEulerDiscreteScheduler`] or [`UniPCMultistepScheduler`]): + Scheduler corresponding to the transformer checkpoint. + proxy_encoder (`DINOv3ViTModel`): + Trained RGB proxy encoder. + image_encoder (`CLIPVisionModel`, *optional*): + Reference image encoder required by the 14B model. + image_processor (`CLIPImageProcessor`, *optional*): + Processor corresponding to the image encoder. + decoder ([`EdenLiveLightTAEDecoder`], *optional*): + Optional faster decoder. When omitted, decode with the full Wan VAE. + proxy_image_height (`int`, defaults to 480): + Height of the proxy images passed to DINO. + proxy_image_width (`int`, defaults to 832): + Width of the proxy images passed to DINO. + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->proxy_encoder->transformer->vae->decoder" + _optional_components = ["image_encoder", "image_processor", "decoder"] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + vae: AutoencoderKLWan, + text_encoder: UMT5EncoderModel, + tokenizer: AutoTokenizer, + transformer: EdenLiveTransformer3DModel, + scheduler: EdenLiveFlowMatchEulerDiscreteScheduler | UniPCMultistepScheduler, + proxy_encoder: DINOv3ViTModel, + image_encoder: CLIPVisionModel | None = None, + image_processor: CLIPImageProcessor | None = None, + decoder: EdenLiveLightTAEDecoder | None = None, + proxy_image_height: int = 480, + proxy_image_width: int = 832, + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + proxy_encoder=proxy_encoder, + image_encoder=image_encoder, + image_processor=image_processor, + decoder=decoder, + ) + self.register_to_config(proxy_image_height=proxy_image_height, proxy_image_width=proxy_image_width) + self._image_conditioned = transformer is not None and transformer.config.image_dim is not None + self.vae_scale_factor_temporal = vae.config.scale_factor_temporal if vae is not None else 4 + self.vae_scale_factor_spatial = vae.config.scale_factor_spatial if vae is not None else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + self.proxy_video_processor = EdenLiveVideoProcessor( + vae_scale_factor=1, do_normalize=False, do_convert_rgb=True + ) + + def _get_t5_prompt_embeds(self, texts, max_sequence_length, device, dtype): + if is_ftfy_available(): + texts = [ftfy.fix_text(text) for text in texts] + texts = [re.sub(r"\s+", " ", html.unescape(html.unescape(text))).strip() for text in texts] + inputs = self.tokenizer( + texts, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + mask = inputs.attention_mask.to(device) + embeddings = self.text_encoder(inputs.input_ids.to(device), attention_mask=mask).last_hidden_state + embeddings = embeddings.to(device=device, dtype=dtype) + return embeddings.masked_fill(~mask.bool().unsqueeze(-1), 0) + + def encode_prompt( + self, + prompt=None, + negative_prompt=None, + do_classifier_free_guidance=True, + prompt_embeds=None, + negative_prompt_embeds=None, + max_sequence_length=512, + device=None, + dtype=None, + ): + """Encode positive and optional negative text with UMT5, zeroing padded token embeddings.""" + device = device or self._execution_device + dtype = dtype or (self.transformer.dtype if self.transformer is not None else self.text_encoder.dtype) + + if prompt_embeds is None: + prompt = [prompt] if isinstance(prompt, str) else prompt + if not prompt or not all(isinstance(text, str) for text in prompt): + raise ValueError("Provide a string prompt, a list of strings, or `prompt_embeds`.") + prompt_embeds = self._get_t5_prompt_embeds(prompt, max_sequence_length, device, dtype) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + if do_classifier_free_guidance: + if negative_prompt_embeds is None: + negative_prompt = "" if negative_prompt is None else negative_prompt + negative_prompt = ( + [negative_prompt] * prompt_embeds.shape[0] if isinstance(negative_prompt, str) else negative_prompt + ) + if len(negative_prompt) != prompt_embeds.shape[0]: + raise ValueError("Negative and positive prompts must have the same batch size.") + negative_prompt_embeds = self._get_t5_prompt_embeds( + negative_prompt, max_sequence_length, device, dtype + ) + else: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + if negative_prompt_embeds.shape != prompt_embeds.shape: + raise ValueError("Positive and negative prompt embeddings must have matching shapes.") + return prompt_embeds, negative_prompt_embeds + + def check_inputs( + self, + prompt, + prompt_embeds, + negative_prompt, + negative_prompt_embeds, + image, + video, + height, + width, + output_type, + callback_inputs, + ): + """Validate conditioning inputs, output dimensions, and callback tensor names.""" + if image is None or video is None: + raise ValueError("Both a first-frame `image` and a full-rate RGB proxy `video` are required.") + if prompt is not None and prompt_embeds is not None: + raise ValueError("Pass either `prompt` or `prompt_embeds`, not both.") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError("Pass either `negative_prompt` or `negative_prompt_embeds`, not both.") + patch = self.transformer.config.patch_size + if ( + height <= 0 + or width <= 0 + or height % (self.vae_scale_factor_spatial * patch[1]) + or width % (self.vae_scale_factor_spatial * patch[2]) + ): + raise ValueError( + "Height and width must be positive and divisible by the VAE scale times the spatial patch size." + ) + if output_type not in ("np", "pt", "pil", "latent"): + raise ValueError("`output_type` must be 'np', 'pt', 'pil', or 'latent'.") + if not set(callback_inputs).issubset(self._callback_tensor_inputs): + raise ValueError(f"Callback tensor inputs must belong to {self._callback_tensor_inputs}.") + + def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): + """Prepare full-video initial noise, including the fixed first frame's noise draw.""" + shape = ( + batch_size, + self.vae.config.z_dim, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + if latents is not None: + if tuple(latents.shape) != shape: + raise ValueError(f"Expected latents with shape {shape}, got {tuple(latents.shape)}.") + return latents.to(device=device, dtype=dtype) + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + def prepare_conditions(self, image, proxy_video, height, width, num_frames, batch_size, device, dtype): + """Encode the reference image and sampled RGB proxy frames for the video transformer.""" + pixels = self.proxy_video_processor.preprocess(image).float() * 255 + scale = max(height / pixels.shape[-2], width / pixels.shape[-1]) + resized_height = max(height, math.ceil(pixels.shape[-2] * scale)) + resized_width = max(width, math.ceil(pixels.shape[-1] * scale)) + pixels = F.interpolate(pixels, (resized_height, resized_width), mode="bilinear", align_corners=False) + top, left = (resized_height - height) // 2, (resized_width - width) // 2 + pixels = pixels[:, :, top : top + height, left : left + width] / 127.5 - 1 + pixels = pixels.to(device=device) + if pixels.shape[0] == 1: + pixels = pixels.expand(batch_size, -1, -1, -1) + if pixels.shape[0] != batch_size: + raise ValueError("Image batch size must be one or match the prompt batch.") + if proxy_video.shape[0] not in (1, batch_size): + raise ValueError("Proxy-video batch size must be one or match the prompt batch.") + clip_embeds = None + if self._image_conditioned: + if self.image_encoder is None or self.image_processor is None: + raise ValueError("The image-conditioned model requires a CLIP image encoder and processor.") + rgb = ((pixels.float().clamp(-1, 1) + 1) * 127.5).round().to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy() + clip_inputs = self.image_processor(images=list(rgb), return_tensors="pt").pixel_values + clip_embeds = ( + self.image_encoder( + clip_inputs.to(device=device, dtype=self.image_encoder.dtype), output_hidden_states=True + ) + .hidden_states[-2] + .to(dtype) + ) + if self._image_conditioned: + condition_video = torch.cat( + [pixels.unsqueeze(2), pixels.new_zeros(batch_size, 3, num_frames - 1, height, width)], dim=2 + ) + else: + condition_video = pixels.unsqueeze(2) + condition = self.vae.encode(condition_video.to(self.vae.dtype)).latent_dist.mode() + mean = condition.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / condition.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + condition = ((condition - mean) * inverse_std).to(dtype) + first_frame = condition[:, :, :1] + if self._image_conditioned: + mask = condition.new_zeros(batch_size, self.vae_scale_factor_temporal, *condition.shape[2:]) + mask[:, :, 0] = 1 + condition = torch.cat([mask, condition], dim=1) + proxy_video = proxy_video.expand(batch_size, -1, -1, -1, -1) + proxy_height, proxy_width = self.config.proxy_image_height, self.config.proxy_image_width + proxy_features = [] + for frame in proxy_video.unbind(2): + frame = frame.to(device=device, dtype=torch.float32) + if not self._image_conditioned: + # Match the original FP32 [-1, 1] -> [0, 1] round trip before DINO. + frame = (frame * 2 - 1 + 1) * 0.5 + inputs = self.proxy_video_processor.preprocess_proxy_frames( + frame, proxy_height, proxy_width, "bicubic" if self._image_conditioned else "bilinear" + ) + with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + hidden = self.proxy_encoder(pixel_values=inputs).last_hidden_state + patch_size = self.proxy_encoder.config.patch_size + grid_height, grid_width = proxy_height // patch_size, proxy_width // patch_size + start = 1 + self.proxy_encoder.config.num_register_tokens + hidden = hidden[:, start : start + grid_height * grid_width].reshape( + batch_size, grid_height, grid_width, -1 + ) + proxy_features.append(hidden.permute(0, 3, 1, 2)) + proxy_features = torch.stack(proxy_features, dim=2) + return condition, first_frame, clip_embeds, proxy_features + + def decode_latents(self, latents, output_type): + """Decode normalized diffusion latents with the VAE or optional LightTAE.""" + if output_type == "latent": + return latents + latents = latents.to(self.vae.dtype) + mean = latents.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + std = latents.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + if isinstance(self.decoder, EdenLiveLightTAEDecoder): + raw_latents = latents * std + mean + else: + raw_latents = latents / (1 / std) + mean + decoder = self.decoder if self.decoder is not None else self.vae + video = decoder.decode(raw_latents.to(dtype=decoder.dtype), return_dict=False)[0] + return self.video_processor.postprocess_video(video, output_type=output_type) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] | None = None, + image=None, + video=None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 512, + output_type: str = "np", + return_dict: bool = True, + attention_kwargs: dict | None = None, + callback_on_step_end: Callable | None = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + ) -> EdenLivePipelineOutput | tuple: + """Generate a complete bidirectional video. + + Args: + prompt (`str` or `list[str]`, *optional*): + Text describing the output video. Supply this or `prompt_embeds`. + image (`PIL.Image.Image`, `np.ndarray`, or `torch.Tensor`): + Photoreal first-frame reference. Tensor and NumPy RGB values are in [0, 1]. The image is resized with + tensor bilinear interpolation and center-cropped to the output dimensions. + video (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`): + Full-rate decoded RGB proxy video, already sampled at the desired output frame rate. Tensor inputs have + shape `(T, C, H, W)` or `(B, T, C, H, W)`; NumPy inputs use channels last. Floating-point values are in + [0, 1]. Proxy frames are sampled at the VAE temporal stride exactly once. + height (`int`, *optional*): + Output height. Defaults to 704 for 5B or 480 for 14B. + width (`int`, *optional*): + Output width. Defaults to 1280 for 5B or 832 for 14B. + num_frames (`int`, *optional*): + Number of decoded output frames, of the form `4k + 1`. When omitted, use the largest valid input + prefix. The requested length cannot exceed the input video; missing frames are never repeated. + num_inference_steps (`int`, *optional*): + Denoising steps; defaults to 48 for 5B or 30 for 14B. + guidance_scale (`float`, *optional*): + Classifier-free guidance strength; defaults to 3. Values greater than one enable the negative prompt + branch. + negative_prompt (`str` or `list[str]`, *optional*): + Text describing content to discourage when classifier-free guidance is enabled. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator, or one per batch item. CPU generators are supported. + latents (`torch.Tensor`, *optional*): + Initial noise with shape `(B, C, T_latent, H_latent, W_latent)`. When omitted, noise for the whole + video is sampled before image conditioning. The fixed first frame still consumes its noise draw. + prompt_embeds (`torch.Tensor`, *optional*): + Precomputed positive text embeddings in place of `prompt`. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Precomputed negative text embeddings in place of `negative_prompt`. + max_sequence_length (`int`, defaults to 512): + Token count for each text prompt, including zero padding. + output_type (`str`, defaults to `"np"`): + Return NumPy videos (`"np"`, shape `(B, T, H, W, 3)`), PyTorch videos (`"pt"`, shape `(B, T, 3, H, + W)`), nested PIL frame lists (`"pil"`), or normalized diffusion latents (`"latent"`, shape `(B, C, + T_latent, H_latent, W_latent)`). Decoded tensor and NumPy values are in [0, 1]. + return_dict (`bool`, defaults to `True`): + Return `EdenLivePipelineOutput`; otherwise return a one-element tuple containing the frames. + attention_kwargs (`dict`, *optional*): + Transformer options. The `scale` entry sets the LoRA scale when adapters are loaded. + callback_on_step_end (`Callable`, *optional*): + Function called as `callback(self, step_index, timestep, tensors)` after each denoising step. Return a + dictionary containing any replacement tensors. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Tensors supplied to the callback, selected from `latents`, `prompt_embeds`, and + `negative_prompt_embeds`. + + Examples: + + Returns: + [`EdenLivePipelineOutput`] or `tuple`: + Generated frames, or latent videos when `output_type="latent"`. + """ + callback_on_step_end_tensor_inputs = ( + ["latents"] if callback_on_step_end_tensor_inputs is None else callback_on_step_end_tensor_inputs + ) + if guidance_scale is None: + guidance_scale = 3.0 + if num_inference_steps is None: + num_inference_steps = 30 if self._image_conditioned else 48 + height = (480 if self._image_conditioned else 704) if height is None else height + width = (832 if self._image_conditioned else 1280) if width is None else width + self.check_inputs( + prompt, + prompt_embeds, + negative_prompt, + negative_prompt_embeds, + image, + video, + height, + width, + output_type, + callback_on_step_end_tensor_inputs, + ) + self._guidance_scale = guidance_scale + device, dtype = self._execution_device, self.transformer.dtype + proxy_video, num_frames = self.proxy_video_processor.preprocess_proxy_video( + video, num_frames, self.vae_scale_factor_temporal + ) + batch_size = ( + prompt_embeds.shape[0] + if prompt_embeds is not None + else (1 if isinstance(prompt, str) else len(prompt or [])) + ) + if batch_size < 1: + raise ValueError("Provide a positive prompt or prompt embeddings.") + latents = self.prepare_latents(batch_size, height, width, num_frames, dtype, device, generator, latents) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + negative_prompt, + self.do_classifier_free_guidance, + prompt_embeds, + negative_prompt_embeds, + max_sequence_length, + ) + condition, first_frame, image_embeds, proxy_features = self.prepare_conditions( + image, proxy_video, height, width, num_frames, batch_size, device, dtype + ) + if not self._image_conditioned: + latents[:, :, :1] = first_frame + self.scheduler.set_timesteps(num_inference_steps, device=device) + if not self._image_conditioned and isinstance(self.scheduler, UniPCMultistepScheduler): + # CPU and CUDA scalar sigmas round BF16 model outputs differently. + self.scheduler.sigmas = self.scheduler.sigmas.to(device=device) + self._num_timesteps = len(self.scheduler.timesteps) + with self.progress_bar(total=self._num_timesteps) as progress_bar: + for step_index, timestep in enumerate(self.scheduler.timesteps): + model_input, model_timestep = prepare_model_input(self, latents, condition, first_frame, timestep) + prediction = self.transformer( + model_input, + model_timestep, + prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=proxy_features, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + if self.do_classifier_free_guidance: + negative_prediction = self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=proxy_features, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + prediction = negative_prediction + self.guidance_scale * (prediction - negative_prediction) + latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0] + if not self._image_conditioned: + latents[:, :, :1] = first_frame + if callback_on_step_end is not None: + callback_kwargs = {} + for key in callback_on_step_end_tensor_inputs: + callback_kwargs[key] = locals()[key] + callback_outputs = callback_on_step_end(self, step_index, timestep, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + progress_bar.update() + frames = self.decode_latents(latents, output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (frames,) + return EdenLivePipelineOutput(frames=frames) diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py new file mode 100644 index 000000000000..45fb9b86b76a --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py @@ -0,0 +1,621 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import html +import math +import re +from typing import Callable + +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel + +from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveTransformer3DModel +from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler +from ...utils import is_ftfy_available, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import EdenLivePipelineOutput +from .pipeline_utils import prepare_model_input +from .video_processor import EdenLiveVideoProcessor + + +if is_ftfy_available(): + import ftfy + + +# docstyle-ignore +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import EdenLiveCausalPipeline + >>> from diffusers.utils import export_to_video, load_image, load_video + >>> pipe = EdenLiveCausalPipeline.from_pretrained( + ... "VAST-AI/EdenLive-5B-Causal-Diffusers", + ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, + ... ).to("cuda") + >>> frames = pipe( + ... image=load_image("reference.png"), video=load_video("proxy.mp4"), + ... prompt="A character walks through snow.", num_frames=49, + ... ).frames[0] + >>> export_to_video(frames, "edenlive.mp4", fps=24) + ``` +""" + + +class EdenLiveCausalPipeline(DiffusionPipeline): + """Causal video generation conditioned on a reference image, text, and an RGB proxy video. + + Args: + vae ([`AutoencoderKLWan`]): + Wan VAE for reference encoding and full-quality video decoding. + text_encoder (`UMT5EncoderModel`): + Text encoder. + tokenizer (`AutoTokenizer`): + Tokenizer corresponding to the text encoder. + transformer ([`EdenLiveTransformer3DModel`]): + EdenLive causal video denoiser. + scheduler ([`EdenLiveFlowMatchEulerDiscreteScheduler`]): + Scheduler corresponding to the transformer checkpoint. + proxy_encoder (`DINOv3ViTModel`): + Trained RGB proxy encoder. + image_encoder (`CLIPVisionModel`, *optional*): + Reference image encoder required by the 14B model. + image_processor (`CLIPImageProcessor`, *optional*): + Processor corresponding to the image encoder. + decoder ([`EdenLiveLightTAEDecoder`], *optional*): + Optional faster decoder. When omitted, decode with the full Wan VAE. + proxy_image_height (`int`, defaults to 480): + Height of the proxy images passed to DINO. + proxy_image_width (`int`, defaults to 832): + Width of the proxy images passed to DINO. + chunk_size (`int`, defaults to 3): + Latent frames generated in each chunk. + max_context_frames (`int`, defaults to 18): + Maximum number of latent frames in the attention window, including the current chunk. + sink_frames (`int`, defaults to 6): + Initial latent frames retained when the attention window fills. + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->proxy_encoder->transformer->vae->decoder" + _optional_components = ["image_encoder", "image_processor", "decoder"] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + vae: AutoencoderKLWan, + text_encoder: UMT5EncoderModel, + tokenizer: AutoTokenizer, + transformer: EdenLiveTransformer3DModel, + scheduler: EdenLiveFlowMatchEulerDiscreteScheduler, + proxy_encoder: DINOv3ViTModel, + image_encoder: CLIPVisionModel | None = None, + image_processor: CLIPImageProcessor | None = None, + decoder: EdenLiveLightTAEDecoder | None = None, + proxy_image_height: int = 480, + proxy_image_width: int = 832, + chunk_size: int = 3, + max_context_frames: int = 18, + sink_frames: int = 6, + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + proxy_encoder=proxy_encoder, + image_encoder=image_encoder, + image_processor=image_processor, + decoder=decoder, + ) + self.register_to_config(proxy_image_height=proxy_image_height, proxy_image_width=proxy_image_width) + if chunk_size < 1 or sink_frames < 0 or max_context_frames < sink_frames + chunk_size: + raise ValueError("The context must fit the sink and a positive-size chunk.") + self.register_to_config(chunk_size=chunk_size, max_context_frames=max_context_frames, sink_frames=sink_frames) + self._image_conditioned = transformer is not None and transformer.config.image_dim is not None + self.vae_scale_factor_temporal = vae.config.scale_factor_temporal if vae is not None else 4 + self.vae_scale_factor_spatial = vae.config.scale_factor_spatial if vae is not None else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + self.proxy_video_processor = EdenLiveVideoProcessor( + vae_scale_factor=1, do_normalize=False, do_convert_rgb=True + ) + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline._get_t5_prompt_embeds + def _get_t5_prompt_embeds(self, texts, max_sequence_length, device, dtype): + if is_ftfy_available(): + texts = [ftfy.fix_text(text) for text in texts] + texts = [re.sub(r"\s+", " ", html.unescape(html.unescape(text))).strip() for text in texts] + inputs = self.tokenizer( + texts, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + mask = inputs.attention_mask.to(device) + embeddings = self.text_encoder(inputs.input_ids.to(device), attention_mask=mask).last_hidden_state + embeddings = embeddings.to(device=device, dtype=dtype) + return embeddings.masked_fill(~mask.bool().unsqueeze(-1), 0) + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline.encode_prompt + def encode_prompt( + self, + prompt=None, + negative_prompt=None, + do_classifier_free_guidance=True, + prompt_embeds=None, + negative_prompt_embeds=None, + max_sequence_length=512, + device=None, + dtype=None, + ): + """Encode positive and optional negative text with UMT5, zeroing padded token embeddings.""" + device = device or self._execution_device + dtype = dtype or (self.transformer.dtype if self.transformer is not None else self.text_encoder.dtype) + + if prompt_embeds is None: + prompt = [prompt] if isinstance(prompt, str) else prompt + if not prompt or not all(isinstance(text, str) for text in prompt): + raise ValueError("Provide a string prompt, a list of strings, or `prompt_embeds`.") + prompt_embeds = self._get_t5_prompt_embeds(prompt, max_sequence_length, device, dtype) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + if do_classifier_free_guidance: + if negative_prompt_embeds is None: + negative_prompt = "" if negative_prompt is None else negative_prompt + negative_prompt = ( + [negative_prompt] * prompt_embeds.shape[0] if isinstance(negative_prompt, str) else negative_prompt + ) + if len(negative_prompt) != prompt_embeds.shape[0]: + raise ValueError("Negative and positive prompts must have the same batch size.") + negative_prompt_embeds = self._get_t5_prompt_embeds( + negative_prompt, max_sequence_length, device, dtype + ) + else: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + if negative_prompt_embeds.shape != prompt_embeds.shape: + raise ValueError("Positive and negative prompt embeddings must have matching shapes.") + return prompt_embeds, negative_prompt_embeds + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline.check_inputs + def check_inputs( + self, + prompt, + prompt_embeds, + negative_prompt, + negative_prompt_embeds, + image, + video, + height, + width, + output_type, + callback_inputs, + ): + """Validate conditioning inputs, output dimensions, and callback tensor names.""" + if image is None or video is None: + raise ValueError("Both a first-frame `image` and a full-rate RGB proxy `video` are required.") + if prompt is not None and prompt_embeds is not None: + raise ValueError("Pass either `prompt` or `prompt_embeds`, not both.") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError("Pass either `negative_prompt` or `negative_prompt_embeds`, not both.") + patch = self.transformer.config.patch_size + if ( + height <= 0 + or width <= 0 + or height % (self.vae_scale_factor_spatial * patch[1]) + or width % (self.vae_scale_factor_spatial * patch[2]) + ): + raise ValueError( + "Height and width must be positive and divisible by the VAE scale times the spatial patch size." + ) + if output_type not in ("np", "pt", "pil", "latent"): + raise ValueError("`output_type` must be 'np', 'pt', 'pil', or 'latent'.") + if not set(callback_inputs).issubset(self._callback_tensor_inputs): + raise ValueError(f"Callback tensor inputs must belong to {self._callback_tensor_inputs}.") + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline.prepare_latents + def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): + """Prepare full-video initial noise, including the fixed first frame's noise draw.""" + shape = ( + batch_size, + self.vae.config.z_dim, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + if latents is not None: + if tuple(latents.shape) != shape: + raise ValueError(f"Expected latents with shape {shape}, got {tuple(latents.shape)}.") + return latents.to(device=device, dtype=dtype) + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline.prepare_conditions + def prepare_conditions(self, image, proxy_video, height, width, num_frames, batch_size, device, dtype): + """Encode the reference image and sampled RGB proxy frames for the video transformer.""" + pixels = self.proxy_video_processor.preprocess(image).float() * 255 + scale = max(height / pixels.shape[-2], width / pixels.shape[-1]) + resized_height = max(height, math.ceil(pixels.shape[-2] * scale)) + resized_width = max(width, math.ceil(pixels.shape[-1] * scale)) + pixels = F.interpolate(pixels, (resized_height, resized_width), mode="bilinear", align_corners=False) + top, left = (resized_height - height) // 2, (resized_width - width) // 2 + pixels = pixels[:, :, top : top + height, left : left + width] / 127.5 - 1 + pixels = pixels.to(device=device) + if pixels.shape[0] == 1: + pixels = pixels.expand(batch_size, -1, -1, -1) + if pixels.shape[0] != batch_size: + raise ValueError("Image batch size must be one or match the prompt batch.") + if proxy_video.shape[0] not in (1, batch_size): + raise ValueError("Proxy-video batch size must be one or match the prompt batch.") + clip_embeds = None + if self._image_conditioned: + if self.image_encoder is None or self.image_processor is None: + raise ValueError("The image-conditioned model requires a CLIP image encoder and processor.") + rgb = ((pixels.float().clamp(-1, 1) + 1) * 127.5).round().to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy() + clip_inputs = self.image_processor(images=list(rgb), return_tensors="pt").pixel_values + clip_embeds = ( + self.image_encoder( + clip_inputs.to(device=device, dtype=self.image_encoder.dtype), output_hidden_states=True + ) + .hidden_states[-2] + .to(dtype) + ) + if self._image_conditioned: + condition_video = torch.cat( + [pixels.unsqueeze(2), pixels.new_zeros(batch_size, 3, num_frames - 1, height, width)], dim=2 + ) + else: + condition_video = pixels.unsqueeze(2) + condition = self.vae.encode(condition_video.to(self.vae.dtype)).latent_dist.mode() + mean = condition.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / condition.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + condition = ((condition - mean) * inverse_std).to(dtype) + first_frame = condition[:, :, :1] + if self._image_conditioned: + mask = condition.new_zeros(batch_size, self.vae_scale_factor_temporal, *condition.shape[2:]) + mask[:, :, 0] = 1 + condition = torch.cat([mask, condition], dim=1) + proxy_video = proxy_video.expand(batch_size, -1, -1, -1, -1) + proxy_height, proxy_width = self.config.proxy_image_height, self.config.proxy_image_width + proxy_features = [] + for frame in proxy_video.unbind(2): + frame = frame.to(device=device, dtype=torch.float32) + if not self._image_conditioned: + # Match the original FP32 [-1, 1] -> [0, 1] round trip before DINO. + frame = (frame * 2 - 1 + 1) * 0.5 + inputs = self.proxy_video_processor.preprocess_proxy_frames( + frame, proxy_height, proxy_width, "bicubic" if self._image_conditioned else "bilinear" + ) + with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + hidden = self.proxy_encoder(pixel_values=inputs).last_hidden_state + patch_size = self.proxy_encoder.config.patch_size + grid_height, grid_width = proxy_height // patch_size, proxy_width // patch_size + start = 1 + self.proxy_encoder.config.num_register_tokens + hidden = hidden[:, start : start + grid_height * grid_width].reshape( + batch_size, grid_height, grid_width, -1 + ) + proxy_features.append(hidden.permute(0, 3, 1, 2)) + proxy_features = torch.stack(proxy_features, dim=2) + return condition, first_frame, clip_embeds, proxy_features + + # Copied from diffusers.pipelines.edenlive.pipeline_edenlive.EdenLivePipeline.decode_latents + def decode_latents(self, latents, output_type): + """Decode normalized diffusion latents with the VAE or optional LightTAE.""" + if output_type == "latent": + return latents + latents = latents.to(self.vae.dtype) + mean = latents.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + std = latents.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + if isinstance(self.decoder, EdenLiveLightTAEDecoder): + raw_latents = latents * std + mean + else: + raw_latents = latents / (1 / std) + mean + decoder = self.decoder if self.decoder is not None else self.vae + video = decoder.decode(raw_latents.to(dtype=decoder.dtype), return_dict=False)[0] + return self.video_processor.postprocess_video(video, output_type=output_type) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] | None = None, + image=None, + video=None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + chunk_prompts: list[str] | None = None, + max_sequence_length: int = 512, + output_type: str = "np", + return_dict: bool = True, + attention_kwargs: dict | None = None, + callback_on_step_end: Callable | None = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + ) -> EdenLivePipelineOutput | tuple: + """Generate a complete causal video. + + Args: + prompt (`str` or `list[str]`, *optional*): + Text describing the output video. Supply this or `prompt_embeds` unless using `chunk_prompts`. + image (`PIL.Image.Image`, `np.ndarray`, or `torch.Tensor`): + Photoreal first-frame reference. Tensor and NumPy RGB values are in [0, 1]. The image is resized with + tensor bilinear interpolation and center-cropped to the output dimensions. + video (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`): + Full-rate decoded RGB proxy video, already sampled at the desired output frame rate. Tensor inputs have + shape `(T, C, H, W)` or `(B, T, C, H, W)`; NumPy inputs use channels last. Floating-point values are in + [0, 1]. Proxy frames are sampled at the VAE temporal stride exactly once. + height (`int`, *optional*): + Output height. Defaults to 704 for 5B or 480 for 14B. + width (`int`, *optional*): + Output width. Defaults to 1280 for 5B or 832 for 14B. + num_frames (`int`, *optional*): + Number of decoded output frames, of the form `4k + 1`. When omitted, use the largest valid input + prefix. The requested length cannot exceed the input video; missing frames are never repeated. + num_inference_steps (`int`, *optional*): + Denoising steps per chunk; released causal schedules use four fixed evaluation points. + guidance_scale (`float`, *optional*): + Classifier-free guidance strength; defaults to 1. Values greater than one enable the negative prompt + branch. + negative_prompt (`str` or `list[str]`, *optional*): + Text describing content to discourage when classifier-free guidance is enabled. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator, or one per batch item. CPU generators are supported. + latents (`torch.Tensor`, *optional*): + Initial noise with shape `(B, C, T_latent, H_latent, W_latent)`. When omitted, noise for the whole + video is sampled before image conditioning. The fixed first frame still consumes its noise draw. + prompt_embeds (`torch.Tensor`, *optional*): + Precomputed positive text embeddings in place of `prompt`. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Precomputed negative text embeddings in place of `negative_prompt`. + chunk_prompts (`list[str]`, *optional*): + One positive prompt per latent chunk, including a short final chunk. Overrides a single `prompt`. + Requires one video and cannot be combined with `prompt_embeds`. The negative prompt stays fixed. Text + changes take effect at chunk boundaries while retaining previously generated video history. + max_sequence_length (`int`, defaults to 512): + Token count for each text prompt, including zero padding. + output_type (`str`, defaults to `"np"`): + Return NumPy videos (`"np"`, shape `(B, T, H, W, 3)`), PyTorch videos (`"pt"`, shape `(B, T, 3, H, + W)`), nested PIL frame lists (`"pil"`), or normalized diffusion latents (`"latent"`, shape `(B, C, + T_latent, H_latent, W_latent)`). Decoded tensor and NumPy values are in [0, 1]. + return_dict (`bool`, defaults to `True`): + Return `EdenLivePipelineOutput`; otherwise return a one-element tuple containing the frames. + attention_kwargs (`dict`, *optional*): + Transformer options. The `scale` entry sets the LoRA scale when adapters are loaded. + callback_on_step_end (`Callable`, *optional*): + Function called as `callback(self, step_index, timestep, tensors)` after each denoising step. Return a + dictionary containing any replacement tensors. In causal mode the tensors contain the current chunk and + step indices count across all chunks. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Tensors supplied to the callback, selected from `latents`, `prompt_embeds`, and + `negative_prompt_embeds`. + + Examples: + + Returns: + [`EdenLivePipelineOutput`] or `tuple`: + Generated frames, or latent videos when `output_type="latent"`. + """ + callback_on_step_end_tensor_inputs = ( + ["latents"] if callback_on_step_end_tensor_inputs is None else callback_on_step_end_tensor_inputs + ) + if guidance_scale is None: + guidance_scale = 1.0 + if num_inference_steps is None: + num_inference_steps = 4 + height = (480 if self._image_conditioned else 704) if height is None else height + width = (832 if self._image_conditioned else 1280) if width is None else width + self.check_inputs( + prompt, + prompt_embeds, + negative_prompt, + negative_prompt_embeds, + image, + video, + height, + width, + output_type, + callback_on_step_end_tensor_inputs, + ) + self._guidance_scale = guidance_scale + device, dtype = self._execution_device, self.transformer.dtype + proxy_video, num_frames = self.proxy_video_processor.preprocess_proxy_video( + video, num_frames, self.vae_scale_factor_temporal + ) + batch_size = ( + prompt_embeds.shape[0] + if prompt_embeds is not None + else (1 if isinstance(prompt, str) else len(prompt or [])) + ) + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + num_chunks = (num_latent_frames + self.config.chunk_size - 1) // self.config.chunk_size + if chunk_prompts is not None: + if prompt is None and prompt_embeds is None: + batch_size = 1 + if ( + batch_size != 1 + or prompt_embeds is not None + or not isinstance(chunk_prompts, (list, tuple)) + or len(chunk_prompts) != num_chunks + or not all(isinstance(text, str) for text in chunk_prompts) + ): + raise ValueError( + "`chunk_prompts` requires batch size one, no prompt embeddings, and one string per chunk." + ) + if batch_size < 1: + raise ValueError("Provide a positive prompt or prompt embeddings or chunk_prompts.") + latents = self.prepare_latents(batch_size, height, width, num_frames, dtype, device, generator, latents) + chunk_embeds = None + if chunk_prompts is not None: + unique_embeds = {} + for text in chunk_prompts: + if text not in unique_embeds: + unique_embeds[text], _ = self.encode_prompt( + text, do_classifier_free_guidance=False, max_sequence_length=max_sequence_length + ) + chunk_embeds = torch.cat([unique_embeds[text] for text in chunk_prompts]) + prompt, prompt_embeds = chunk_prompts[0], chunk_embeds[:1] + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + negative_prompt, + self.do_classifier_free_guidance, + prompt_embeds, + negative_prompt_embeds, + max_sequence_length, + ) + condition, first_frame, image_embeds, proxy_features = self.prepare_conditions( + image, proxy_video, height, width, num_frames, batch_size, device, dtype + ) + states = [ + self.transformer.create_causal_state( + batch_size=batch_size, + height=latents.shape[3], + width=latents.shape[4], + chunk_size=self.config.chunk_size, + max_context_frames=self.config.max_context_frames, + sink_frames=self.config.sink_frames, + device=device, + dtype=dtype, + ) + for _ in range(2 if self.do_classifier_free_guidance else 1) + ] + self._num_timesteps = num_inference_steps * num_chunks + generated = [] + with self.progress_bar(total=num_chunks) as progress_bar: + for chunk_index, start in enumerate(range(0, latents.shape[2], self.config.chunk_size)): + if chunk_embeds is not None: + prompt_embeds = chunk_embeds[chunk_index : chunk_index + 1] + current = latents[:, :, start : start + self.config.chunk_size].clone() + current_proxy = proxy_features[:, :, start : start + self.config.chunk_size] + if start == 0: + current[:, :, :1] = first_frame + self.scheduler.set_timesteps(num_inference_steps, device=current.device) + for step_index, timestep in enumerate(self.scheduler.timesteps): + model_input, model_timestep = prepare_model_input( + self, current, condition, first_frame, timestep, start, causal=True + ) + prediction = self.transformer( + model_input, + model_timestep, + prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=current_proxy, + causal_state=states[0], + cache_update=False, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + if self.do_classifier_free_guidance: + negative_prediction = self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=current_proxy, + causal_state=states[1], + cache_update=False, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + prediction = negative_prediction + self.guidance_scale * (prediction - negative_prediction) + variance_noise = None + if self.scheduler.config.stochastic_sampling and step_index < len(self.scheduler.timesteps) - 1: + variance_noise = randn_tensor( + current.shape, generator=generator, device=current.device, dtype=current.dtype + ) + current = self.scheduler.step( + prediction, + timestep, + current, + generator=None if variance_noise is not None else generator, + variance_noise=variance_noise, + return_dict=False, + )[0] + if start == 0: + current[:, :, :1] = first_frame + if callback_on_step_end is not None: + callback_kwargs = {} + for key in callback_on_step_end_tensor_inputs: + callback_kwargs[key] = current if key == "latents" else locals()[key] + callback_outputs = callback_on_step_end( + self, chunk_index * num_inference_steps + step_index, timestep, callback_kwargs + ) + current = callback_outputs.pop("latents", current) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + model_input, model_timestep = prepare_model_input( + self, + current, + condition, + first_frame, + current.new_zeros(()), + start, + causal=True, + ) + # A clean cache commit uses one zero timestep per sample, not per token. + model_timestep = torch.zeros(current.shape[0], device=current.device, dtype=torch.float32) + self.transformer( + model_input, + model_timestep, + prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=current_proxy, + causal_state=states[0], + cache_update=True, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + if self.do_classifier_free_guidance: + self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=current_proxy, + causal_state=states[1], + cache_update=True, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + generated.append(current) + progress_bar.update() + latents = torch.cat(generated, dim=2) + frames = self.decode_latents(latents, output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (frames,) + return EdenLivePipelineOutput(frames=frames) diff --git a/src/diffusers/pipelines/edenlive/pipeline_output.py b/src/diffusers/pipelines/edenlive/pipeline_output.py new file mode 100644 index 000000000000..fe61e1d3184b --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_output.py @@ -0,0 +1,37 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +from dataclasses import dataclass + +import numpy as np +import PIL.Image +import torch + +from ...utils import BaseOutput + + +@dataclass +class EdenLivePipelineOutput(BaseOutput): + """Output of [`EdenLivePipeline`]. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or `list[list[PIL.Image.Image]]`): + Generated videos selected by `output_type`. `"np"` returns shape `(B, T, H, W, 3)` and `"pt"` returns `(B, + T, 3, H, W)`, both with floating-point RGB values in [0, 1]. `"pil"` returns one list of RGB images per + video. `"latent"` bypasses decoding and returns normalized diffusion latents with shape `(B, C, T_latent, + H_latent, W_latent)`; these values are not bounded RGB pixels. `B` is the video batch size and `T` is the + decoded frame count. The latent channel count `C` is 48 for 5B and 16 for 14B. + """ + + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] diff --git a/src/diffusers/pipelines/edenlive/pipeline_utils.py b/src/diffusers/pipelines/edenlive/pipeline_utils.py new file mode 100644 index 000000000000..8141f2856087 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_utils.py @@ -0,0 +1,32 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import torch + + +def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0, causal=False): + model_input = latents + pin_first = start == 0 and (causal or not pipe._image_conditioned) + if pin_first: + model_input = model_input.clone() + model_input[:, :, :1] = first_frame + if pipe._image_conditioned: + model_input = torch.cat([model_input, condition[:, :, start : start + latents.shape[2]]], dim=1) + timestep = timestep.expand(latents.shape[0]) + if pin_first: + patch = pipe.transformer.config.patch_size + tokens_per_frame = (latents.shape[3] // patch[1]) * (latents.shape[4] // patch[2]) + timestep = timestep[:, None].expand(-1, latents.shape[2] * tokens_per_frame).clone() + timestep[:, :tokens_per_frame] = 0 + return model_input, timestep diff --git a/src/diffusers/pipelines/edenlive/video_processor.py b/src/diffusers/pipelines/edenlive/video_processor.py new file mode 100644 index 000000000000..8822346e46a0 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/video_processor.py @@ -0,0 +1,70 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F + +from ...image_processor import is_valid_image +from ...video_processor import VideoProcessor + + +class EdenLiveVideoProcessor(VideoProcessor): + """Preprocess full-rate RGB proxy videos for EdenLive's latent-time DINO conditioning.""" + + def __init__(self, vae_scale_factor=1, do_normalize=False, do_convert_rgb=True): + super().__init__(vae_scale_factor=vae_scale_factor, do_normalize=do_normalize, do_convert_rgb=do_convert_rgb) + + def preprocess_proxy_video(self, video, num_frames=None, temporal_scale_factor=4): + if isinstance(video, (np.ndarray, torch.Tensor)) and video.ndim in (4, 5): + videos = [video] if video.ndim == 4 else list(video) + elif isinstance(video, list) and video: + videos = [video] if is_valid_image(video[0]) and getattr(video[0], "ndim", 3) == 3 else video + else: + raise ValueError( + "Pass a nonempty RGB video or a batch of videos following VideoProcessor input conventions." + ) + available = min(len(sample) for sample in videos) + if num_frames is None: + num_frames = (available - 1) // temporal_scale_factor * temporal_scale_factor + 1 + if not isinstance(num_frames, int) or num_frames < 1 or (num_frames - 1) % temporal_scale_factor: + raise ValueError(f"`num_frames` must be a positive integer of the form {temporal_scale_factor}k + 1.") + if num_frames > available: + raise ValueError(f"Requested {num_frames} frames, but the proxy video contains only {available}.") + video = self.preprocess_video([sample[:num_frames:temporal_scale_factor] for sample in videos]) + if video.shape[1] != 3: + raise ValueError("Proxy video must contain RGB frames.") + if not torch.isfinite(video).all() or video.min() < 0 or video.max() > 1: + raise ValueError("Proxy video must contain finite RGB values in [0, 1].") + return video, num_frames + + def preprocess_proxy_frames(self, images, height, width, resize_mode="bilinear"): + if resize_mode == "bilinear": + pixels = F.interpolate( + images.float(), (height, width), mode="bilinear", align_corners=False, antialias=True + ) + elif resize_mode == "bicubic": + arrays = images.float().clamp(0, 1).mul(255).round().to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy() + processed = [] + for array in arrays: + image = PIL.Image.fromarray(array) + image = image.resize((width, height), PIL.Image.Resampling.BICUBIC) + processed.append(torch.from_numpy(np.array(image)).permute(2, 0, 1)) + pixels = torch.stack(processed).to(device=images.device, dtype=torch.float32) / 255 + else: + raise ValueError(f"Unsupported proxy resize mode: {resize_mode}.") + mean = pixels.new_tensor([0.485, 0.456, 0.406])[None, :, None, None] + std = pixels.new_tensor([0.229, 0.224, 0.225])[None, :, None, None] + return (pixels - mean) / std diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 47f4a0857090..ff99bb95d8ff 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -55,6 +55,7 @@ _import_structure["scheduling_dpmsolver_multistep"] = ["DPMSolverMultistepScheduler"] _import_structure["scheduling_dpmsolver_multistep_inverse"] = ["DPMSolverMultistepInverseScheduler"] _import_structure["scheduling_dpmsolver_singlestep"] = ["DPMSolverSinglestepScheduler"] + _import_structure["scheduling_edenlive"] = ["EdenLiveFlowMatchEulerDiscreteScheduler"] _import_structure["scheduling_edm_dpmsolver_multistep"] = ["EDMDPMSolverMultistepScheduler"] _import_structure["scheduling_edm_euler"] = ["EDMEulerScheduler"] _import_structure["scheduling_entropy_bound"] = ["EntropyBoundScheduler", "EntropyBoundSchedulerOutput"] @@ -140,6 +141,7 @@ from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler + from .scheduling_edenlive import EdenLiveFlowMatchEulerDiscreteScheduler from .scheduling_edm_dpmsolver_multistep import EDMDPMSolverMultistepScheduler from .scheduling_edm_euler import EDMEulerScheduler from .scheduling_entropy_bound import EntropyBoundScheduler, EntropyBoundSchedulerOutput diff --git a/src/diffusers/schedulers/scheduling_edenlive.py b/src/diffusers/schedulers/scheduling_edenlive.py new file mode 100644 index 000000000000..1a1a8e6f1188 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_edenlive.py @@ -0,0 +1,211 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils.torch_utils import randn_tensor +from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteSchedulerOutput +from .scheduling_utils import SchedulerMixin + + +class EdenLiveFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Flow Euler scheduler for EdenLive, including its distilled stochastic trajectory. + + Args: + num_train_timesteps (`int`, defaults to 1000): + Scale used to convert model timesteps to noise levels. + shift (`float`, defaults to 3.0): + Flow shift applied to generated schedules. Explicit timesteps and sigmas are already in model space. + stochastic_sampling (`bool`, defaults to `False`): + Predict a clean endpoint and re-noise it between model evaluations. + default_timesteps (`list[float]`, *optional*): + Model timesteps used by a distilled checkpoint. Their length fixes the default number of inference steps. + """ + + order = 1 + _compatibles = [] + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 3.0, + stochastic_sampling: bool = False, + default_timesteps: list[float] | None = None, + ): + if num_train_timesteps < 1 or shift <= 0: + raise ValueError("`num_train_timesteps` and `shift` must be positive.") + self.num_inference_steps = None + self.init_noise_sigma = 1.0 + self._step_index = None + self._begin_index = None + self.timesteps = torch.empty(0, dtype=torch.float32) + self.sigmas = torch.empty(0, dtype=torch.float32) + + @property + def step_index(self): + return self._step_index + + @property + def begin_index(self): + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + if not 0 <= begin_index < len(self.timesteps): + raise ValueError("`begin_index` must select an inference timestep.") + self._begin_index = begin_index + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[float] | None = None, + sigmas: list[float] | None = None, + ): + """Set model evaluation points and reset the step and begin indices for a new video or chunk. + + Args: + num_inference_steps (`int`, *optional*): + Number of model evaluations. Required when neither an explicit nor a configured default schedule is + available. When a schedule is provided, this value must match its number of positive evaluation points. + device (`str` or `torch.device`, *optional*): + Device for the FP32 `timesteps` and `sigmas` tensors. + timesteps (`list[float]`, *optional*): + Strictly decreasing model timesteps in [0, `num_train_timesteps`]. They are divided by + `num_train_timesteps` to obtain sigmas and are not shifted again. An optional final zero is removed + from the model evaluation points. + sigmas (`list[float]`, *optional*): + Strictly decreasing noise levels in [0, 1], optionally ending in zero. Mutually exclusive with + `timesteps`. All values must be finite. When both arguments are omitted, `default_timesteps` takes + precedence; otherwise a schedule is generated using `shift` and `num_inference_steps`. + + The internal `sigmas` tensor includes a terminal zero after the positive model evaluation points. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Pass only one of `timesteps` or `sigmas`.") + if timesteps is None and sigmas is None and self.config.default_timesteps is not None: + timesteps = self.config.default_timesteps + if timesteps is not None or sigmas is not None: + values = torch.as_tensor( + timesteps if timesteps is not None else sigmas, dtype=torch.float32, device=device + ) + if values.ndim != 1 or values.numel() == 0: + raise ValueError("The schedule must be a non-empty one-dimensional sequence.") + if timesteps is not None: + values = values / self.config.num_train_timesteps + if not torch.isfinite(values).all() or torch.any(values < 0) or torch.any(values > 1): + raise ValueError("Schedule sigmas must be finite and lie in [0, 1].") + if torch.any(values[1:] >= values[:-1]): + raise ValueError("Schedule sigmas must be strictly decreasing.") + if values[-1] == 0: + values = values[:-1] + if values.numel() == 0 or torch.any(values == 0): + raise ValueError("The schedule must contain positive model evaluation sigmas.") + if num_inference_steps is not None and num_inference_steps != len(values): + raise ValueError("`num_inference_steps` must match the number of model evaluation points.") + else: + if num_inference_steps is None or num_inference_steps < 1: + raise ValueError("Pass a positive `num_inference_steps` or an explicit schedule.") + values = torch.linspace( + 1.0, 1.0 / self.config.num_train_timesteps, num_inference_steps + 1, dtype=torch.float32, device=device + )[:-1] + values = self.config.shift * values / (1.0 + (self.config.shift - 1.0) * values) + values[0] -= 1.0e-6 + self.num_inference_steps = len(values) + self.timesteps = values * self.config.num_train_timesteps + self.sigmas = torch.cat([values, values.new_zeros(1)]) + self._step_index = None + self._begin_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + schedule_timesteps = self.timesteps if schedule_timesteps is None else schedule_timesteps + timestep = torch.as_tensor(timestep, device=schedule_timesteps.device, dtype=schedule_timesteps.dtype) + indices = (schedule_timesteps == timestep).nonzero().flatten() + if len(indices) == 0: + raise ValueError("The timestep is not in this scheduler's current schedule.") + return indices[1 if len(indices) > 1 else 0].item() + + def scale_model_input(self, sample: torch.Tensor, timestep=None): + return sample + + def step( + self, + model_output: torch.Tensor, + timestep: float | torch.Tensor, + sample: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, + variance_noise: torch.Tensor | None = None, + return_dict: bool = True, + ): + """Advance one velocity prediction, preserving the sample dtype. + + Call `set_timesteps` before stepping through a video or causal chunk. Stochastic sampling predicts a clean + endpoint and adds fresh noise between model evaluations; the terminal step does not draw noise. + + Args: + model_output (`torch.Tensor`): + Predicted flow velocity, with the same shape as `sample`. + timestep (`float` or `torch.Tensor`): + Current model evaluation point from `scheduler.timesteps`. + sample (`torch.Tensor`): + Current noisy latent sample. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator, or one per batch item, for stochastic noise draws. Mutually exclusive with + `variance_noise`. Deterministic sampling does not draw noise. + variance_noise (`torch.Tensor`, *optional*): + Precomputed noise for a nonterminal stochastic step, for example noise shared across distributed ranks. + Must have the same shape and dtype as `sample` and be on its device. Ignored when no noise is needed. + return_dict (`bool`, defaults to `True`): + Return the scheduler output when true; otherwise return a one-element tuple. + + Returns: + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: + The next sample in `prev_sample`, or the first tuple entry, with the same dtype as `sample`. + """ + if self.num_inference_steps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if variance_noise is not None and generator is not None: + raise ValueError("Pass only one of `generator` or `variance_noise`.") + if self._step_index is None: + self._step_index = self.index_for_timestep(timestep) if self._begin_index is None else self._begin_index + if self._step_index >= self.num_inference_steps: + raise ValueError("The schedule is exhausted. Call `set_timesteps` for the next chunk.") + sigma = self.sigmas[self._step_index].to(sample.device) + sigma_next = self.sigmas[self._step_index + 1].to(sample.device) + if self.config.stochastic_sampling: + endpoint = sample - sigma.to(sample.dtype) * model_output + if self._step_index + 1 == self.num_inference_steps: + prev_sample = endpoint + else: + if variance_noise is None: + variance_noise = randn_tensor( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype + ) + elif variance_noise.shape != sample.shape or variance_noise.dtype != sample.dtype: + raise ValueError("`variance_noise` must have the same shape and dtype as `sample`.") + sigma_next = sigma_next.to(sample.dtype) + prev_sample = (1.0 - sigma_next) * endpoint + sigma_next * variance_noise + else: + delta = (sigma_next - sigma).view(*([1] * sample.ndim)) + prev_sample = sample.to(model_output.dtype) + model_output * delta + prev_sample = prev_sample.to(sample.dtype) + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 84fd853d17ab..ad94d5ea2d1f 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1324,6 +1324,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class EdenLiveCausalState(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class EdenLiveLightTAEDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class EdenLiveTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ErnieImageTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] @@ -3516,6 +3561,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class EdenLiveFlowMatchEulerDiscreteScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class EDMDPMSolverMultistepScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index c066286de27b..54d6ac8b76ba 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2057,6 +2057,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class EdenLiveCausalPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class EdenLivePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class EdenLivePipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class ErnieImagePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_edenlive.py b/tests/models/autoencoders/test_models_autoencoder_edenlive.py new file mode 100644 index 000000000000..6e752da1ea1c --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_edenlive.py @@ -0,0 +1,94 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import pytest +import torch + +from diffusers import EdenLiveLightTAEDecoder +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin + + +enable_full_determinism() + + +class EdenLiveLightTAEDecoderTesterConfig(BaseModelTesterConfig): + model_class = EdenLiveLightTAEDecoder + main_input_name = "latents" + input_shape = (4, 4, 2, 2) + output_shape = (3, 13, 16, 16) + + def get_init_dict(self): + return {"latent_channels": 4, "block_out_channels": (4, 4, 4, 4)} + + def get_dummy_inputs(self, height=2, width=2): + return { + "latents": randn_tensor( + (1, 4, 4, height, width), + generator=torch.Generator().manual_seed(0), + device=torch_device, + dtype=self.torch_dtype, + ) + } + + +class TestEdenLiveLightTAEDecoderModel(EdenLiveLightTAEDecoderTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("patch_size", [1, 2]) + @pytest.mark.parametrize("frames", [1, 3, 4, 5]) + @torch.no_grad() + def test_tiny_frames_causality_and_call_isolation(self, patch_size, frames): + torch.manual_seed(0) + decoder = EdenLiveLightTAEDecoder( + latent_channels=4, patch_size=patch_size, block_out_channels=(4, 4, 4, 4) + ).to(torch_device) + latents = randn_tensor((1, 4, frames, 2, 2), generator=torch.Generator().manual_seed(9), device=torch_device) + first = decoder.decode(latents).sample + # A separate, different decode must not contaminate the next call's temporal memory. + decoder.decode(latents + 1) + repeated = decoder.decode(latents, return_dict=False)[0] + torch.testing.assert_close(first, repeated, atol=0, rtol=0) + assert first.shape == (1, 3, (frames - 1) * 4 + 1, 16 * patch_size, 16 * patch_size) + assert torch.isfinite(first).all() and first.min() >= -1 and first.max() <= 1 + if frames > 1: + changed = latents.clone() + changed[:, :, -1] += 1 + prefix = decoder.decode(changed).sample[:, :, :-4] + torch.testing.assert_close(prefix, first[:, :, :-4], atol=0, rtol=0) + assert not any("cache" in key for key in decoder.state_dict()) + + +class TestEdenLiveLightTAEDecoderMemory(EdenLiveLightTAEDecoderTesterConfig, MemoryTesterMixin): + @pytest.mark.skipif(torch_device == "mps", reason="MPS does not support float8 casting.") + @torch.no_grad() + def test_layerwise_casting_memory(self): + # Every parameter is castable, so the storage footprint is independent of compute dtype. + sizes = [] + for storage, compute in ( + (torch.float32, torch.float32), + (torch.float8_e4m3fn, torch.float32), + (torch.float8_e4m3fn, torch.bfloat16), + ): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device, dtype=compute) + model.enable_layerwise_casting(storage_dtype=storage, compute_dtype=compute) + inputs = {key: value.to(compute) for key, value in self.get_dummy_inputs().items()} + assert torch.isfinite(model(**inputs).sample).all() + sizes.append(model.get_memory_footprint()) + assert sizes[1] == sizes[2] < sizes[0] + + +class TestEdenLiveLightTAEDecoderTorchCompile(EdenLiveLightTAEDecoderTesterConfig, TorchCompileTesterMixin): + different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] diff --git a/tests/models/transformers/test_models_transformer_edenlive.py b/tests/models/transformers/test_models_transformer_edenlive.py new file mode 100644 index 000000000000..d28defc9ebc1 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_edenlive.py @@ -0,0 +1,153 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import pytest +import torch + +from diffusers import EdenLiveTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class EdenLiveTransformerTesterConfig(BaseModelTesterConfig): + model_class = EdenLiveTransformer3DModel + main_input_name = "hidden_states" + input_shape = (4, 3, 8, 8) + output_shape = (4, 3, 8, 8) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self): + return { + "patch_size": (1, 2, 2), + "num_attention_heads": 4, + "attention_head_dim": 12, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "num_layers": 2, + "proxy_channels": 8, + "proxy_input_channels": 16, + "rope_max_seq_len": 64, + } + + def get_dummy_inputs(self, batch_size=1, height=8, width=8): + generator = self.generator + return { + "hidden_states": randn_tensor( + (batch_size, 4, 3, height, width), generator=generator, device=torch_device, dtype=self.torch_dtype + ), + "proxy_features": randn_tensor( + (batch_size, 16, 3, 4, 4), generator=generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (batch_size, 12, 16), generator=generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([500.0], device=torch_device).expand(batch_size), + } + + +class TestEdenLiveTransformerModel(EdenLiveTransformerTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("proxy_dtype", [torch.float32, torch.bfloat16]) + @torch.no_grad() + def test_proxy_projection_changes_conditioning(self, proxy_dtype): + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["proxy_features"] = inputs["proxy_features"].to(proxy_dtype) + expected = model(**inputs).sample + model.proxy_projection.weight.zero_() + model.proxy_projection.bias.zero_() + unconditioned = model(**inputs).sample + assert torch.isfinite(expected).all() and torch.isfinite(unconditioned).all() + assert (expected - unconditioned).abs().max() > 1e-6 + + @torch.no_grad() + def test_causal_state_commits_only_clean_chunks_and_preserves_sink(self): + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + inputs = self.get_dummy_inputs() + tokens_per_frame = (inputs["hidden_states"].shape[-2] // model.config.patch_size[1]) * ( + inputs["hidden_states"].shape[-1] // model.config.patch_size[2] + ) + state = model.create_causal_state( + batch_size=1, + height=8, + width=8, + chunk_size=3, + max_context_frames=18, + sink_frames=6, + device=torch_device, + dtype=torch.float32, + ) + assert state.next_frame == 0 + model(**inputs, causal_state=state) + assert state.next_frame == 0 + assert all(layer.num_tokens == 0 for layer in state.layers) + sink_keys = None + for index in range(14): + inputs["timestep"].zero_() + model(**inputs, causal_state=state, cache_update=True) + assert state.next_frame == 3 * (index + 1) + assert all(layer.num_tokens <= 18 * tokens_per_frame for layer in state.layers) + if index == 1: + sink_keys = [layer.key[:, : 6 * tokens_per_frame].clone() for layer in state.layers] + elif index > 1: + for expected, layer in zip(sink_keys, state.layers): + torch.testing.assert_close(layer.key[:, : 6 * tokens_per_frame], expected, atol=0, rtol=0) + state.reset() + assert state.next_frame == 0 and all(layer.num_tokens == 0 for layer in state.layers) + assert not any("cache" in key for key in model.state_dict()) + + +class TestEdenLiveTransformerMemory(EdenLiveTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestEdenLiveTransformerTorchCompile(EdenLiveTransformerTesterConfig, TorchCompileTesterMixin): + # Keep all shapes in the proxy-resize branch for the zero-recompilation check. + different_shapes_for_compilation = [(8, 8), (8, 16), (16, 16)] + + @torch.no_grad() + def test_torch_compile_without_proxy_resize(self): + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + inputs = self.get_dummy_inputs(height=4, width=4) + expected = model(**inputs).sample + compiled = torch.compile(model, fullgraph=True) + with torch._dynamo.config.patch(error_on_recompile=True): + for _ in range(2): + actual = compiled(**inputs).sample + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + + +class TestEdenLiveTransformerTraining(EdenLiveTransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestEdenLiveTransformerAttention(EdenLiveTransformerTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_edenlive_causal.py b/tests/models/transformers/test_models_transformer_edenlive_causal.py new file mode 100644 index 000000000000..9197fc187854 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_edenlive_causal.py @@ -0,0 +1,167 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import pytest +import torch + +from diffusers.models.transformers import transformer_edenlive + +from ...testing_utils import require_torch_multi_gpu +from .test_models_transformer_edenlive import EdenLiveTransformerTesterConfig + + +class TestEdenLiveCausalCacheOrder: + @pytest.mark.parametrize("use_rotary_emb", [False, True]) + @torch.no_grad() + def test_without_cache_matches_attention_reference(self, use_rotary_emb): + torch.manual_seed(0) + attn = transformer_edenlive.EdenLiveAttention( + dim=8, heads=2, dim_head=4, processor=transformer_edenlive.EdenLiveCausalAttnProcessor() + ).eval() + hidden_states = torch.randn(1, 6, 8) + query = attn.norm_q(attn.to_q(hidden_states)).unflatten(-1, (2, 4)) + key = attn.norm_k(attn.to_k(hidden_states)).unflatten(-1, (2, 4)) + value = attn.to_v(hidden_states).unflatten(-1, (2, 4)) + rotary_emb = None + if use_rotary_emb: + angles = torch.randn(1, 6, 1, 2) + rotations = torch.polar(torch.ones_like(angles), angles) + query = torch.view_as_real(torch.view_as_complex(query.unflatten(-1, (2, 2))) * rotations).flatten(-2) + key = torch.view_as_real(torch.view_as_complex(key.unflatten(-1, (2, 2))) * rotations).flatten(-2) + rotary_emb = (angles.cos().repeat_interleave(2, dim=-1), angles.sin().repeat_interleave(2, dim=-1)) + expected = ( + torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + ) + .transpose(1, 2) + .flatten(2) + ) + expected = attn.to_out[1](attn.to_out[0](expected)) + actual = attn(hidden_states, rotary_emb=rotary_emb) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + + @torch.no_grad() + def test_local_chunk_order_survives_eviction_and_short_tail(self): + torch.manual_seed(0) + attn = transformer_edenlive.EdenLiveAttention( + dim=2, + heads=1, + dim_head=2, + processor=transformer_edenlive.EdenLiveCausalAttnProcessor(), + ).eval() + cache = transformer_edenlive._EdenLiveLayerCache(torch.zeros(1, 6, 1, 2), torch.zeros(1, 6, 1, 2)) + + def hidden_states(ids): + values = torch.tensor(ids, dtype=torch.float32) + return torch.stack((values / 16, (values.remainder(3) - 1) / 2), dim=-1).unsqueeze(0) + + def call(ids, expected_ids, commit): + hidden = hidden_states(ids) + visible = hidden_states(expected_ids) + rope = (torch.ones(1, 1, 1, 2), torch.zeros(1, 1, 1, 2)) + stored_key, stored_value = cache.key.clone(), cache.value.clone() + stored_count = cache.num_tokens + output = attn( + hidden, + rotary_emb=rope, + key_rotary_emb=rope, + layer_cache=cache, + cache_update=commit, + cache_window_tokens=6, + cache_sink_tokens=2, + ) + query = attn.norm_q(attn.to_q(hidden)).unsqueeze(1) + key = attn.norm_k(attn.to_k(visible)).unsqueeze(1) + value = attn.to_v(visible).unsqueeze(1) + expected = torch.nn.functional.scaled_dot_product_attention(query, key, value).squeeze(1) + expected = attn.to_out[1](attn.to_out[0](expected)) + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-6) + if commit: + assert cache.num_tokens == len(expected_ids) + torch.testing.assert_close(cache.key[:, : cache.num_tokens], key.transpose(1, 2), rtol=0, atol=0) + torch.testing.assert_close(cache.value[:, : cache.num_tokens], value.transpose(1, 2), rtol=0, atol=0) + else: + assert cache.num_tokens == stored_count + torch.testing.assert_close(cache.key, stored_key, rtol=0, atol=0) + torch.testing.assert_close(cache.value, stored_value, rtol=0, atol=0) + + call([0, 1], [0, 1], True) + call([4, 5], [0, 1, 4, 5], True) + call([8, 9], [0, 1, 4, 5, 8, 9], True) + call([12, 13], [0, 1, 8, 9, 12, 13], True) + call([16], [0, 1, 9, 12, 13, 16], False) + call([16], [0, 1, 9, 12, 13, 16], True) + + +class TestEdenLiveCausalDeviceMap(EdenLiveTransformerTesterConfig): + @require_torch_multi_gpu + @pytest.mark.parametrize("block_device", ["cuda:1", "cpu"]) + @pytest.mark.parametrize("head_device", ["cuda:0", "cuda:1"]) + @torch.no_grad() + def test_causal_rollout_matches_single_device(self, tmp_path, block_device, head_device): + torch.manual_seed(0) + baseline = self.model_class(**self.get_init_dict()).eval() + baseline.save_pretrained(tmp_path) + device_map = {name: "cuda:0" for name, _ in baseline.named_children() if name != "blocks"} + device_map.update( + { + "scale_shift_table": "cuda:0", + "blocks.0": "cuda:0", + "blocks.1": block_device, + "norm_out": head_device, + "proj_out": head_device, + } + ) + dispatched = self.model_class.from_pretrained(tmp_path, device_map=device_map).eval() + baseline.to("cuda:0") + inputs = {name: value.to("cuda:0") for name, value in self.get_dummy_inputs(height=4, width=4).items()} + torch.testing.assert_close(dispatched(**inputs).sample, baseline(**inputs).sample, atol=1e-5, rtol=1e-5) + state_args = { + "batch_size": 1, + "height": 4, + "width": 4, + "chunk_size": 3, + "max_context_frames": 6, + "sink_frames": 3, + } + reference_state = baseline.create_causal_state(**state_args, device="cuda:0") + state = dispatched.create_causal_state(**state_args, device="cuda:0") + expected_devices = ["cuda:0", "cuda:0" if block_device == "cpu" else block_device] + assert [str(layer.key.device) for layer in state.layers] == expected_devices + assert [str(layer.value.device) for layer in state.layers] == expected_devices + + for num_frames in (3, 3, 3, 1): + chunk_inputs = dict(inputs) + chunk_inputs["hidden_states"] = inputs["hidden_states"][:, :, :num_frames] + chunk_inputs["proxy_features"] = inputs["proxy_features"][:, :, :num_frames] + for commit in (False, True): + chunk_inputs["timestep"] = torch.tensor([0.0 if commit else 500.0], device="cuda:0") + expected = baseline(**chunk_inputs, causal_state=reference_state, cache_update=commit).sample + actual = dispatched(**chunk_inputs, causal_state=state, cache_update=commit).sample + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + assert state.next_frame == reference_state.next_frame + for layer, reference_layer in zip(state.layers, reference_state.layers): + assert layer.num_tokens == reference_layer.num_tokens + torch.testing.assert_close(layer.key.cpu(), reference_layer.key.cpu(), atol=1e-5, rtol=1e-5) + torch.testing.assert_close(layer.value.cpu(), reference_layer.value.cpu(), atol=1e-5, rtol=1e-5) + + state.reset() + reference_state.reset() + assert state.next_frame == 0 and all(layer.num_tokens == 0 for layer in state.layers) + torch.testing.assert_close( + dispatched(**inputs, causal_state=state).sample, + baseline(**inputs, causal_state=reference_state).sample, + atol=1e-5, + rtol=1e-5, + ) diff --git a/tests/pipelines/edenlive/__init__.py b/tests/pipelines/edenlive/__init__.py new file mode 100644 index 000000000000..b8ca19521280 --- /dev/null +++ b/tests/pipelines/edenlive/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. diff --git a/tests/pipelines/edenlive/test_pipeline_edenlive.py b/tests/pipelines/edenlive/test_pipeline_edenlive.py new file mode 100644 index 000000000000..fb89a2b6b33f --- /dev/null +++ b/tests/pipelines/edenlive/test_pipeline_edenlive.py @@ -0,0 +1,240 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import numpy as np +import pytest +import torch +from PIL import Image +from transformers import ( + AutoTokenizer, + CLIPImageProcessor, + CLIPVisionConfig, + CLIPVisionModel, + DINOv3ViTConfig, + DINOv3ViTModel, + UMT5Config, + UMT5EncoderModel, +) + +from diffusers import ( + AutoencoderKLWan, + EdenLiveFlowMatchEulerDiscreteScheduler, + EdenLiveLightTAEDecoder, + EdenLivePipeline, + EdenLiveTransformer3DModel, + UniPCMultistepScheduler, +) + +from ...testing_utils import torch_device +from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_PARAMS, VIDEO_TO_VIDEO_BATCH_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +class EdenLivePipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = EdenLivePipeline + profile = "5b" + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS | {"video"} + optional_input_params = frozenset( + ["num_frames", "num_inference_steps", "guidance_scale", "generator", "latents", "output_type", "return_dict"] + ) + batch_input_params = VIDEO_TO_VIDEO_BATCH_PARAMS | {"image"} + output_shape = (9, 3, 16, 16) + + def get_dummy_components(self): + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=4, + z_dim=4, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + ) + torch.manual_seed(0) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = UMT5EncoderModel( + UMT5Config( + vocab_size=len(tokenizer), d_model=32, d_ff=32, d_kv=8, num_layers=1, num_heads=4, dropout_rate=0.0 + ) + ).eval() + torch.manual_seed(0) + proxy_encoder = DINOv3ViTModel( + DINOv3ViTConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + patch_size=16, + image_size=32, + num_register_tokens=4, + ) + ).eval() + torch.manual_seed(0) + transformer = EdenLiveTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=4, + attention_head_dim=12, + in_channels=4 if self.profile == "5b" else 12, + out_channels=4, + proxy_input_channels=32, + proxy_channels=8, + condition_rmsnorm=self.profile == "14b", + text_dim=32, + freq_dim=16, + ffn_dim=64, + num_layers=1, + image_dim=32 if self.profile == "14b" else None, + added_kv_proj_dim=48 if self.profile == "14b" else None, + rope_max_seq_len=64, + ) + if self.profile == "14b": + torch.manual_seed(0) + image_encoder = CLIPVisionModel( + CLIPVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + image_size=16, + patch_size=8, + projection_dim=32, + ) + ).eval() + image_processor = CLIPImageProcessor(size={"shortest_edge": 16}, crop_size={"height": 16, "width": 16}) + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(shift=3) + else: + image_encoder, image_processor = None, None + scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5) + return { + "vae": vae, + "decoder": None, + "proxy_image_height": 32, + "proxy_image_width": 32, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "scheduler": scheduler, + "proxy_encoder": proxy_encoder, + "image_encoder": image_encoder, + "image_processor": image_processor, + } + + def get_dummy_inputs(self): + generator = np.random.default_rng(0) + return { + "prompt": "a character walks through snow", + "negative_prompt": "blurry", + "image": Image.fromarray(generator.integers(0, 256, (16, 16, 3), dtype=np.uint8)), + "video": generator.random((9, 16, 16, 3), dtype=np.float32), + "height": 16, + "width": 16, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "generator": self.get_generator(0), + "max_sequence_length": 16, + "output_type": "pt", + } + + +class EdenLive14BRequiredComponentsTesterMixin: + def test_save_load_optional_components(self, tmp_path): + pipe = self.get_pipeline() + pipe.register_modules(image_encoder=None, image_processor=None) + pipe.save_pretrained(tmp_path) + restored = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) + assert restored.image_encoder is None and restored.image_processor is None + with pytest.raises(ValueError, match="CLIP image encoder"): + restored(**self.get_dummy_inputs()) + + +class EdenLiveOptionalDecoderTesterMixin: + def test_lighttae_save_load_and_short_tail(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + torch.manual_seed(0) + pipe.register_modules( + decoder=EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)).to(torch_device) + ) + inputs = self.get_dummy_inputs() + inputs["video"] = np.concatenate([inputs["video"], inputs["video"][:4]]) + inputs["num_frames"] = 13 + expected = pipe(**inputs).frames + assert expected.shape == (1, 13, 3, 16, 16) + pipe.save_pretrained(tmp_path) + restored = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) + assert isinstance(restored.decoder, EdenLiveLightTAEDecoder) + inputs["generator"] = self.get_generator(0) + torch.testing.assert_close(restored(**inputs).frames, expected, atol=1e-5, rtol=0) + + +class TestEdenLivePipeline(EdenLiveOptionalDecoderTesterMixin, EdenLivePipelineTesterConfig, PipelineTesterMixin): + @pytest.mark.parametrize("num_frames", [0, 4, 8, 13]) + def test_invalid_frame_count(self, num_frames): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["num_frames"] = num_frames + with pytest.raises(ValueError, match="frames|frame"): + pipe(**inputs) + + def test_automatic_frame_prefix(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["num_frames"] = None + inputs["video"] = np.concatenate([inputs["video"], inputs["video"][:2]]) + output = pipe(**inputs).frames + assert output.shape == (1, *self.output_shape) + + @pytest.mark.parametrize("output_type", ["np", "pt", "pil", "latent"]) + def test_output_types(self, output_type): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["output_type"] = output_type + output = pipe(**inputs).frames + if output_type == "np": + assert output.shape == (1, 9, 16, 16, 3) + elif output_type == "pt": + assert output.shape == (1, *self.output_shape) + elif output_type == "latent": + assert output.shape == (1, 4, 3, 2, 2) + else: + assert len(output) == 1 and len(output[0]) == 9 + assert isinstance(output[0][0], Image.Image) + + @pytest.mark.parametrize("condition", ["image", "video", "prompt"]) + def test_conditions_change_generation(self, condition): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + expected = pipe(**inputs).frames + inputs = self.get_dummy_inputs() + if condition == "prompt": + inputs[condition] = "a racing car on a sunny road" + elif condition == "image": + inputs[condition] = Image.new("RGB", (16, 16), "white") + else: + inputs[condition] = 1 - inputs[condition] + actual = pipe(**inputs).frames + assert (actual - expected).abs().max() > 1e-6 + + +class TestEdenLive14BPipeline(EdenLive14BRequiredComponentsTesterMixin, TestEdenLivePipeline): + profile = "14b" + + +class TestEdenLivePipelineMemory(EdenLivePipelineTesterConfig, MemoryTesterMixin): + pass + + +class TestEdenLive14BPipelineMemory(EdenLivePipelineTesterConfig, MemoryTesterMixin): + profile = "14b" diff --git a/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py b/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py new file mode 100644 index 000000000000..7ca96012dc7e --- /dev/null +++ b/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py @@ -0,0 +1,234 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import numpy as np +import pytest +import torch + +from diffusers import ( + EdenLiveCausalPipeline, + EdenLiveFlowMatchEulerDiscreteScheduler, +) + +from ...testing_utils import torch_device +from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils.common import cast_pipeline_to_dtype +from .test_pipeline_edenlive import ( + EdenLive14BRequiredComponentsTesterMixin, + EdenLiveOptionalDecoderTesterMixin, + EdenLivePipelineTesterConfig, +) + + +class EdenLiveCausalPipelineTesterConfig(EdenLivePipelineTesterConfig): + pipeline_class = EdenLiveCausalPipeline + + def get_dummy_components(self): + components = super().get_dummy_components() + components["scheduler"] = EdenLiveFlowMatchEulerDiscreteScheduler( + stochastic_sampling=True, default_timesteps=[1000, 500] + ) + return components + + +class TestEdenLiveCausalPipeline( + EdenLiveOptionalDecoderTesterMixin, EdenLiveCausalPipelineTesterConfig, PipelineTesterMixin +): + def test_load_causal_pipeline_from_root(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + expected = pipe(**self.get_dummy_inputs()).frames + pipe.save_pretrained(tmp_path) + restored = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) + assert restored.decoder is None + assert isinstance(restored.scheduler, EdenLiveFlowMatchEulerDiscreteScheduler) + torch.testing.assert_close(restored(**self.get_dummy_inputs()).frames, expected, atol=1e-5, rtol=0) + + def test_short_tail_and_no_state_leak(self): + pipe = self.get_pipeline().to(torch_device) + + def inputs(): + result = self.get_dummy_inputs() + result["video"] = np.concatenate([result["video"], result["video"][:4]]) + result["num_frames"] = 13 + result["chunk_prompts"] = ["snow", "fire"] + return result + + expected = pipe(**inputs()).frames + actual = pipe(**inputs()).frames + assert actual.shape == (1, 13, 3, 16, 16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert not any("cache" in name for name in pipe.transformer.state_dict()) + + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_chunk_prompt_changes_tail(self, guidance): + pipe = self.get_pipeline().to(torch_device) + + def inputs(tail): + result = self.get_dummy_inputs() + result["video"] = np.concatenate([result["video"], result["video"][:4]]) + result.update(num_frames=13, chunk_prompts=["snow", tail], output_type="latent", guidance_scale=guidance) + return result + + first = pipe(**inputs("fire")).frames + second = pipe(**inputs("ocean")).frames + torch.testing.assert_close(first[:, :, :3], second[:, :, :3], atol=0, rtol=0) + assert (first[:, :, 3:] - second[:, :, 3:]).abs().max() > 1e-6 + + def test_invalid_chunk_prompts(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["chunk_prompts"] = ["one", "too many"] + before = inputs["generator"].get_state().clone() + with pytest.raises(ValueError, match="chunk_prompts"): + pipe(**inputs) + assert torch.equal(before, inputs["generator"].get_state()) + + def test_chunk_prompts_reject_batched_prompt(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs.update(prompt=["snow", "fire"], chunk_prompts=["ocean"]) + with pytest.raises(ValueError, match="batch size one"): + pipe(**inputs) + + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_chunk_prompts_encode_unique_texts_individually(self, guidance): + fixture = self + pipe = fixture.get_pipeline().to(torch_device) + encoded, used = [], [] + + def record_text(module, args, kwargs): + encoded.append((args[0].clone(), kwargs["attention_mask"].clone())) + + def record_dit(module, args, kwargs): + if kwargs.get("cache_update"): + used.append(args[2].clone()) + + pipe.text_encoder.register_forward_pre_hook(record_text, with_kwargs=True) + pipe.transformer.register_forward_pre_hook(record_dit, with_kwargs=True) + inputs = fixture.get_dummy_inputs() + inputs.update( + video=np.tile(inputs["video"], (4, 1, 1, 1))[:29], + num_frames=29, + chunk_prompts=["snow", "fire", "snow"], + guidance_scale=guidance, + output_type="latent", + ) + pipe(**inputs) + expected_texts = ["snow", "fire"] + ([inputs["negative_prompt"]] if guidance > 1 else []) + assert len(encoded) == len(expected_texts) + for (ids, mask), text in zip(encoded, expected_texts): + tokens = pipe.tokenizer( + [text], + padding="max_length", + max_length=inputs["max_sequence_length"], + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + torch.testing.assert_close(ids.cpu(), tokens.input_ids) + torch.testing.assert_close(mask.cpu(), tokens.attention_mask) + conditional = used[::2] if guidance > 1 else used + assert len(conditional) == 3 + torch.testing.assert_close(conditional[0], conditional[2], rtol=0, atol=0) + assert not torch.equal(conditional[0], conditional[1]) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_clean_commit_timestep_shape_and_cfg(self, dtype): + fixture = self + pipe = cast_pipeline_to_dtype(fixture.get_pipeline().to(torch_device), dtype) + calls = [] + + def capture(module, args, kwargs): + calls.append((args[1].clone(), kwargs.get("cache_update", False))) + + pipe.transformer.register_forward_pre_hook(capture, with_kwargs=True) + inputs = fixture.get_dummy_inputs() + inputs.update( + video=np.concatenate([inputs["video"], inputs["video"][:4]]), + num_frames=13, + output_type="latent", + guidance_scale=3.0, + ) + pipe(**inputs) + commits = [time for time, commit in calls if commit] + assert len(commits) == 4 + assert all(time.dtype == torch.float32 and not time.any() for time in commits) + assert all(time.shape == (1,) for time in commits) + noisy = calls[0][0] + assert noisy.shape == (1, 3) and noisy[0, 0] == 0 and (noisy[0, 1:] > 0).all() + + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_repeated_chunk_prompt_matches_static(self, guidance): + pipe = self.get_pipeline().to(torch_device) + + def inputs(): + result = self.get_dummy_inputs() + result.update( + video=np.tile(result["video"], (2, 1, 1, 1))[:13], + num_frames=13, + prompt="snow", + guidance_scale=guidance, + output_type="latent", + ) + return result + + expected = pipe(**inputs()).frames + actual = pipe(**inputs(), chunk_prompts=["snow", "snow"]).frames + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_next_call_is_independent_after_failure(self): + pipe = self.get_pipeline().to(torch_device) + + def inputs(): + result = self.get_dummy_inputs() + result.update( + video=np.tile(result["video"], (2, 1, 1, 1))[:13], + num_frames=13, + chunk_prompts=["snow", "fire"], + guidance_scale=3.0, + output_type="latent", + ) + return result + + expected = pipe(**inputs()).frames + states = [] + + def fail_in_second_chunk(module, args, kwargs): + state = kwargs["causal_state"] + if not any(state is seen for seen in states): + states.append(state) + if state.next_frame >= 3: + raise RuntimeError("injected model failure") + + handle = pipe.transformer.register_forward_pre_hook(fail_in_second_chunk, with_kwargs=True) + try: + with pytest.raises(RuntimeError, match="injected model failure"): + pipe(**inputs()) + finally: + handle.remove() + assert len(states) == 2 + actual = pipe(**inputs()).frames + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +class TestEdenLive14BCausalPipeline(EdenLive14BRequiredComponentsTesterMixin, TestEdenLiveCausalPipeline): + profile = "14b" + + +class TestEdenLiveCausalPipelineMemory(EdenLiveCausalPipelineTesterConfig, MemoryTesterMixin): + pass + + +class TestEdenLive14BCausalPipelineMemory(EdenLiveCausalPipelineTesterConfig, MemoryTesterMixin): + profile = "14b" diff --git a/tests/schedulers/test_scheduler_edenlive.py b/tests/schedulers/test_scheduler_edenlive.py new file mode 100644 index 000000000000..be5af9b4cb2d --- /dev/null +++ b/tests/schedulers/test_scheduler_edenlive.py @@ -0,0 +1,104 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. +# +# 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. + +import tempfile +import unittest + +import torch + +from diffusers import EdenLiveFlowMatchEulerDiscreteScheduler + + +class EdenLiveFlowMatchEulerDiscreteSchedulerTests(unittest.TestCase): + def test_shifted_schedule(self): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(shift=3.0) + scheduler.set_timesteps(30) + self.assertEqual(scheduler.sigmas.dtype, torch.float32) + self.assertEqual(len(scheduler.timesteps), 30) + torch.testing.assert_close( + scheduler.sigmas[[0, 1, 15, 29, 30]], + torch.tensor([0.9999989867210388, 0.9886479377746582, 0.7503749132156372, 0.09629422426223755, 0.0]), + rtol=0, + atol=0, + ) + + def test_explicit_timesteps_are_not_shifted(self): + for points in ([1000, 937.5, 833.3333333333, 625], [1000, 900, 750, 500]): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(shift=9.0, default_timesteps=points) + scheduler.set_timesteps(4) + torch.testing.assert_close(scheduler.sigmas[:-1], torch.tensor(points) / 1000, rtol=0, atol=0) + scheduler.set_timesteps(sigmas=[1.0, 0.5, 0.0]) + torch.testing.assert_close(scheduler.sigmas, torch.tensor([1.0, 0.5, 0.0]), rtol=0, atol=0) + self.assertEqual(len(scheduler.timesteps), 2) + + def test_stochastic_dtype_and_variance_noise(self): + for dtype in (torch.float32, torch.bfloat16): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True) + scheduler.set_timesteps(timesteps=[1000, 900]) + generator = torch.Generator().manual_seed(9) + sample = torch.randn((2, 3, 4), generator=generator).to(dtype) + prediction = torch.randn(sample.shape, generator=generator).to(dtype) + noise_generator = torch.Generator().manual_seed(12) + noise = torch.randn(sample.shape, generator=noise_generator, dtype=dtype) + explicit = scheduler.step(prediction, scheduler.timesteps[0], sample, variance_noise=noise).prev_sample + scheduler.set_timesteps(timesteps=[1000, 900]) + generated = scheduler.step( + prediction, scheduler.timesteps[0], sample, generator=torch.Generator().manual_seed(12) + ).prev_sample + self.assertEqual(explicit.dtype, dtype) + torch.testing.assert_close(explicit, generated, rtol=0, atol=0) + + def test_final_step_does_not_advance_generator(self): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True) + scheduler.set_timesteps(timesteps=[500]) + generator = torch.Generator().manual_seed(7) + state = generator.get_state().clone() + sample = torch.ones((1, 2, 3), dtype=torch.bfloat16) + result = scheduler.step(sample, scheduler.timesteps[0], sample, generator=generator).prev_sample + self.assertTrue(torch.equal(state, generator.get_state())) + self.assertTrue(torch.equal(result, sample * 0.5)) + + def test_chunk_reset_and_serialization(self): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler( + stochastic_sampling=True, default_timesteps=[1000, 900, 750, 500] + ) + sample = torch.ones((1, 2, 3)) + scheduler.set_timesteps(4) + result = scheduler.step(sample, scheduler.timesteps[0], sample, variance_noise=sample).prev_sample + self.assertEqual(scheduler.step_index, 1) + scheduler.set_timesteps(4) + self.assertIsNone(scheduler.step_index) + repeated = scheduler.step(sample, scheduler.timesteps[0], sample, variance_noise=sample).prev_sample + torch.testing.assert_close(result, repeated, rtol=0, atol=0) + with tempfile.TemporaryDirectory() as directory: + scheduler.save_pretrained(directory) + restored = EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained(directory) + restored.set_timesteps(4) + torch.testing.assert_close(restored.sigmas, scheduler.sigmas, rtol=0, atol=0) + self.assertIsNone(restored.step_index) + self.assertTrue(restored.config.stochastic_sampling) + + def test_invalid_schedule(self): + scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(default_timesteps=[1000, 500]) + for kwargs in ( + {"timesteps": [1000], "sigmas": [1]}, + {"timesteps": [500, 1000]}, + {"timesteps": [1000, 1000, 500]}, + {"timesteps": [float("nan")]}, + {"sigmas": []}, + {"sigmas": [1.1, 0.5]}, + {"num_inference_steps": 4}, + ): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + scheduler.set_timesteps(**kwargs)