From 04262407f2ed37ff7fce8bb1f9ad627e1b7883e2 Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Fri, 16 Oct 2026 05:22:27 +0800 Subject: [PATCH 1/6] feat: add EdenLive v2v bidir&causal inference pipelines --- docs/source/en/_toctree.yml | 6 + .../en/api/models/edenlive_transformer_3d.md | 95 ++ docs/source/en/api/pipelines/edenlive.md | 227 +++ docs/source/en/api/schedulers/edenlive.md | 81 + examples/edenlive/README.md | 113 ++ examples/edenlive/inference_edenlive.py | 220 +++ examples/edenlive/requirements.txt | 9 + src/diffusers/__init__.py | 14 + src/diffusers/models/__init__.py | 10 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_edenlive.py | 182 ++ src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_edenlive.py | 1040 ++++++++++++ src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/edenlive/__init__.py | 41 + .../pipelines/edenlive/loading_utils.py | 523 ++++++ .../pipelines/edenlive/pipeline_edenlive.py | 740 ++++++++ .../pipelines/edenlive/pipeline_output.py | 37 + .../pipelines/edenlive/pipeline_utils.py | 377 +++++ .../pipelines/edenlive/video_processor.py | 79 + src/diffusers/schedulers/__init__.py | 2 + .../schedulers/scheduling_edenlive.py | 211 +++ src/diffusers/utils/dummy_pt_objects.py | 75 + .../dummy_torch_and_transformers_objects.py | 30 + tests/pipelines/edenlive/__init__.py | 0 tests/pipelines/edenlive/test_edenlive.py | 1499 +++++++++++++++++ 26 files changed, 5615 insertions(+) create mode 100644 docs/source/en/api/models/edenlive_transformer_3d.md create mode 100644 docs/source/en/api/pipelines/edenlive.md create mode 100644 docs/source/en/api/schedulers/edenlive.md create mode 100644 examples/edenlive/README.md create mode 100644 examples/edenlive/inference_edenlive.py create mode 100644 examples/edenlive/requirements.txt create mode 100644 src/diffusers/models/autoencoders/autoencoder_edenlive.py create mode 100644 src/diffusers/models/transformers/transformer_edenlive.py create mode 100644 src/diffusers/pipelines/edenlive/__init__.py create mode 100644 src/diffusers/pipelines/edenlive/loading_utils.py create mode 100644 src/diffusers/pipelines/edenlive/pipeline_edenlive.py create mode 100644 src/diffusers/pipelines/edenlive/pipeline_output.py create mode 100644 src/diffusers/pipelines/edenlive/pipeline_utils.py create mode 100644 src/diffusers/pipelines/edenlive/video_processor.py create mode 100644 src/diffusers/schedulers/scheduling_edenlive.py create mode 100644 tests/pipelines/edenlive/__init__.py create mode 100644 tests/pipelines/edenlive/test_edenlive.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index ab98a5243059..e6eb72fe4fb9 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -345,6 +345,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 @@ -721,6 +723,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 @@ -790,6 +794,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..ea8f29733647 --- /dev/null +++ b/docs/source/en/api/models/edenlive_transformer_3d.md @@ -0,0 +1,95 @@ + + +# EdenLiveTransformer3DModel + +The EdenLive diffusion transformer predicts flow velocities from video latents, text embeddings, and projected +DINO proxy features. The same class supports full-video bidirectional attention and causal attention over successive +chunks. The pipeline loads the configuration and weights for the selected profile; see the +[pipeline guide](../pipelines/edenlive) for original-checkpoint loading and video generation. + +## Inputs and output + +Call `forward()` with the following inputs. Here `B` is batch size, `T` is the number of latent frames in the video or +current chunk, and `H` and `W` are latent spatial dimensions. + +| Argument | Shape or format | Meaning | +|---|---|---| +| `hidden_states` | `(B, in_channels, T, H, W)` | Noisy latents, plus image-conditioning channels for 14B | +| `timestep` | Scalar, `(B,)`, `(B, T)`, or `(B, num_patch_tokens)` | Model time; retain FP32 even with BF16 model weights | +| `encoder_hidden_states` | `(B, sequence_length, text_dim)` | UMT5 text embeddings | +| `encoder_hidden_states_image` | `(B, image_sequence_length, image_dim)` | CLIP reference-image embeddings, required for 14B and unused by 5B | +| `proxy_features` | `(B, proxy_channels, T, proxy_height, proxy_width)` | Required projected DINO features; spatial dimensions are interpolated to match the latents | + +The temporal patch size is one; latent height and width must be divisible by the spatial patch size, which is two +for the released profiles. Proxy features must match the latent batch and frame counts. The 5B model uses 48 input +and output latent channels. The 14B model receives 36 input channels: 16 noisy latent channels, 16 image latent +channels, and four mask channels; it predicts 16 output channels. Proxy channels are appended internally and are +not included in `in_channels`. + +With `return_dict=True`, `forward()` returns `Transformer2DModelOutput.sample` with shape +`(B, out_channels, T, H, W)`. With `return_dict=False`, it returns the same prediction as the first tuple entry. +The prediction is a latent flow velocity, not a decoded video. + +The transformer's `text_timeline` is an internal token layout for 5B bidirectional inference. It contains `lengths` +(positive query-token counts per segment and sample) and `text_length` (the number of tokens per text prompt). +User-facing timelines in seconds are accepted and converted by `EdenLivePipeline`; they cannot be passed directly +to the transformer. + +## Causal attention state + +Use `create_causal_state(batch_size, height, width, ...)` to allocate state for a rollout, with height and width in +latent coordinates. The released defaults are three latent frames per chunk, a maximum context of 18 latent frames, +and six retained sink frames. Allocate a separate state for each classifier-free guidance branch. + +Denoising calls pass `causal_state=state` and `cache_update=False`. After a chunk has been denoised, a clean call +at timestep zero with `cache_update=True` commits its keys and values and advances the frame cursor. A short chunk +is supported only at the end of the rollout. The pipeline manages these calls and releases its states after every +invocation. For custom loops, call `state.reset()` before beginning an independent rollout. + +## EdenLiveTransformer3DModel API + +[[autodoc]] EdenLiveTransformer3DModel + - forward + - create_causal_state + +## EdenLiveProxyProjection + +The trained projection maps the last channel dimension of DINO features to `out_channels`, applies L2 normalization +on that dimension, and returns FP32 features. Inputs can retain arbitrary leading batch, frame, or spatial token +dimensions. The pipeline reshapes the result into the channel-first video layout expected by the transformer. + +Its `in_channels` must match the DINO encoder feature width and `out_channels` must match the transformer's +`proxy_channels`. Use the released trained projection together with its DINO encoder. + +[[autodoc]] EdenLiveProxyProjection + - forward + +## EdenLiveCausalState + +State belongs to one rollout and guidance branch and is not part of the model checkpoint. Create it through the +transformer so cache allocation matches the model's layers, dtype, device, and context-parallel layout. + +| Field | Meaning | +|---|---| +| `layers` | Per-layer unrotated keys, values, and valid-token counts | +| `batch_size`, `height`, `width` | Batch and latent spatial dimensions used to allocate the caches | +| `chunk_size`, `max_context_frames`, `sink_frames` | Causal context limits in latent frames | +| `tokens_per_frame` | Spatial patch-token count per latent frame | +| `cp_size`, `cp_rank` | Context-parallel degree and local rank | +| `next_frame` | Number of latent frames committed so far; also the next chunk's position offset | + +`reset()` clears valid-token counts and resets `next_frame` to zero, allowing the allocation to be reused for an +independent rollout with the same shape and parallel layout. + +[[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..b3f4bc648e35 --- /dev/null +++ b/docs/source/en/api/pipelines/edenlive.md @@ -0,0 +1,227 @@ + + +# EdenLive + +[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) generates video from a reference image, text, and a +proxy video. It uses a trained DINOv3 encoder and projection to condition a Wan-based diffusion transformer. + +[`EdenLivePipeline`] supports both bidirectional generation and successive causal chunks with clean latent history. +Select a released configuration with `profile` when loading the original checkpoint repository. Defaults apply when +`num_inference_steps` and `guidance_scale` are omitted: + +| Profile | Sampling | Default steps | Guidance scale | Default height x width | +|---|---|---|---|---| +| `5b-bidir` | UniPC over the full video | 48 | 3 | 704 x 1280 | +| `14b-bidir` | Euler over the full video | 30 | 3 | 480 x 832 | +| `5b-causal` | Stochastic Euler per chunk | 4 | 1 | 704 x 1280 | +| `14b-causal` | Stochastic Euler per chunk | 4 | 1 | 480 x 832 | + +The released causal checkpoints use fixed four-step schedules. Changing `num_inference_steps` alone to another +value raises an error; it does not select a new distilled schedule. Guidance greater than one enables the negative +prompt branch. + +## Generate a video + +```py +import torch +from diffusers import EdenLivePipeline +from diffusers.utils import export_to_video, load_image, load_video + +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive", profile="5b-causal" +).to("cuda") + +frames = pipe( + prompt="A character explores a snowy landscape.", + image=load_image("reference.png"), + video=load_video("proxy-24fps.mp4"), + num_frames=49, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(frames, "edenlive.mp4", fps=24) +``` + +## Load and save checkpoints + +Use `EdenLivePipeline.from_pretrained()` for both formats: + +| Format | Required metadata | Profile selection | +|---|---|---| +| Original release | Root `model_index.json` with the EdenLive original-format marker | Pass one of the four complete profile names | +| Saved Diffusers pipeline | `model_index.json` and saved component directories | Omit `profile` to use the saved value | + +The original release loader reads the existing checkpoint files without writing converted weights. It fetches only +the files needed for the selected profile and components. The same entry accepts a Hub repository ID or local +directory, plus `revision`, `cache_dir`, `token`, and `local_files_only=True` for cached or offline loading. + +The original release's default precision is BF16 for the transformer, VAE, UMT5 text encoder, CLIP encoder, and +LightTAE, with FP32 parameters for the DINO encoder and projection. An explicit `dtype` (or `torch_dtype`) can be +a single dtype or a mapping by component name. Initialized component overrides retain their identity and dtype. +Active `device_map`, `variant`, and quantization options are not supported by the original-format loader. Load on CPU +and then place the pipeline on the desired device, as in the example above. + +Save and reload a standard pipeline with: + +```py +pipe.save_pretrained("edenlive-saved") +pipe = EdenLivePipeline.from_pretrained("edenlive-saved").to("cuda") +``` + +The saved configuration retains the full profile and decoder choice. A conflicting explicit profile raises an error. +Older standard snapshots with size-only profiles or the former causal class name are normalized in memory; newly +saved snapshots use `EdenLivePipeline`. The generic `DiffusionPipeline.from_pretrained()` entry is not supported for +original releases. + +## Prepare the inputs + +The reference image and proxy video should depict the same composition. Supply a full RGB proxy video already +sampled at the desired output frame rate. The Python pipeline does not resample timestamps; the command-line +example below does. The pipeline then selects proxy frames at the VAE temporal stride once. + +| Input | Accepted layout and values | +|---|---| +| `image` | RGB PIL image, or floating-point NumPy/tensor image with values in [0, 1] | +| `video` | List of RGB PIL frames; tensor `(T, 3, H, W)` or `(B, T, 3, H, W)`; NumPy `(T, H, W, 3)` or `(B, T, H, W, 3)` with floating-point values in [0, 1] | +| `prompt` | String or list of strings; alternatively provide `prompt_embeds` or a supported temporal prompt format | +| `height`, `width` | Output dimensions, both multiples of 32 for 5B or 16 for 14B | +| `num_frames` | Positive frame count of the form \\(4k+1\\), no greater than the input video length | + +Here `B` is the video batch size, `T` the decoded frame count, and `H` and `W` the frame height and width. The reference +image is resized and center-cropped to the output dimensions. Proxy RGB frames are processed at `proxy_image_size` +(480 x 832 by default) with the profile's `proxy_resize_mode` before DINO encoding. The command-line example also +crops and resizes proxy frames to the output aspect ratio before passing them to the pipeline. Omitting `num_frames` +selects the largest valid input prefix; missing frames are not padded or repeated by the pipeline. + +Both model sizes use temporal compression by four. A video of \\(F\\) frames has +\\(T_{\mathrm{latent}}=(F-1)/4+1\\) latent frames. The 5B model uses 48 latent channels and spatial compression by 16; +the 14B model uses 16 channels and spatial compression by 8. + +## Choose a decoder + +Original releases load the matching [`EdenLiveLightTAEDecoder`] by default. Use the full Wan VAE with: + +```py +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive", profile="14b-causal", decoder=None +) +``` + +Pass `decoder=my_decoder` to use an initialized compatible decoder. The full VAE remains available to encode the +reference image. A custom decoder must implement `decode(latents, return_dict=False)` and return a tuple whose +first item is an RGB tensor in `(B, 3, T, H, W)` layout and `[-1, 1]` range. Its `config.latent_channels`, `scale_factor_temporal`, and +`scale_factor_spatial` must match the VAE. Use a compatible `ModelMixin` decoder for pipeline-wide device placement +and standard serialization. Plain `torch.nn.Module` decoders support instance loading and decoding with explicit +component placement. A decoder receives raw, unnormalized VAE latents; the pipeline applies the inverse latent +normalization before calling it. Selecting `output_type="latent"` bypasses this conversion and decoding. + +## Temporal prompts + +The `5b-bidir` profile accepts `text_timeline` instead of a single prompt. Each segment contains `start_sec`, `end_sec`, +and `text`. For example, this timeline describes 49 frames sampled at 24 fps: + +```py +timeline = { + "fps": 24, + "num_frames": 49, + "segments": [ + {"start_sec": 0.0, "end_sec": 1.0, "text": "A character walks through snow."}, + {"start_sec": 1.0, "end_sec": 2.0, "text": "The character approaches a warm campfire."}, + ], +} +``` + +Pass it as `text_timeline=timeline` with the image and video. Segments must start at zero, be contiguous, and cover +the timestamp of the last requested frame. Each segment must receive at least one latent frame. A latent frame +exactly at a segment end belongs to the segment on its left. The timeline frame count must match the requested +video, and its `fps` describes the supplied frames; it does not trigger resampling. For a batch, provide one timeline +dictionary per video. Timeline prompts cannot be combined with precomputed text embeddings. Other profiles reject +`text_timeline`. + +Causal profiles accept `chunk_prompts`, a list containing one complete prompt per latent chunk. For the released +three-frame chunk size, 49 decoded frames produce 13 latent frames and require five prompts; the last chunk has +one latent frame. This interface requires batch size one and cannot be combined with precomputed positive text +embeddings. A single `prompt` applies to all chunks when `chunk_prompts` is omitted. Bidirectional profiles reject +`chunk_prompts`. + +Each pipeline call starts with fresh causal history. The released 18-frame latent context includes six retained +sink frames and the current chunk. Within a call, completed clean chunks are committed to the cache before the +next chunk is generated. + +## Multiple GPUs and the command-line example + +The inference example resamples video timestamps and uses Ulysses context parallelism with PyTorch FSDP2 on the +same ranks. Run with Python for one GPU or `torchrun` for two or four GPUs: + +```bash +torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ + --model VAST-AI/EdenLive --profile 14b-causal \ + --video proxy.mp4 --image reference.png \ + --prompt "A character explores a snowy landscape." \ + --num-frames 81 --output edenlive.mp4 +``` + +Use `--full-vae` for full VAE decoding. Only rank zero exports the video; the example also records timing and peak +memory. Keep the parallel degree fixed when comparing reproducible outputs: BF16 matrix multiplication shapes +can change predictions, as also observed in the reference implementation. + +The example defaults to 24 fps for 5B and 16 fps for 14B. See the +[example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive) for installation, single-GPU +usage, temporal prompt JSON files, and the timing report format. + +## EdenLivePipeline + +The constructor registers initialized components. `from_pretrained()` loads a release or saved pipeline; +`__call__()` prepares the image, proxy features, text, and initial noise before running the profile's sampling path. +Use `generator` for reproducible noise and `callback_on_step_end` to inspect or replace supported tensors after a +step. In causal mode callbacks receive the current chunk, while their step indices count across all chunks. + +[[autodoc]] EdenLivePipeline + - all + - __call__ + +## EdenLiveLightTAEDecoder + +This decoder implements the released Tiny decoding path without an encoder. The 5B variant uses 48 latent channels +and `patch_size=2`; the 14B variant uses 16 channels and `patch_size=1`. `decode()` accepts raw VAE latents in +`(B, C, T_latent, H_latent, W_latent)` layout and returns `DecoderOutput.sample` in `(B, 3, T, H, W)` layout with +values in `[-1, 1]`. Unlike pipeline frame output, its channel axis precedes the time axis. + +To decode `output_type="latent"` yourself, invert the VAE normalization first: + +\\[ +z_{\mathrm{raw}} = z_{\mathrm{diffusion}} \odot \sigma + \mu. +\\] + +Here \\(\mu\\) and \\(\sigma\\) are the channel-wise `vae.config.latents_mean` and `vae.config.latents_std`, and +\\(\odot\\) denotes element-wise multiplication with broadcasting. Move the unnormalized latents to the decoder's +device and dtype before decoding. The pipeline handles this when returning decoded frames. Every `decode()` call +creates fresh temporal state and supports a short final chunk; separate calls do not share decoder history. + +[[autodoc]] EdenLiveLightTAEDecoder + - decode + +## EdenLivePipelineOutput + +With `return_dict=True`, read the result from `.frames`: + +| `output_type` | `.frames` layout | Values | +|---|---|---| +| `"np"` (default) | NumPy `(B, T, H, W, 3)` | Floating-point RGB in [0, 1] | +| `"pt"` | Tensor `(B, T, 3, H, W)` | Floating-point RGB in [0, 1] | +| `"pil"` | One list of PIL RGB frames per video | PIL images | +| `"latent"` | Tensor `(B, C, T_latent, H_latent, W_latent)` | Normalized diffusion latents, without RGB bounds | + +With `return_dict=False`, the pipeline returns `(frames,)` with the same selected representation. Select +`result.frames[0]` to export the first decoded video with `export_to_video()`. + +[[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..8564bf33e377 --- /dev/null +++ b/docs/source/en/api/schedulers/edenlive.md @@ -0,0 +1,81 @@ + + +# EdenLiveFlowMatchEulerDiscreteScheduler + +This scheduler implements the deterministic flow Euler trajectory used by EdenLive 14B bidirectional inference and +the stochastic trajectory used by the distilled causal checkpoints. The pipeline selects the appropriate +configuration when loading a profile. The 5B bidirectional profile uses [`UniPCMultistepScheduler`] instead. + +## Released schedules + +| Profile | Scheduler mode | Default evaluations | Model timesteps | +|---|---|---|---| +| `14b-bidir` | Deterministic Euler, `shift=3.0` | 30 per video | Generated from the requested step count | +| `5b-causal` | `stochastic_sampling=True` | 4 per chunk | `[1000, 937.5, 833.3333333333, 625]` | +| `14b-causal` | `stochastic_sampling=True` | 4 per chunk | `[1000, 900, 750, 500]` | + +The causal schedules are stored as `default_timesteps` and already contain the released time shift. Passing a +different `num_inference_steps` without replacing the schedule raises an error. Arbitrary step counts are not a +supported inference setting for these distilled checkpoints. + +## Set a schedule + +Call `set_timesteps()` before stepping through a video or a new causal chunk. It resets both the step index and the +begin index. Choose one schedule source: + +- `timesteps`: explicit model times in [0, `num_train_timesteps`], converted to noise levels by dividing by + `num_train_timesteps` (1000 by default). +- `sigmas`: explicit noise levels in [0, 1]. Do not supply both `timesteps` and `sigmas`. +- Neither argument: use `default_timesteps` if configured; otherwise generate a schedule from a positive + `num_inference_steps` and apply `shift`. + +Explicit values must be finite and strictly decreasing. They are not shifted again. An optional terminal zero is +removed from the model evaluation points, and the internal `sigmas` tensor always receives one terminal zero. +If `num_inference_steps` is provided with an explicit or default schedule, it must match the number of positive +evaluation points. `timesteps` and `sigmas` are stored in FP32 on the requested `device`. + +For example, the 14B causal schedule can be constructed independently for a custom inference loop: + +```py +from diffusers import EdenLiveFlowMatchEulerDiscreteScheduler + +scheduler = EdenLiveFlowMatchEulerDiscreteScheduler( + stochastic_sampling=True, + default_timesteps=[1000, 900, 750, 500], +) +scheduler.set_timesteps(device="cuda") +``` + +For a loaded pipeline, use its existing `pipe.scheduler` to preserve the profile configuration. + +## Advance the sample + +Pass the model's flow velocity, the current value from `scheduler.timesteps`, and the noisy latent sample to +`step(model_output, timestep, sample, ...)`. Deterministic mode takes an Euler step. Stochastic mode first predicts +a clean endpoint and then adds fresh noise at the next noise level. The terminal step returns the clean endpoint +without drawing noise. + +Use `generator` to control stochastic draws, or supply `variance_noise` for a nonterminal stochastic step. The two +arguments are mutually exclusive. Supplied noise must have the same shape and dtype as the sample and be on its +device. This allows a distributed caller to generate noise once and share it across ranks. Noise is unused in +deterministic mode and at the terminal stochastic step. + +`step()` returns `FlowMatchEulerDiscreteSchedulerOutput.prev_sample`, or `(prev_sample,)` when +`return_dict=False`. The result preserves the input sample's dtype. `scale_model_input()` returns its input +unchanged. Once the schedule is exhausted, call `set_timesteps()` before processing the next chunk. + +## API + +[[autodoc]] EdenLiveFlowMatchEulerDiscreteScheduler + - set_timesteps + - step diff --git a/examples/edenlive/README.md b/examples/edenlive/README.md new file mode 100644 index 000000000000..62dd10b83a75 --- /dev/null +++ b/examples/edenlive/README.md @@ -0,0 +1,113 @@ +# EdenLive inference + +`inference_edenlive.py` runs offline video generation from a reference image, text, and a proxy video. It supports +the four released profiles through `EdenLivePipeline`, with single-GPU execution or Ulysses context parallelism +and PyTorch FSDP2 across two or four GPUs. See the [pipeline guide](../../docs/source/en/api/pipelines/edenlive.md) +for the Python API, component loading, decoder overrides, and output formats. + +## Setup + +Run these commands from the Diffusers repository root in an environment with a CUDA-compatible PyTorch build: + +```bash +pip install -e . +pip install --upgrade -r examples/edenlive/requirements.txt +``` + +Transformers must provide `DINOv3ViTModel`, `UMT5EncoderModel`, and `CLIPVisionModel`. Multi-GPU execution also +requires PyTorch's composable FSDP2 API (`torch.distributed.fsdp.fully_shard`). The release directory or Hub repository +must contain EdenLive's root `model_index.json` and the original weights for the selected profile. + +## Single GPU + +```bash +python examples/edenlive/inference_edenlive.py \ + --model VAST-AI/EdenLive --profile 5b-causal \ + --video proxy.mp4 --image reference.png \ + --prompt "A character explores a snowy landscape." \ + --num-frames 49 --output edenlive.mp4 +``` + +`--model` also accepts a local release or saved Diffusers pipeline. Add `--local-files-only` to use only local or +cached weights. The CLI requires `--profile`; for a saved pipeline it must match the saved configuration. + +| Profile | Default height x width | Default fps | Default steps | Guidance scale | +|---|---|---|---|---| +| `5b-bidir` | 704 x 1280 | 24 | 48 | 3 | +| `14b-bidir` | 480 x 832 | 16 | 30 | 3 | +| `5b-causal` | 704 x 1280 | 24 | 4 per chunk | 1 | +| `14b-causal` | 480 x 832 | 16 | 4 per chunk | 1 | + +The CLI resamples source timestamps to `--fps`, center-crops to the output aspect ratio, and resizes the frames. +`--num-frames` must be positive, have the form $4k+1$, and fit the clip at the selected frame rate. If omitted, the +largest valid prefix is used. Output height and width must be multiples of 32 for 5B or 16 for 14B. Unlike this CLI, +the Python pipeline expects frames that are already sampled at the desired output frame rate. + +The released causal schedules contain four fixed evaluation points. Keep `--num-inference-steps` at four or omit it +for causal profiles. The bidirectional profiles accept a different positive step count. `--guidance-scale` above +one enables the negative prompt branch; supply its text with `--negative-prompt`. + +## Multiple GPUs + +Launch one process per GPU, with two or four processes: + +```bash +torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ + --model VAST-AI/EdenLive --profile 14b-causal \ + --video proxy.mp4 --image reference.png \ + --prompt "A character explores a snowy landscape." \ + --num-frames 81 --output edenlive.mp4 +``` + +Set `--nproc_per_node=4` for four GPUs. The transformer uses Ulysses context parallelism and FSDP2 over the same +ranks; other components are replicated. Add `--no-fsdp` to retain context parallelism with replicated transformer +weights. Only rank zero writes the video and JSON report. When comparing seeded results, keep the parallel degree +and attention backend fixed because BF16 matrix multiplication shapes can affect predictions. + +## Temporal prompts + +For `5b-bidir`, pass `--text-timeline timeline.json`. This example describes 49 frames at 24 fps: + +```json +{ + "fps": 24, + "num_frames": 49, + "segments": [ + {"start_sec": 0.0, "end_sec": 1.0, "text": "A character walks through snow."}, + {"start_sec": 1.0, "end_sec": 2.0, "text": "The character approaches a warm campfire."} + ] +} +``` + +Match the JSON `fps` and `num_frames` to `--fps` and `--num-frames`. Segments start at zero, are contiguous, and +cover the last frame timestamp; each segment must include at least one latent frame. A latent frame at an exact +segment boundary uses the segment on its left. Other profiles do not accept timelines. + +For either causal profile, pass `--chunk-prompts prompts.json` containing a JSON list of complete prompts. A 49-frame +video has 13 latent frames, so the released three-frame chunks require 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 final chunk may be shorter. The CLI still requires `--prompt` when either temporal format is supplied; the +timeline or chunk prompts determine the positive captions used for those segments or chunks. Each invocation +starts a fresh causal history. + +## Decoding and reports + +LightTAE is the default decoder for original releases. Add `--full-vae` to use the full Wan VAE for output decoding. +The VAE always encodes the reference image. `--vae-tiling` enables tiling on that VAE; it does not enable LightTAE +tiling. Use `--seed` to select the noise seed and `--attention-backend` to select an available Diffusers attention +backend. Run `python examples/edenlive/inference_edenlive.py --help` for all options. + +The output path selects the video filename. A JSON file with the same stem records the CLI arguments, resampled +source frame indices, output fps and frame count, world size, FSDP2 status, and per-rank elapsed seconds and peak +allocated/reserved GPU memory in bytes. Timing covers the pipeline call, including conditioning and decoding; +it excludes checkpoint loading, input video preprocessing, and video export. diff --git a/examples/edenlive/inference_edenlive.py b/examples/edenlive/inference_edenlive.py new file mode 100644 index 000000000000..c549ffb56640 --- /dev/null +++ b/examples/edenlive/inference_edenlive.py @@ -0,0 +1,220 @@ +# Copyright 2026 The HuggingFace 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, optionally using Ulysses and FSDP2 through torchrun.""" + +import argparse +import json +import math +import os +import time +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +from PIL import Image + +from diffusers import ContextParallelConfig, 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 configure_parallelism(transformer, device, world_size, use_fsdp): + if world_size == 1: + transformer.to(device) + return + if world_size not in (2, 4): + raise ValueError("This example supports 1, 2, or 4 processes.") + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + + mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) + transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) + if use_fsdp: + policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=torch.bfloat16, + cast_forward_inputs=False, + ) + for block in transformer.blocks: + fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + fully_shard(transformer, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + + def native_timestep(module, args, kwargs): + # Match native FSDP time quantization without copying mutable cache dataclasses. + if "timestep" in kwargs: + kwargs = {**kwargs, "timestep": kwargs["timestep"].float().to(torch.bfloat16)} + elif len(args) > 1: + args = (args[0], args[1].float().to(torch.bfloat16), *args[2:]) + return args, kwargs + + transformer.register_forward_pre_hook(native_timestep, with_kwargs=True, prepend=True) + else: + transformer.to(device) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="VAST-AI/EdenLive", help="EdenLive release directory or Hub repository.") + parser.add_argument("--profile", choices=("5b-bidir", "14b-bidir", "5b-causal", "14b-causal"), required=True) + parser.add_argument("--full-vae", action="store_true", help="Decode with the full Wan VAE instead of LightTAE.") + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--image", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--text-timeline", type=Path, help="JSON timeline for 5B bidirectional inference.") + 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, default=0) + parser.add_argument("--attention-backend") + parser.add_argument("--no-fsdp", action="store_true", help="Use replicated DiT weights with context parallelism.") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--vae-tiling", action="store_true") + return parser.parse_args() + + +def main(): + args = parse_args() + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if world_size not in (1, 2, 4): + raise ValueError("Launch this example with 1, 2, or 4 processes.") + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if world_size > 1: + dist.init_process_group("nccl", device_id=device) + rank = dist.get_rank() if dist.is_initialized() else 0 + try: + loading_kwargs = {"decoder": None} if args.full_vae else {} + pipeline = EdenLivePipeline.from_pretrained( + args.model, + profile=args.profile, + local_files_only=args.local_files_only, + **loading_kwargs, + ) + if args.attention_backend is not None: + pipeline.transformer.set_attention_backend(args.attention_backend) + configure_parallelism(pipeline.transformer, device, world_size, not args.no_fsdp) + for name, component in pipeline.components.items(): + if name != "transformer" and isinstance(component, torch.nn.Module): + component.to(device=device) + if args.vae_tiling: + pipeline.vae.enable_tiling() + pipeline.set_progress_bar_config(disable=rank != 0) + is_5b = args.profile.startswith("5b-") + 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) + 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) + for name in ("text_timeline", "chunk_prompts"): + path = getattr(args, name) + if path is not None: + inputs[name] = json.loads(path.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 = { + "rank": rank, + "seconds": time.perf_counter() - started, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), + } + all_metrics = [None] * world_size + if world_size > 1: + dist.all_gather_object(all_metrics, metrics) + else: + all_metrics[0] = metrics + if rank == 0: + 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, + "world_size": world_size, + "fsdp2": world_size > 1 and not args.no_fsdp, + "metrics": all_metrics, + } + args.output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n") + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/edenlive/requirements.txt b/examples/edenlive/requirements.txt new file mode 100644 index 000000000000..9935e2adae5f --- /dev/null +++ b/examples/edenlive/requirements.txt @@ -0,0 +1,9 @@ +torch +torchvision +transformers +accelerate +av +imageio +imageio-ffmpeg +sentencepiece +ftfy diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 36cd6d283884..896c324230ac 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -286,6 +286,10 @@ "DreamLiteTransformer2DModel", "DreamLiteUNetModel", "EasyAnimateTransformer3DModel", + "EdenLiveCausalState", + "EdenLiveLightTAEDecoder", + "EdenLiveProxyProjection", + "EdenLiveTransformer3DModel", "ErnieImageTransformer2DModel", "Flux2Transformer2DModel", "FluxControlNetModel", @@ -449,6 +453,7 @@ "DPMSolverMultistepInverseScheduler", "DPMSolverMultistepScheduler", "DPMSolverSinglestepScheduler", + "EdenLiveFlowMatchEulerDiscreteScheduler", "EDMDPMSolverMultistepScheduler", "EDMEulerScheduler", "EntropyBoundScheduler", @@ -659,6 +664,8 @@ "EasyAnimateControlPipeline", "EasyAnimateInpaintPipeline", "EasyAnimatePipeline", + "EdenLivePipeline", + "EdenLivePipelineOutput", "ErnieImagePipeline", "Flux2KleinInpaintPipeline", "Flux2KleinKVPipeline", @@ -1184,6 +1191,10 @@ DreamLiteTransformer2DModel, DreamLiteUNetModel, EasyAnimateTransformer3DModel, + EdenLiveCausalState, + EdenLiveLightTAEDecoder, + EdenLiveProxyProjection, + EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, FluxControlNetModel, @@ -1343,6 +1354,7 @@ DPMSolverMultistepInverseScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, + EdenLiveFlowMatchEulerDiscreteScheduler, EDMDPMSolverMultistepScheduler, EDMEulerScheduler, EntropyBoundScheduler, @@ -1532,6 +1544,8 @@ EasyAnimateControlPipeline, EasyAnimateInpaintPipeline, EasyAnimatePipeline, + EdenLivePipeline, + EdenLivePipelineOutput, ErnieImagePipeline, Flux2KleinInpaintPipeline, Flux2KleinKVPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index cbe283b73954..b89ad4f8aaad 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,11 @@ _import_structure["transformers.transformer_cosmos"] = ["CosmosTransformer3DModel"] _import_structure["transformers.transformer_cosmos3"] = ["Cosmos3OmniTransformer"] _import_structure["transformers.transformer_easyanimate"] = ["EasyAnimateTransformer3DModel"] + _import_structure["transformers.transformer_edenlive"] = [ + "EdenLiveCausalState", + "EdenLiveProxyProjection", + "EdenLiveTransformer3DModel", + ] _import_structure["transformers.transformer_ernie_image"] = ["ErnieImageTransformer2DModel"] _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"] _import_structure["transformers.transformer_flux2"] = ["Flux2Transformer2DModel"] @@ -215,6 +221,7 @@ AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, + EdenLiveLightTAEDecoder, Kandinsky6SRVAE, LongCatAudioDiTVae, LTX2VideoDiffusionDecoderModel, @@ -268,6 +275,9 @@ DreamLiteTransformer2DModel, DualTransformer2DModel, EasyAnimateTransformer3DModel, + EdenLiveCausalState, + EdenLiveProxyProjection, + EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, FluxTransformer2DModel, 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..14fad3d18ed8 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_edenlive.py @@ -0,0 +1,182 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES and The HuggingFace 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 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 + + +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.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[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. + """ + + @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 _convert_original_state_dict(self, state_dict): + converted = {} + for key, value in state_dict.items(): + if key.startswith("encoder."): + continue + key = key.replace("decoder.blocks.", "decoder.", 1) + converted[key] = value + for index, layer in enumerate(self.decoder): + if isinstance(layer, EdenLiveTAETemporalUpsample): + key = f"decoder.{index}.conv.weight" + expected = layer.conv.weight.shape[0] + if key in converted and converted[key].shape[0] > expected: + # Original LightTAE checkpoints retain a stride-two tensor at the stride-one first stage. + converted[key] = converted[key][-expected:] + return converted + + def _decode_chunk(self, latents, state): + first_chunk = not state + 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): + """Decode raw VAE latents using fresh temporal state for each call. + + 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 only within this call, 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.") + state = {} + 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 079c4a37a236..fb7990cee394 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, EdenLiveProxyProjection, 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..c05482a2b987 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_edenlive.py @@ -0,0 +1,1040 @@ +# Copyright 2026 The EdenLive Authors and The HuggingFace 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 maybe_adjust_dtype_for_device, maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput +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: + """Per-rollout, per-guidance-branch KV storage. Positions count latent frames.""" + + layers: list[_EdenLiveLayerCache] + batch_size: int + height: int + width: int + chunk_size: int + max_context_frames: int + sink_frames: int + tokens_per_frame: int + cp_size: int = 1 + cp_rank: int = 0 + 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 + _parallel_config = None + + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + rotary_emb=None, + *, + text_timeline=None, + query_start=0, + 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.") + if text_timeline is not None: + raise ValueError("Temporal text conditioning is supported by the 5B model only.") + 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)) + if rotary_emb is not None: + query = _apply_rotary_emb(query, rotary_emb) + key = _apply_rotary_emb(key, rotary_emb) + + 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, + parallel_config=None, + ) + .flatten(2, 3) + .type_as(query) + ) + + if text_timeline is None: + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config if encoder_hidden_states is None else None, + ) + else: + if encoder_hidden_states is None or attention_mask is not None: + raise ValueError("Temporal text requires cross-attention without an additional attention mask.") + lengths = text_timeline["lengths"] + text_length = int(text_timeline["text_length"]) + query_end = query_start + query.shape[1] + samples = [] + for sample_index, segments in enumerate(lengths): + pieces = [] + segment_start = 0 + for segment_index, segment_length in enumerate(segments): + segment_end = segment_start + segment_length + start, end = max(query_start, segment_start), min(query_end, segment_end) + if start < end: + context_start = segment_index * text_length + pieces.append( + dispatch_attention_fn( + query[sample_index : sample_index + 1, start - query_start : end - query_start], + key[sample_index : sample_index + 1, context_start : context_start + text_length], + value[sample_index : sample_index + 1, context_start : context_start + text_length], + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=None, + ) + ) + segment_start = segment_end + samples.append(torch.cat(pieces, dim=1)) + output = torch.cat(samples, dim=0) + + 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(EdenLiveAttnProcessor): + 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 layer_cache is None: + return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, rotary_emb) + if encoder_hidden_states is not None or attention_mask is not None: + raise ValueError("Causal KV storage is only used by 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)) + query = _apply_rotary_emb(query, rotary_emb) + 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) + # Ulysses retains the native rank-major KV order, including across chunks. + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + output = output.flatten(2, 3).type_as(query) + return attn.to_out[1](attn.to_out[0](output)) + + +class EdenLiveProxyProjection(ModelMixin, ConfigMixin): + r""" + Learned DINO patch projection, separate from the Transformers vision backbone. + + Args: + in_channels (`int`, defaults to `1024`): + Number of input DINO feature channels. + out_channels (`int`, defaults to `1024`): + Number of projected proxy feature channels. + """ + + @register_to_config + def __init__(self, in_channels: int = 1024, out_channels: int = 1024): + super().__init__() + self.patch_projection = nn.Linear(in_channels, out_channels) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.Tensor`): + DINO patch features with `in_channels` in the final dimension. + + Returns: + `torch.Tensor`: + Unit-normalized projected features in float32, with `out_channels` in the final dimension. + """ + hidden_states = self.patch_projection(hidden_states) + return F.normalize(hidden_states, p=2, dim=-1).float() + + +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, + is_cross_attention=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) + + if is_cross_attention is not None: + self.is_cross_attention = is_cross_attention + else: + 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 + + # Copied from diffusers.models.transformers.transformer_wan.WanAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **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)) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_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, cross_attn_norm=True, 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) if cross_attn_norm else nn.Identity() + 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_timeline=None, + query_start=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_timeline=text_timeline, + query_start=query_start, + 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. + cross_attn_norm (`bool`, defaults to `True`): + Whether to normalize the hidden states before cross-attention. + 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. + 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"] + _repeated_blocks = ["EdenLiveTransformerBlock"] + _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] + _skip_keys = ["causal_state"] + _cp_plan = { + "rope": { + 0: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + "condition_embedder": { + 1: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), + }, + "blocks.0": {"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3)}, + "head_input": ContextParallelOutput(gather_dim=1, expected_dims=3), + } + + @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, + cross_attn_norm: bool = True, + 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, + 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 proxy_channels < 1 or num_layers < 1: + raise ValueError("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.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, cross_attn_norm, eps, added_kv_proj_dim + ) + for _ in range(num_layers) + ] + ) + self.head_input = nn.Identity() + 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 _context_parallel_info(self): + config = self._parallel_config + if config is None or config.context_parallel_config is None: + return 1, 0 + config = config.context_parallel_config + if config.ring_degree != 1 or config.ulysses_anything or config.ring_anything: + raise ValueError( + "EdenLive supports equal-sized Ulysses context parallelism without Ring or Anything mode." + ) + size = config.ulysses_degree + if self.config.num_attention_heads % size: + raise ValueError("num_attention_heads must be divisible by the Ulysses degree.") + return size, config._ulysses_local_rank + + 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. + + 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; defaults to the model device. + dtype (`torch.dtype`, *optional*): Cache dtype; defaults to the model dtype. + + Returns: + `EdenLiveCausalState`: Empty KV storage for the current context-parallel layout. + """ + _, 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.") + cp_size, cp_rank = self._context_parallel_info() + tokens_per_frame = (height // patch_height) * (width // patch_width) + if tokens_per_frame % cp_size: + raise ValueError("Latent spatial token count must be divisible by the Ulysses degree.") + shape = ( + batch_size, + max_context_frames * tokens_per_frame // cp_size, + self.config.num_attention_heads, + self.config.attention_head_dim, + ) + layers = [ + _EdenLiveLayerCache( + torch.zeros(shape, device=device or self.device, dtype=dtype or self.dtype), + torch.zeros(shape, device=device or self.device, dtype=dtype or self.dtype), + ) + for _ in self.blocks + ] + return EdenLiveCausalState( + layers, + batch_size, + height, + width, + chunk_size, + max_context_frames, + sink_frames, + tokens_per_frame, + cp_size, + cp_rank, + ) + + @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, + text_timeline: dict[str, Any] | 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 proxy features of shape `(batch_size, proxy_channels, num_frames, proxy_height, proxy_width)`. + Spatial dimensions are interpolated to match the latent input. + text_timeline (`dict[str, Any]`, *optional*): + Bidirectional 5B prompt layout with `lengths`, a sequence of positive segment token counts per sample, + and `text_length`, the token count of each prompt in the concatenated text embedding bank. + 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_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.") + + cp_size, cp_rank = self._context_parallel_info() + tokens_per_frame = (height // patch_height) * (width // patch_width) + num_tokens = num_frames * tokens_per_frame + if num_tokens % cp_size: + raise ValueError("The input token count must be divisible by the Ulysses degree.") + query_start = cp_rank * (num_tokens // cp_size) + text_context_length = encoder_hidden_states.shape[1] + if text_timeline is not None: + if causal_state is not None or self.config.image_dim is not None: + raise ValueError("Temporal prompt banks are supported by bidirectional 5B inference only.") + lengths = text_timeline["lengths"] + text_length = int(text_timeline["text_length"]) + if len(lengths) != batch_size or text_length < 1: + raise ValueError("A timeline must contain one nonempty prompt bank per sample.") + if any(not sample or min(sample) < 1 or sum(sample) != num_tokens for sample in lengths): + raise ValueError("Each timeline must partition all query tokens into positive-length segments.") + if text_context_length != max(map(len, lengths)) * text_length: + raise ValueError("The text embedding bank does not match the timeline layout.") + + 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.") + if (batch_size, height, width, cp_size, cp_rank) != ( + state.batch_size, + state.height, + state.width, + state.cp_size, + state.cp_rank, + ) or len(state.layers) != len(self.blocks): + raise ValueError("The causal state does not match this model's batch, latent dimensions or CP layout.") + 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 // cp_size + cache_sink_tokens = state.sink_frames * tokens_per_frame // cp_size + key_rotary_emb = self.cache_rope(hidden_states, num_frames=visible_frames) + if cp_size > 1: + # Each rank caches its contiguous slice of every chunk. Keep the same + # chunk-local order for RoPE; Ulysses exchanges the sharded KV tensors. + chunk_tokens = state.chunk_size * tokens_per_frame + key_rotary_emb = tuple( + torch.cat( + [ + chunk.narrow(1, cp_rank * (chunk.shape[1] // cp_size), chunk.shape[1] // cp_size) + for chunk in frequency.split(chunk_tokens, dim=1) + ], + dim=1, + ) + for frequency in key_rotary_emb + ) + + rotary_emb = self.rope(hidden_states, start_frame=query_start_frame) + 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_timeline": text_timeline, + "query_start": query_start, + "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) + + hidden_states = self.head_input(hidden_states) + 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 = ( + self.norm_out(hidden_states.float()) * (1 + scale.to(hidden_states.device)) + + shift.to(hidden_states.device) + ).type_as(hidden_states) + 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..90f792b5f946 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"] = ["EdenLivePipeline", "EdenLivePipelineOutput"] _import_structure["easyanimate"] = [ "EasyAnimatePipeline", "EasyAnimateInpaintPipeline", @@ -724,6 +725,7 @@ EasyAnimateInpaintPipeline, EasyAnimatePipeline, ) + from .edenlive import 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..13360a497a1d --- /dev/null +++ b/src/diffusers/pipelines/edenlive/__init__.py @@ -0,0 +1,41 @@ +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_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_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/loading_utils.py b/src/diffusers/pipelines/edenlive/loading_utils.py new file mode 100644 index 000000000000..e61b4772cc85 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/loading_utils.py @@ -0,0 +1,523 @@ +# Copyright 2026 The HuggingFace 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. + +"""Read the four original EdenLive releases without writing converted checkpoints.""" + +import copy +import json +from collections import Counter +from contextlib import nullcontext +from pathlib import Path, PurePosixPath + +import torch +from huggingface_hub import hf_hub_download, snapshot_download +from transformers import ( + AutoTokenizer, + CLIPImageProcessor, + CLIPVisionModel, + DINOv3ViTConfig, + DINOv3ViTModel, + UMT5EncoderModel, +) + +from ...loaders.single_file_utils import convert_wan_vae_to_diffusers +from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel +from ...models.model_loading_utils import load_state_dict +from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler +from ...utils import _resolve_dtype +from ...utils.hub_utils import _resolve_revision + + +_EDENLIVE_PROFILES = { + "5b-bidir": {"size": "5b", "causal": False, "transformer": "dit-5b-bidir", "resize": "bilinear"}, + "14b-bidir": {"size": "14b", "causal": False, "transformer": "dit-14b-bidir", "resize": "bicubic"}, + "5b-causal": {"size": "5b", "causal": True, "transformer": "dit-5b", "resize": "bilinear"}, + "14b-causal": {"size": "14b", "causal": True, "transformer": "dit-14b", "resize": "bicubic_center_crop"}, +} + +_COMPONENT_CLASSES = { + "transformer": ("diffusers", "EdenLiveTransformer3DModel"), + "vae": ("diffusers", "AutoencoderKLWan"), + "decoder": ("diffusers", "EdenLiveLightTAEDecoder"), + "text_encoder": ("transformers", "UMT5EncoderModel"), + "tokenizer": ("transformers", "AutoTokenizer"), + "proxy_encoder": ("transformers", "DINOv3ViTModel"), + "proxy_projection": ("diffusers", "EdenLiveProxyProjection"), + "image_encoder": ("transformers", "CLIPVisionModel"), + "image_processor": ("transformers", "CLIPImageProcessor"), +} + +_HUB_OPTIONS = {"cache_dir", "force_download", "proxies", "local_files_only", "token", "revision"} + +_UNSUPPORTED_LOADING_OPTIONS = { + "device_map", + "max_memory", + "offload_folder", + "offload_state_dict", + "variant", + "quantization_config", + "custom_pipeline", + "custom_revision", + "use_onnx", + "use_flashpack", + "load_connected_pipeline", + "provider", + "sess_options", + "provider_options", +} + + +def read_checkpoint_config(source, loading_options): + options = loading_options.copy() + if Path(source).is_dir(): + config_file = Path(source) / "model_index.json" + else: + options["revision"] = _resolve_revision( + source, + **{ + key: value + for key, value in options.items() + if key in {"revision", "cache_dir", "local_files_only", "token"} + }, + ) + config_file = Path( + hf_hub_download( + str(source), + "model_index.json", + **{key: value for key, value in options.items() if key in _HUB_OPTIONS}, + ) + ) + with config_file.open() as handle: + config = json.load(handle) + if not isinstance(config, dict): + raise ValueError("EdenLive model_index.json must contain an object.") + return config, config_file, options + + +def normalize_saved_profile(config, requested_profile): + profile = config.get("profile") + if profile in ("5b", "14b"): + modes = {"EdenLivePipeline": "bidir", "EdenLiveChunkCausalPipeline": "causal"} + mode = modes.get(config.get("_class_name")) + if mode is None: + raise ValueError("Cannot infer an EdenLive profile from this legacy pipeline class.") + profile = f"{profile}-{mode}" + if profile not in _EDENLIVE_PROFILES: + raise ValueError(f"Saved EdenLive pipeline has an invalid profile: {profile!r}.") + if requested_profile is not None and requested_profile != profile: + raise ValueError(f"Requested profile {requested_profile!r} conflicts with saved profile {profile!r}.") + return profile + + +def load_original_pipeline(pipeline_class, source, config_file, config, profile, loading_options): + component_names, option_names = pipeline_class._get_signature_keys(pipeline_class) + passed_components = {name: loading_options[name] for name in component_names if name in loading_options} + context = resolve_original_checkpoint( + config, profile=profile, passed_components=passed_components, loading_options=loading_options + ) + supported_options = ( + _HUB_OPTIONS + | _UNSUPPORTED_LOADING_OPTIONS + | { + "dtype", + "torch_dtype", + "low_cpu_mem_usage", + "disable_mmap", + "use_safetensors", + "trust_remote_code", + } + ) + unknown = set(loading_options) - set(component_names) - set(option_names) - supported_options + if unknown: + raise ValueError(f"Original EdenLive loading does not support {', '.join(sorted(unknown))}.") + dtype = _resolve_dtype(loading_options.get("dtype"), loading_options.get("torch_dtype")) + if dtype is None: + dtype = context["default_dtype"] + if not isinstance(dtype, (dict, torch.dtype)) or ( + isinstance(dtype, dict) and any(not isinstance(value, torch.dtype) for value in dtype.values()) + ): + raise ValueError("`dtype` must be a torch.dtype or a mapping of component names to torch.dtype values.") + folder = config_file.parent + if not Path(source).is_dir() and ( + loading_options.get("force_download", False) + or not all((folder / filename).is_file() for filename in context["files"]) + ): + folder = Path( + snapshot_download( + str(source), + allow_patterns=context["files"], + user_agent={"pipeline_class": pipeline_class.__name__}, + **{key: value for key, value in loading_options.items() if key in _HUB_OPTIONS}, + ) + ) + missing = [filename for filename in context["files"] if not (folder / filename).is_file()] + if missing: + raise FileNotFoundError(f"Original checkpoint files are missing from {folder}: {missing}") + init_kwargs = {name: context["config"][name] for name in option_names if name in context["config"]} + init_kwargs.update({name: loading_options[name] for name in option_names if name in loading_options}) + try: + for name in component_names: + if name in passed_components: + init_kwargs[name] = passed_components[name] + elif context["config"][name][0] is None: + init_kwargs[name] = None + else: + component_dtype = ( + dtype.get(name, dtype.get("default", torch.float32)) if isinstance(dtype, dict) else dtype + ) + init_kwargs[name] = load_original_component( + name, cached_folder=folder, context=context, dtype=component_dtype, loading_options=loading_options + ) + pipeline = pipeline_class(**init_kwargs) + finally: + context["state_dicts"].clear() + pipeline.register_to_config(_name_or_path=str(source)) + return pipeline + + +def _relative_path(value): + if not isinstance(value, str) or not value or "\\" in value or any(c in value for c in "*?["): + raise ValueError("EdenLive path overrides must be literal repository-relative paths.") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or not path.parts: + raise ValueError(f"EdenLive path must stay within the checkpoint repository: {value!r}.") + return path.as_posix() + + +def _component_sources(profile_name, path_overrides): + profile = _EDENLIVE_PROFILES[profile_name] + size, transformer = profile["size"], profile["transformer"] + sources = { + "transformer": { + "config": f"{transformer}/config.json", + "weights": f"{transformer}/diffusion_pytorch_model.safetensors", + }, + "vae": { + "config": f"vae-{size}/config.json", + "weights": f"vae-{size}/Wan2.{'2' if size == '5b' else '1'}_VAE.pth", + }, + "decoder": {"weights": f"decoder-{size}/lighttaew2_{'2' if size == '5b' else '1'}.pth"}, + "text_encoder": {"config": "umt5/text_encoder/config.json", "weights": "umt5/text_encoder/model.safetensors"}, + "tokenizer": {"directory": "umt5/tokenizer"}, + "proxy_encoder": { + "config": "vision-encoder/backbone/config.json", + "weights": "vision-encoder/projection.safetensors", + }, + "proxy_projection": {"weights": "vision-encoder/projection.safetensors"}, + "image_encoder": { + "config": "clip/image_encoder/config.json", + "weights": "clip/image_encoder/model.safetensors", + }, + "image_processor": {"directory": "clip/image_processor"}, + } + if not isinstance(path_overrides, dict) or set(path_overrides) - set(sources): + raise ValueError("`_path_overrides` must map EdenLive component names to paths.") + for name, override in path_overrides.items(): + if not isinstance(override, dict): + raise ValueError(f"Path overrides for {name} must be an object.") + # Only transformers vary by full profile; VAE and decoder files vary by size. + selectors = ( + set(_EDENLIVE_PROFILES) + if name == "transformer" + else {"5b", "14b"} + if name in ("vae", "decoder") + else set() + ) + if selectors and set(override) <= selectors: + selector = profile_name if name == "transformer" else size + for value in override.values(): + if not isinstance(value, dict) or set(value) - (set(sources[name]) | {"directory"}): + raise ValueError(f"Invalid path fields for {name}.") + for path in value.values(): + _relative_path(path) + override = override.get(selector, {}) + if set(override) - (set(sources[name]) | {"directory"}): + raise ValueError(f"Invalid path fields for {name}: {sorted(override)}.") + override = {key: _relative_path(value) for key, value in override.items()} + if "directory" in override and "directory" not in sources[name]: + directory = override.pop("directory") + sources[name] = {key: f"{directory}/{PurePosixPath(value).name}" for key, value in sources[name].items()} + sources[name].update(override) + if size == "5b": + sources.pop("image_encoder") + sources.pop("image_processor") + return sources + + +def _source_files(name, source): + if name == "tokenizer": + return [ + f"{source['directory']}/{filename}" + for filename in ("tokenizer_config.json", "special_tokens_map.json", "tokenizer.json", "spiece.model") + ] + if name == "image_processor": + return [f"{source['directory']}/preprocessor_config.json"] + return list(source.values()) + + +def _read_config(folder, source): + with open(Path(folder) / source["config"]) as handle: + config = json.load(handle) + if not isinstance(config, dict): + raise ValueError(f"Expected an object in {source['config']}.") + return config + + +def _read_component_state(folder, name, context, disable_mmap): + filename = context["sources"][name]["weights"] + state = context["state_dicts"].get(filename) + if state is None: + state = load_state_dict(str(Path(folder) / filename), disable_mmap=disable_mmap, map_location="cpu") + if not isinstance(state, dict) or any( + not isinstance(k, str) or not isinstance(v, torch.Tensor) for k, v in state.items() + ): + raise ValueError(f"Expected a tensor state dictionary in {filename}.") + context["source_users"][filename] -= 1 + if context["source_users"][filename] > 0: + context["state_dicts"][filename] = state + else: + context["state_dicts"].pop(filename, None) + return state + + +def _strict_assign(model, state, dtype): + model.load_state_dict(state, strict=True, assign=True) + return model.to(dtype=dtype).eval().requires_grad_(False) + + +def _empty_weights(loading_options): + return torch.device("meta") if loading_options.get("low_cpu_mem_usage", True) else nullcontext() + + +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 _load_transformers_component(model_class, config, state, dtype, loading_options, allowed_unused=()): + state = {key: value for key, value in state.items() if key not in allowed_unused} + model, info = model_class.from_pretrained( + None, + config=model_class.config_class.from_dict(config), + state_dict=state, + dtype=dtype, + local_files_only=True, + low_cpu_mem_usage=loading_options.get("low_cpu_mem_usage", True), + output_loading_info=True, + ) + if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): + raise ValueError(f"{model_class.__name__} checkpoint did not load strictly: {info}.") + return model.eval().requires_grad_(False) + + +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 resolve_original_checkpoint(config_dict, *, profile, passed_components, loading_options): + if "_edenlive_original" not in config_dict: + return None + original = config_dict["_edenlive_original"] + if not isinstance(original, dict) or set(original) - {"_path_overrides"}: + raise ValueError("`_edenlive_original` must be an object with optional `_path_overrides`.") + if profile not in _EDENLIVE_PROFILES: + raise ValueError( + f"Original EdenLive checkpoints require `profile` in {tuple(_EDENLIVE_PROFILES)}; got {profile!r}." + ) + preset = _EDENLIVE_PROFILES[profile] + unsupported = { + key: loading_options.get(key) + for key in _UNSUPPORTED_LOADING_OPTIONS + if loading_options.get(key) is not None and loading_options.get(key) is not False + } + if unsupported: + raise ValueError( + f"Original EdenLive loading does not support {', '.join(sorted(unsupported))}; load first and place the pipeline explicitly." + ) + sources = _component_sources(profile, original.get("_path_overrides", {})) + sources = {name: source for name, source in sources.items() if name not in passed_components} + safe_only = loading_options.get("use_safetensors") + weights = [source["weights"] for source in sources.values() if "weights" in source] + if safe_only is True and any(not path.endswith(".safetensors") for path in weights): + raise ValueError( + "Original EdenLive includes PyTorch .pth components; omit `use_safetensors=True` or provide those components explicitly." + ) + if safe_only is False and any(path.endswith(".safetensors") for path in weights): + raise ValueError( + "Original EdenLive includes Safetensors components; omit `use_safetensors=False` or provide those components explicitly." + ) + config = copy.deepcopy(config_dict) + config.pop("_edenlive_original") + config.update( + _class_name="EdenLivePipeline", + profile=profile, + proxy_resize_mode=preset["resize"], + proxy_unconditional="zero" if profile == "14b-causal" else "same", + proxy_projection_autocast=True, + proxy_image_size=(480, 832), + ) + config.update(_COMPONENT_CLASSES) + if preset["size"] == "5b": + config.update(image_encoder=(None, None), image_processor=(None, None)) + scheduler_name = "UniPCMultistepScheduler" if profile == "5b-bidir" else "EdenLiveFlowMatchEulerDiscreteScheduler" + config["scheduler"] = ("diffusers", scheduler_name) + if preset["causal"]: + config.update(chunk_size=3, max_context_frames=18, sink_frames=6) + return { + "config": config, + "profile": profile, + "preset": preset, + "sources": sources, + "files": sorted( + {"model_index.json"} + | {filename for name, source in sources.items() for filename in _source_files(name, source)} + ), + "default_dtype": { + "default": torch.bfloat16, + "proxy_encoder": torch.float32, + "proxy_projection": torch.float32, + }, + "state_dicts": {}, + "source_users": Counter(source["weights"] for source in sources.values() if "weights" in source), + } + + +def load_original_component(name, *, cached_folder, context, dtype, loading_options): + preset = context["preset"] + if name == "scheduler": + if context["profile"] == "5b-bidir": + return UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5.0) + if preset["causal"]: + timesteps = [1000, 937.5, 833.3333333333, 625] if preset["size"] == "5b" else [1000, 900, 750, 500] + return EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True, default_timesteps=timesteps) + return EdenLiveFlowMatchEulerDiscreteScheduler(shift=3.0) + source = context["sources"][name] + if name in ("tokenizer", "image_processor"): + model_class = AutoTokenizer if name == "tokenizer" else CLIPImageProcessor + return model_class.from_pretrained(Path(cached_folder) / source["directory"], local_files_only=True) + state = _read_component_state(cached_folder, name, context, loading_options.get("disable_mmap", False)) + config = _read_config(cached_folder, source) if "config" in source else None + if name == "transformer": + proxy_channels = state["patch_embedding.weight"].shape[1] - config["in_channels"] + expected_proxy_channels = config.get("proxy_channels", 1024) + condition_rmsnorm = "patch_embedding.condition_norm.weight" in state + if proxy_channels != expected_proxy_channels or condition_rmsnorm != config.get( + "condition_rmsnorm", preset["size"] == "14b" + ): + raise ValueError( + "Transformer proxy channels or normalization do not match its configuration and selected profile." + ) + with _empty_weights(loading_options): + model = EdenLiveTransformer3DModel.from_config( + config, proxy_channels=proxy_channels, condition_rmsnorm=condition_rmsnorm + ) + # These rotary buffers are derived, not serialized checkpoint tensors. + 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)) + return _strict_assign(model, state, dtype) + if name == "vae": + mapped = {} + for key, value in state.items(): + target = _remap_wan22_vae_key(key) + if target in mapped: + raise ValueError(f"Multiple VAE keys map to {target}.") + mapped[target] = value + mapped = convert_wan_vae_to_diffusers(mapped) + if len(mapped) != len(state): + raise ValueError("Multiple VAE checkpoint keys map to the same parameter.") + # Wan 2.2 stores folded RGB convolutions; the current native class + # takes their post-patchify/pre-unpatchify channel counts. + patch_size = config.get("patch_size") or 1 + for field, key, dimension in ( + ("in_channels", "encoder.conv_in.weight", 1), + ("out_channels", "decoder.conv_out.weight", 0), + ): + channels = mapped[key].shape[dimension] + declared = config.get(field, 3) + if channels not in (declared, declared * patch_size**2): + raise ValueError(f"VAE {field} does not match the checkpoint convolutions.") + config[field] = channels + with _empty_weights(loading_options): + model = AutoencoderKLWan.from_config(config) + return _strict_assign(model, mapped, dtype) + if name == "decoder": + with _empty_weights(loading_options): + model = EdenLiveLightTAEDecoder( + latent_channels=48 if preset["size"] == "5b" else 16, patch_size=2 if preset["size"] == "5b" else 1 + ) + return _strict_assign(model, model._convert_original_state_dict(state), dtype) + if name in ("proxy_encoder", "proxy_projection"): + state = _proxy_state(state) + if name == "proxy_encoder": + config = DINOv3ViTConfig.from_dict(config) + with _empty_weights(loading_options): + model = DINOv3ViTModel(config) + expected = set(model.state_dict()) + mapped = {} + for key, value in state.items(): + if key.startswith("backbone."): + target = key.removeprefix("backbone.") + if target not in expected and f"model.{target}" in expected: + target = f"model.{target}" + mapped[target] = value + model.rope_embeddings = type(model.rope_embeddings)(config) + else: + mapped = {key: value for key, value in state.items() if key.startswith("patch_projection.")} + weight = mapped["patch_projection.weight"] + with _empty_weights(loading_options): + model = EdenLiveProxyProjection(in_channels=weight.shape[1], out_channels=weight.shape[0]) + return _strict_assign(model, mapped, dtype) + if name == "text_encoder": + return _load_transformers_component(UMT5EncoderModel, config, state, dtype, loading_options) + if name == "image_encoder": + return _load_transformers_component( + CLIPVisionModel, config, state, dtype, loading_options, ("visual_projection.weight",) + ) + raise ValueError(f"Unknown EdenLive component: {name}.") diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py new file mode 100644 index 000000000000..e758d24e9b9a --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -0,0 +1,740 @@ +# Copyright 2026 The HuggingFace 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 typing import Callable + +import torch +from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel + +from ...models import AutoencoderKLWan, EdenLiveProxyProjection, EdenLiveTransformer3DModel +from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler +from ..pipeline_utils import DiffusionPipeline +from . import pipeline_utils as utils +from .loading_utils import load_original_pipeline, normalize_saved_profile, read_checkpoint_config +from .pipeline_output import EdenLivePipelineOutput + + +class EdenLivePipeline(DiffusionPipeline): + """Bidirectional and causal video generation conditioned on a reference image, text, and an RGB proxy video. + + Load released weights with [`EdenLivePipeline.from_pretrained`]. The constructor accepts initialized components and + does not read checkpoints. The full `profile` selects the model size and sampling path. + + Args: + vae ([`AutoencoderKLWan`]): + Wan VAE for reference-image encoding and, when `decoder=None`, video decoding. + text_encoder (`UMT5EncoderModel`): + UMT5 text encoder. + tokenizer (`AutoTokenizer`): + Tokenizer corresponding to the text encoder. + transformer ([`EdenLiveTransformer3DModel`]): + Video denoiser with the weights and configuration for the selected profile. + scheduler ([`EdenLiveFlowMatchEulerDiscreteScheduler`] or [`UniPCMultistepScheduler`]): + Sampling schedule compatible with the transformer. The 5B bidirectional release uses UniPC; the other + profiles use the EdenLive Euler scheduler, with fixed stochastic schedules for causal checkpoints. + proxy_encoder (`DINOv3ViTModel`): + Trained DINOv3 encoder for the RGB proxy video. + proxy_projection ([`EdenLiveProxyProjection`]): + Projection from DINO features to transformer conditioning features. + image_encoder (`CLIPVisionModel`, *optional*): + CLIP reference-image encoder required by 14B profiles. + image_processor (`CLIPImageProcessor`, *optional*): + Image processor corresponding to the CLIP encoder. + profile (`str`, defaults to `"5b-bidir"`): + One of `"5b-bidir"`, `"14b-bidir"`, `"5b-causal"`, or `"14b-causal"`. + proxy_resize_mode (`str`, defaults to `"bilinear"`): + RGB proxy preprocessing before DINO encoding: `"bilinear"`, `"bicubic"`, or `"bicubic_center_crop"`. + proxy_unconditional (`str`, defaults to `"same"`): + Use the same proxy features (`"same"`) or zeros (`"zero"`) in the negative guidance branch. + proxy_projection_autocast (`bool`, defaults to `True`): + Allow autocast during proxy projection while retaining FP32 projection parameters. + proxy_image_size (`tuple[int, int]`, defaults to `(480, 832)`): + Height and width used to preprocess RGB frames for the proxy encoder. + chunk_size (`int`, defaults to 3): + Number of latent frames per causal chunk. A final chunk may be shorter. + max_context_frames (`int`, defaults to 18): + Maximum causal attention context in latent frames, including sink frames and the current chunk. + sink_frames (`int`, defaults to 6): + Initial latent frames retained as attention sinks when the context window advances. + decoder (`torch.nn.Module`, *optional*): + Optional video decoder compatible with the VAE latent channels and temporal/spatial scales. `None` uses the + VAE. Original-release loading supplies LightTAE by default; pass `decoder=None` when loading to select the + full VAE. Use a compatible `ModelMixin` decoder for pipeline device placement and serialization. + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->proxy_encoder->proxy_projection->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, + proxy_projection: EdenLiveProxyProjection, + image_encoder: CLIPVisionModel | None = None, + image_processor: CLIPImageProcessor | None = None, + profile: str = "5b-bidir", + proxy_resize_mode: str = "bilinear", + proxy_unconditional: str = "same", + proxy_projection_autocast: bool = True, + proxy_image_size: tuple | list = (480, 832), + chunk_size: int = 3, + max_context_frames: int = 18, + sink_frames: int = 6, + decoder: torch.nn.Module | None = None, + ): + super().__init__() + 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.") + utils.initialize_pipeline( + self, + vae, + text_encoder, + tokenizer, + transformer, + scheduler, + proxy_encoder, + proxy_projection, + image_encoder, + image_processor, + decoder, + profile, + proxy_resize_mode, + proxy_unconditional, + proxy_projection_autocast, + proxy_image_size, + ) + + self.register_to_config(chunk_size=chunk_size, max_context_frames=max_context_frames, sink_frames=sink_frames) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *, profile=None, **kwargs): + """Load an original EdenLive release or a standard saved pipeline. + + Original releases are identified by the EdenLive metadata in their root `model_index.json`. Only files needed + by the selected profile and components are loaded; weights are mapped in memory without conversion copies. + Standard snapshots use the inherited Diffusers component loader after normalizing legacy profile metadata. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Hub repository ID or local directory containing an original release or saved Diffusers pipeline. + profile (`str`, *optional*): + Required for original releases: `"5b-bidir"`, `"14b-bidir"`, `"5b-causal"`, or `"14b-causal"`. Standard + snapshots retain their saved profile; a conflicting explicit value raises an error. + dtype (`torch.dtype` or `dict[str, torch.dtype]`, *optional*): + Loading dtype, or a component-name mapping with an optional `"default"` entry. Original releases + default to BF16 for the transformer, VAE, text encoder, CLIP, and LightTAE, and FP32 for DINO and its + projection. `torch_dtype` is also accepted. Supplied component instances retain their own dtype. + decoder (`torch.nn.Module`, *optional*): + Original releases use LightTAE when omitted. Explicit `None` selects the full VAE. A compatible + initialized decoder replaces the saved or released decoder without changing its identity or dtype. + Standard snapshots use their saved decoder when this argument is omitted. + revision (`str`, *optional*): + Hub branch, tag, or commit to resolve before selecting component files. + cache_dir (`str` or `os.PathLike`, *optional*): + Hub cache directory. + local_files_only (`bool`, defaults to `False`): + Load local files or cached Hub files without downloading. + token (`str` or `bool`, *optional*): + Authentication token or Hub token-selection option. + kwargs (`dict`, *optional*): + Initialized component overrides, such as `transformer`, `vae`, or `text_encoder`, and loading options + supported by the checkpoint format. Original loading does not support active `device_map`, `variant`, + or quantization options. Load on CPU and move the pipeline to the desired device afterward. + + Returns: + [`EdenLivePipeline`]: + Pipeline with a canonical full profile and the selected components. Older standard snapshots with + size-only profiles or the former causal class name are normalized in memory. + """ + config, config_file, options = read_checkpoint_config(pretrained_model_name_or_path, kwargs) + if "_edenlive_original" in config: + return load_original_pipeline(cls, pretrained_model_name_or_path, config_file, config, profile, options) + profile = normalize_saved_profile(config, profile) + decoder = options.get("decoder") + if decoder is not None: + options["decoder"] = None + pipeline = super().from_pretrained(pretrained_model_name_or_path, profile=profile, **options) + if decoder is not None: + utils.validate_decoder(decoder, pipeline.vae) + pipeline.register_modules(decoder=decoder) + return pipeline + + 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 prompts with the native UMT5 encoder. + + Args: + prompt (`str` or `list[str]`, *optional*): + Positive text, unless `prompt_embeds` is supplied. + negative_prompt (`str` or `list[str]`, *optional*): + Negative text for classifier-free guidance; defaults to empty strings. + do_classifier_free_guidance (`bool`, defaults to `True`): + Whether to produce negative embeddings. + prompt_embeds (`torch.Tensor`, *optional*): + Precomputed positive embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Precomputed negative embeddings. + max_sequence_length (`int`, defaults to 512): + Fixed token count per prompt. + device (`torch.device`, *optional*): + Device for encoding and returned embeddings. + dtype (`torch.dtype`, *optional*): + Output dtype; defaults to the transformer dtype, or the text encoder dtype when used alone. + + Returns: + `tuple[torch.Tensor, torch.Tensor | None]`: + Positive and negative text embeddings, with padding tokens zeroed. + """ + prompt_embeds, negative_prompt_embeds = utils.encode_prompt( + self, + prompt, + negative_prompt, + do_classifier_free_guidance, + prompt_embeds, + negative_prompt_embeds, + max_sequence_length, + device, + dtype, + ) + return prompt_embeds, negative_prompt_embeds + + def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): + """Prepare full-video initial noise before encoding the image condition. + + Args: + batch_size (`int`): + Number of generated videos. + height (`int`): + Decoded frame height. + width (`int`): + Decoded frame width. + num_frames (`int`): + Decoded frame count of the form `4k + 1`. + dtype (`torch.dtype`): + Latent dtype. + device (`torch.device`): + Latent device. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator, or one per video. + latents (`torch.Tensor`, *optional*): + Existing initial noise to validate and move to the requested device and dtype. + + Returns: + `torch.Tensor`: + Noise with shape `(B, C, T_latent, H_latent, W_latent)`. + """ + return utils.prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator, latents) + + @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() + 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, + text_timeline: dict | list[dict] | 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] = ["latents"], + ) -> EdenLivePipelineOutput | tuple: + """Generate a video from a reference image and a full-rate RGB proxy video. + + Args: + prompt (`str` or `list[str]`, *optional*): + Text describing the output video. Supply this or `prompt_embeds` unless using temporal 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; defaults to 48 for 5B bidirectional, 30 for 14B bidirectional, or 4 per causal chunk. + Released causal schedules contain four fixed evaluation points and reject a different step count. + guidance_scale (`float`, *optional*): + Classifier-free guidance strength; defaults to 3 for bidirectional or 1 for causal generation. 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. Under context parallelism, the + first rank generates and broadcasts each complete noise tensor. + 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`. + text_timeline (`dict` or `list[dict]`, *optional*): + 5B bidirectional temporal captions, one dictionary per video, with `fps`, `num_frames`, and contiguous + `segments`. Each segment has `start_sec`, `end_sec`, and `text`. Exact boundaries belong to the segment + on their left. This option cannot be combined with precomputed embeddings. + chunk_prompts (`list[str]`, *optional*): + Causal-only prompts, one per latent chunk. Requires batch size one and no precomputed positive + embeddings. Each call starts with fresh attention caches. + 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*): + Keyword arguments passed to the transformer's attention processors. + 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`. + + Returns: + [`EdenLivePipelineOutput`] or `tuple`: + Generated frames, or latent videos when `output_type="latent"`. + """ + if text_timeline is not None and (self.is_causal or self.model_size != "5b"): + raise ValueError("`text_timeline` requires the 5B bidirectional profile.") + if chunk_prompts is not None and not self.is_causal: + raise ValueError("`chunk_prompts` requires a causal profile.") + if guidance_scale is None: + guidance_scale = 1.0 if self.is_causal else 3.0 + if num_inference_steps is None: + num_inference_steps = 4 if self.is_causal else (48 if self.model_size == "5b" else 30) + height = (704 if self.model_size == "5b" else 480) if height is None else height + width = (1280 if self.model_size == "5b" else 832) if width is None else width + utils.check_inputs( + self, + 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 + ) + if text_timeline is not None: + if self.model_size != "5b" or prompt_embeds is not None or negative_prompt_embeds is not None: + raise ValueError( + "`text_timeline` requires the 5B profile and text prompts instead of precomputed embeddings." + ) + batch_size = 1 if isinstance(text_timeline, dict) else len(text_timeline) + else: + batch_size = ( + prompt_embeds.shape[0] + if prompt_embeds is not None + else ( + 1 + if isinstance(prompt, str) or (prompt is None and chunk_prompts is not None) + else len(prompt or []) + ) + ) + if batch_size < 1: + raise ValueError("Provide `prompt`, `prompt_embeds`, `text_timeline`, or `chunk_prompts`.") + latents = self.prepare_latents(batch_size, height, width, num_frames, dtype, device, generator, latents) + chunk_embeds = None + num_chunks = (latents.shape[2] + self.config.chunk_size - 1) // self.config.chunk_size + if chunk_prompts is not None: + 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(chunk_prompt, str) for chunk_prompt in chunk_prompts) + ): + raise ValueError( + "`chunk_prompts` requires batch size one, no prompt embeddings, and one string per chunk." + ) + unique_embeds = {} + for chunk_prompt in chunk_prompts: + if chunk_prompt not in unique_embeds: + unique_embeds[chunk_prompt], _ = self.encode_prompt( + chunk_prompt, do_classifier_free_guidance=False, max_sequence_length=max_sequence_length + ) + chunk_embeds = torch.cat([unique_embeds[chunk_prompt] for chunk_prompt in chunk_prompts]) + prompt = chunk_prompts[0] + prompt_embeds = chunk_embeds[:1] + if text_timeline is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + negative_prompt, + self.do_classifier_free_guidance, + prompt_embeds, + negative_prompt_embeds, + max_sequence_length, + ) + layout = None + else: + prompt_embeds, negative_prompt_embeds, layout = utils.encode_timeline( + self, + text_timeline, + negative_prompt, + num_frames, + latents.shape, + max_sequence_length, + self.do_classifier_free_guidance, + ) + condition, first_frame, image_embeds, proxy_features = utils.prepare_conditions( + self, image, proxy_video, height, width, num_frames, batch_size, device, dtype + ) + if self.is_causal: + return self._denoise_causal( + latents, + condition, + first_frame, + proxy_features, + image_embeds, + prompt_embeds, + negative_prompt_embeds, + chunk_embeds, + num_chunks, + batch_size, + device, + dtype, + num_inference_steps, + generator, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + output_type, + return_dict, + ) + return self._denoise_bidir( + latents, + condition, + first_frame, + proxy_features, + image_embeds, + prompt_embeds, + negative_prompt_embeds, + layout, + device, + num_inference_steps, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + output_type, + return_dict, + ) + + def _denoise_bidir( + self, + latents, + condition, + first_frame, + proxy_features, + image_embeds, + prompt_embeds, + negative_prompt_embeds, + layout, + device, + num_inference_steps, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + output_type, + return_dict, + ): + if self.model_size == "5b": + latents[:, :, :1] = first_frame + self.scheduler.set_timesteps(num_inference_steps, device=device) + if self.model_size == "5b" 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 = utils.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, + text_timeline=layout, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + if self.do_classifier_free_guidance: + negative_proxy = ( + proxy_features + if self.config.proxy_unconditional == "same" + else torch.zeros_like(proxy_features) + ) + negative_prediction = self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=negative_proxy, + text_timeline=layout, + 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 self.model_size == "5b": + 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 = utils.decode_latents(self, latents, output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (frames,) + return EdenLivePipelineOutput(frames=frames) + + def _denoise_chunk( + self, + latents, + condition, + first_frame, + proxy_features, + image_embeds, + prompt_embeds, + negative_prompt_embeds, + states, + num_inference_steps, + generator, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + callback_step_offset=0, + ): + start = states[0].next_frame + if start == 0: + latents[:, :, :1] = first_frame + self.scheduler.set_timesteps(num_inference_steps, device=latents.device) + for step_index, timestep in enumerate(self.scheduler.timesteps): + model_input, model_timestep = utils.prepare_model_input( + self, latents, condition, first_frame, timestep, start, causal=True + ) + prediction = self.transformer( + model_input, + model_timestep, + prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=proxy_features, + causal_state=states[0], + cache_update=False, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + if self.do_classifier_free_guidance: + negative_proxy = ( + proxy_features if self.config.proxy_unconditional == "same" else torch.zeros_like(proxy_features) + ) + negative_prediction = self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=negative_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 = utils.sample_noise(self, latents.shape, generator, latents.device, latents.dtype) + latents = self.scheduler.step( + prediction, + timestep, + latents, + generator=None if variance_noise is not None else generator, + variance_noise=variance_noise, + return_dict=False, + )[0] + if start == 0: + 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, callback_step_offset + 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) + model_input, model_timestep = utils.prepare_model_input( + self, latents, condition, first_frame, latents.new_zeros(()), start, causal=True + ) + if self.model_size == "14b": + # A clean cache commit uses one zero timestep per sample, not per token. + model_timestep = torch.zeros(latents.shape[0], device=latents.device, dtype=torch.float32) + self.transformer( + model_input, + model_timestep, + prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=proxy_features, + causal_state=states[0], + cache_update=True, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + if self.do_classifier_free_guidance: + negative_proxy = ( + proxy_features if self.config.proxy_unconditional == "same" else torch.zeros_like(proxy_features) + ) + self.transformer( + model_input, + model_timestep, + negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + proxy_features=negative_proxy, + causal_state=states[1], + cache_update=True, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + return latents, prompt_embeds, negative_prompt_embeds + + def _denoise_causal( + self, + latents, + condition, + first_frame, + proxy_features, + image_embeds, + prompt_embeds, + negative_prompt_embeds, + chunk_embeds, + num_chunks, + batch_size, + device, + dtype, + num_inference_steps, + generator, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + output_type, + return_dict, + ): + 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 = [] + try: + 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, prompt_embeds, negative_prompt_embeds = self._denoise_chunk( + latents[:, :, start : start + self.config.chunk_size].clone(), + condition, + first_frame, + proxy_features[:, :, start : start + self.config.chunk_size], + image_embeds, + prompt_embeds, + negative_prompt_embeds, + states, + num_inference_steps, + generator, + attention_kwargs, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + chunk_index * num_inference_steps, + ) + generated.append(current) + progress_bar.update() + latents = torch.cat(generated, dim=2) + frames = utils.decode_latents(self, latents, output_type) + finally: + for state in states: + state.reset() + 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..3d5956179460 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_output.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace 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..d03db5a66407 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_utils.py @@ -0,0 +1,377 @@ +# Copyright 2026 The HuggingFace 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 bisect +import html +import math +import re + +import torch +import torch.nn.functional as F + +from ...models.autoencoders.autoencoder_edenlive import EdenLiveLightTAEDecoder +from ...utils import is_ftfy_available +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from .video_processor import EdenLiveVideoProcessor + + +if is_ftfy_available(): + import ftfy + + +def initialize_pipeline( + pipe, + vae, + text_encoder, + tokenizer, + transformer, + scheduler, + proxy_encoder, + proxy_projection, + image_encoder, + image_processor, + decoder, + profile, + proxy_resize_mode, + proxy_unconditional, + proxy_projection_autocast, + proxy_image_size, +): + if profile not in ("5b-bidir", "14b-bidir", "5b-causal", "14b-causal"): + raise ValueError("`profile` must be one of 5b-bidir, 14b-bidir, 5b-causal, or 14b-causal.") + pipe.model_size, mode = profile.split("-") + pipe.is_causal = mode == "causal" + if proxy_unconditional not in ("same", "zero"): + raise ValueError("`proxy_unconditional` must be 'same' or 'zero'.") + if decoder is not None: + validate_decoder(decoder, vae) + pipe.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + proxy_encoder=proxy_encoder, + proxy_projection=proxy_projection, + image_encoder=image_encoder, + image_processor=image_processor, + decoder=decoder, + ) + pipe.register_to_config( + profile=profile, + proxy_resize_mode=proxy_resize_mode, + proxy_unconditional=proxy_unconditional, + proxy_projection_autocast=proxy_projection_autocast, + proxy_image_size=proxy_image_size, + ) + pipe.vae_scale_factor_temporal = vae.config.scale_factor_temporal if vae is not None else 4 + pipe.vae_scale_factor_spatial = vae.config.scale_factor_spatial if vae is not None else 8 + pipe.video_processor = VideoProcessor(vae_scale_factor=pipe.vae_scale_factor_spatial) + pipe.proxy_video_processor = EdenLiveVideoProcessor(vae_scale_factor=1, do_normalize=False, do_convert_rgb=True) + + +def validate_decoder(decoder, vae): + if not isinstance(decoder, torch.nn.Module) or not callable(getattr(decoder, "decode", None)): + raise ValueError("`decoder` must be a module with a `decode(latents, return_dict=False)` method.") + config = getattr(decoder, "config", None) + latent_channels = getattr(config, "latent_channels", getattr(config, "z_dim", None)) + temporal = getattr(decoder, "scale_factor_temporal", getattr(config, "scale_factor_temporal", None)) + spatial = getattr(decoder, "scale_factor_spatial", getattr(config, "scale_factor_spatial", None)) + if vae is None or (latent_channels, temporal, spatial) != ( + vae.config.z_dim, + vae.config.scale_factor_temporal, + vae.config.scale_factor_spatial, + ): + raise ValueError("`decoder` latent channels and temporal/spatial scale factors must match the Wan VAE.") + + +def encode_prompt( + pipe, + prompt, + negative_prompt=None, + do_classifier_free_guidance=True, + prompt_embeds=None, + negative_prompt_embeds=None, + max_sequence_length=512, + device=None, + dtype=None, +): + device = device or pipe._execution_device + dtype = dtype or (pipe.transformer.dtype if pipe.transformer is not None else pipe.text_encoder.dtype) + + def encode(texts): + 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 = pipe.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 = pipe.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) + + 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 = encode(prompt) + 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 = encode(negative_prompt) + 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( + pipe, + prompt, + prompt_embeds, + negative_prompt, + negative_prompt_embeds, + image, + video, + height, + width, + output_type, + callback_inputs, +): + 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 = pipe.transformer.config.patch_size + if ( + height <= 0 + or width <= 0 + or height % (pipe.vae_scale_factor_spatial * patch[1]) + or width % (pipe.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(pipe._callback_tensor_inputs): + raise ValueError(f"Callback tensor inputs must belong to {pipe._callback_tensor_inputs}.") + + +def prepare_latents(pipe, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): + shape = ( + batch_size, + pipe.vae.config.z_dim, + (num_frames - 1) // pipe.vae_scale_factor_temporal + 1, + height // pipe.vae_scale_factor_spatial, + width // pipe.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 sample_noise(pipe, shape, generator, device, dtype) + + +def sample_noise(pipe, shape, generator, device, dtype): + parallel_config = getattr(pipe.transformer, "_parallel_config", None) + context_parallel_config = getattr(parallel_config, "context_parallel_config", None) + if context_parallel_config is None: + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + group = context_parallel_config._flattened_mesh.get_group() + source_rank = torch.distributed.get_global_rank(group, 0) + if torch.distributed.get_rank() == source_rank: + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + noise = torch.empty(shape, device=device, dtype=dtype) + torch.distributed.broadcast(noise, src=source_rank, group=group) + return noise + + +def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batch_size, device, dtype): + pixels = pipe.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 or proxy_video.shape[0] not in (1, batch_size): + raise ValueError("Image and proxy-video batch sizes must be one or match the prompt batch.") + clip_embeds = None + if pipe.model_size == "14b": + if pipe.image_encoder is None or pipe.image_processor is None: + raise ValueError("The 14B profile 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 = pipe.image_processor(images=list(rgb), return_tensors="pt").pixel_values + clip_embeds = ( + pipe.image_encoder( + clip_inputs.to(device=device, dtype=pipe.image_encoder.dtype), output_hidden_states=True + ) + .hidden_states[-2] + .to(dtype) + ) + 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 = pipe.vae.encode(condition_video.to(pipe.vae.dtype)).latent_dist.mode() + mean = condition.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / condition.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) + condition = ((condition - mean) * inverse_std).to(dtype) + first_frame = condition[:, :, :1] + if pipe.model_size == "14b": + mask = condition.new_zeros(batch_size, pipe.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 = pipe.config.proxy_image_size + proxy_features = [] + for frame in proxy_video.unbind(2): + frame = frame.to(device=device, dtype=torch.float32) + if pipe.model_size == "5b": + # Match the original FP32 [-1, 1] -> [0, 1] round trip before DINO. + frame = (frame * 2 - 1 + 1) * 0.5 + inputs = pipe.proxy_video_processor.preprocess_proxy_frames( + frame, proxy_height, proxy_width, pipe.config.proxy_resize_mode + ) + with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + hidden = pipe.proxy_encoder(pixel_values=inputs).last_hidden_state + patch_size = pipe.proxy_encoder.config.patch_size + grid_height, grid_width = proxy_height // patch_size, proxy_width // patch_size + start = 1 + pipe.proxy_encoder.config.num_register_tokens + hidden = hidden[:, start : start + grid_height * grid_width].reshape(batch_size, grid_height, grid_width, -1) + with torch.autocast( + device.type, dtype=torch.bfloat16, enabled=device.type == "cuda" and pipe.config.proxy_projection_autocast + ): + projected = pipe.proxy_projection(hidden if pipe.config.proxy_projection_autocast else hidden.float()) + proxy_features.append(projected.float().permute(0, 3, 1, 2)) + return condition, first_frame, clip_embeds, torch.stack(proxy_features, dim=2) + + +def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0, causal=False): + model_input = latents + pin_first = start == 0 and (causal or pipe.model_size == "5b") + if pin_first: + model_input = model_input.clone() + model_input[:, :, :1] = first_frame + if pipe.model_size == "14b": + 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 + + +def decode_latents(pipe, latents, output_type): + if output_type == "latent": + return latents + latents = latents.to(pipe.vae.dtype) + mean = latents.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) + std = latents.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) + if isinstance(pipe.decoder, EdenLiveLightTAEDecoder): + raw_latents = latents * std + mean + else: + raw_latents = latents / (1 / std) + mean + decoder = pipe.decoder if pipe.decoder is not None else pipe.vae + dtype = getattr(decoder, "dtype", None) + if dtype is None: + parameter = next(decoder.parameters(), None) + dtype = parameter.dtype if parameter is not None else raw_latents.dtype + video = decoder.decode(raw_latents.to(dtype=dtype), return_dict=False)[0] + return pipe.video_processor.postprocess_video(video, output_type=output_type) + + +def encode_timeline(pipe, timelines, negative_prompt, num_frames, latent_shape, max_sequence_length, guidance): + timelines = [timelines] if isinstance(timelines, dict) else timelines + batch_size = len(timelines) + negatives = ( + [negative_prompt or ""] * batch_size if isinstance(negative_prompt, (str, type(None))) else negative_prompt + ) + if len(negatives) != batch_size: + raise ValueError("Negative prompt batch must match the timeline batch.") + texts, negative_texts, lengths, counts = [], [], [], [] + tokens_per_frame = (latent_shape[-2] // pipe.transformer.config.patch_size[1]) * ( + latent_shape[-1] // pipe.transformer.config.patch_size[2] + ) + for sample, negative in zip(timelines, negatives): + fps = float(sample["fps"]) + if not math.isfinite(fps) or fps <= 0 or sample["num_frames"] != num_frames: + raise ValueError("Each timeline needs positive finite fps and the requested decoded num_frames.") + duration, previous = (num_frames - 1) / fps, 0.0 + segments = [] + for segment in sample["segments"]: + start, end = float(segment["start_sec"]), float(segment["end_sec"]) + if ( + not math.isfinite(start) + or not math.isfinite(end) + or end <= start + or not math.isclose(start, previous, abs_tol=1e-6) + ): + raise ValueError("Timeline segments must be finite, ordered, contiguous, and start at zero.") + previous = end + if start < duration: + segments.append({"text": segment["text"], "end_sec": min(end, duration)}) + if not segments or previous < duration: + raise ValueError("Timeline segments must cover the requested video.") + frame_counts = [0] * len(segments) + for frame in range(latent_shape[2]): + index = bisect.bisect_left( + [segment["end_sec"] for segment in segments], frame * pipe.vae_scale_factor_temporal / fps - 1e-7 + ) + frame_counts[index] += 1 + if 0 in frame_counts: + raise ValueError("A timeline segment is shorter than the latent time resolution.") + lengths.append(tuple(count * tokens_per_frame for count in frame_counts)) + counts.append(len(segments)) + texts.extend(segment["text"] for segment in segments) + negative_texts.extend([negative] * len(segments)) + positive, negative = encode_prompt(pipe, texts, negative_texts, guidance, max_sequence_length=max_sequence_length) + + def pack(embeddings): + if embeddings is None: + return None + result = embeddings.new_zeros(batch_size, max(counts) * max_sequence_length, embeddings.shape[-1]) + offset = 0 + for index, count in enumerate(counts): + result[index, : count * max_sequence_length] = embeddings[offset : offset + count].flatten(0, 1) + offset += count + return result + + return pack(positive), pack(negative), {"lengths": tuple(lengths), "text_length": max_sequence_length} diff --git a/src/diffusers/pipelines/edenlive/video_processor.py b/src/diffusers/pipelines/edenlive/video_processor.py new file mode 100644 index 000000000000..f5901f4dc0e6 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/video_processor.py @@ -0,0 +1,79 @@ +# Copyright 2026 The HuggingFace 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 in ("bicubic", "bicubic_center_crop"): + 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) + if resize_mode == "bicubic": + image = image.resize((width, height), PIL.Image.Resampling.BICUBIC) + else: + ratio = image.width / image.height + resized_height, resized_width = ( + (height, round(height * ratio)) if ratio > width / height else (round(width / ratio), width) + ) + image = image.resize((resized_width, resized_height), PIL.Image.Resampling.BICUBIC) + left, top = (resized_width - width) // 2, (resized_height - height) // 2 + image = image.crop((left, top, left + width, top + height)) + 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..4d16e4079818 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_edenlive.py @@ -0,0 +1,211 @@ +# Copyright 2026 The HuggingFace 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 9e1f5f027ce0..7ec38ee07677 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1324,6 +1324,66 @@ 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 EdenLiveProxyProjection(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"] @@ -3501,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 a068ef8672e8..3bf80360258f 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1937,6 +1937,36 @@ 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/pipelines/edenlive/__init__.py b/tests/pipelines/edenlive/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/edenlive/test_edenlive.py b/tests/pipelines/edenlive/test_edenlive.py new file mode 100644 index 000000000000..0f8b128c3a9f --- /dev/null +++ b/tests/pipelines/edenlive/test_edenlive.py @@ -0,0 +1,1499 @@ +# Copyright 2026 The HuggingFace 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 hashlib +import importlib.util +import json +import os +import socket +import tempfile +import unittest +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from PIL import Image +from safetensors.torch import load_file, save_file +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import ( + CLIPImageProcessor, + CLIPVisionConfig, + CLIPVisionModel, + DINOv3ViTConfig, + DINOv3ViTModel, + PreTrainedTokenizerFast, + UMT5Config, + UMT5EncoderModel, +) + +from diffusers import ( + AutoencoderKLWan, + ContextParallelConfig, + EdenLiveFlowMatchEulerDiscreteScheduler, + EdenLiveLightTAEDecoder, + EdenLivePipeline, + EdenLiveProxyProjection, + EdenLiveTransformer3DModel, + UniPCMultistepScheduler, +) +from diffusers.models.transformers import transformer_edenlive +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig + + +@pytest.fixture(autouse=True, scope="class") +def full_precision_matmul(): + precision = torch.get_float32_matmul_precision() + matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + cudnn_tf32 = torch.backends.cudnn.allow_tf32 + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.set_float32_matmul_precision(precision) + torch.backends.cuda.matmul.allow_tf32 = matmul_tf32 + torch.backends.cudnn.allow_tf32 = cudnn_tf32 + + +def make_tokenizer(): + words = "a character walks through snow blurry racing car on sunny road fire ocean one too many".split() + vocabulary = {word: index for index, word in enumerate(["", "", "", *words])} + backend = Tokenizer(models.WordLevel(vocabulary, unk_token="")) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + return PreTrainedTokenizerFast( + tokenizer_object=backend, pad_token="", eos_token="", unk_token="", model_max_length=512 + ) + + +def file_digests(directory): + return { + str(path.relative_to(directory)): hashlib.sha256(path.read_bytes()).hexdigest() + for path in directory.rglob("*") + if path.is_file() + } + + +@pytest.fixture(params=["5b-bidir", "14b-bidir", "5b-causal", "14b-causal"]) +def original_checkpoint(tmp_path, request, monkeypatch): + import httpx + import requests + + from diffusers.pipelines.edenlive import loading_utils + + def unexpected_request(*args, **kwargs): + pytest.fail("A local original checkpoint must not request any external Hub configuration.") + + monkeypatch.setattr(httpx.Client, "send", unexpected_request) + monkeypatch.setattr(requests.Session, "send", unexpected_request) + + def tiny_decoder(**_kwargs): + # Production LightTAE dimensions stay fixed by the release profile. Only + # its construction is reduced for this weight-loading regression fixture. + return EdenLiveLightTAEDecoder(latent_channels=4, patch_size=1, block_out_channels=(4, 4, 4, 4)) + + monkeypatch.setattr(loading_utils, "EdenLiveLightTAEDecoder", tiny_decoder) + profile = request.param + size, mode = profile.split("-") + fixture = EdenLivePipelineTesterConfig() + fixture.profile = size + components = fixture.get_dummy_components() + source = tmp_path / "original" + source.mkdir() + (source / "model_index.json").write_text( + json.dumps({"_class_name": "EdenLivePipeline", "_edenlive_original": {"_path_overrides": {}}}) + ) + transformer_dir = source / (f"dit-{size}-bidir" if mode == "bidir" else f"dit-{size}") + transformer_dir.mkdir() + components["transformer"].save_config(transformer_dir) + save_file(components["transformer"].state_dict(), transformer_dir / "diffusion_pytorch_model.safetensors") + vae_dir = source / f"vae-{size}" + vae_dir.mkdir() + components["vae"].save_config(vae_dir) + # Include original Wan parameter names and the original .pth container. + vae_state = {} + replacements = { + "encoder.conv_in.": "encoder.conv1.", + "decoder.conv_out.": "decoder.head.2.", + "quant_conv.": "conv1.", + "post_quant_conv.": "conv2.", + } + for key, value in components["vae"].state_dict().items(): + for prefix, original_prefix in replacements.items(): + if key.startswith(prefix): + key = original_prefix + key[len(prefix) :] + break + vae_state[key] = value + torch.save(vae_state, vae_dir / ("Wan2.2_VAE.pth" if size == "5b" else "Wan2.1_VAE.pth")) + components["text_encoder"].save_pretrained(source / "umt5/text_encoder") + components["tokenizer"].save_pretrained(source / "umt5/tokenizer") + # This fast fixture consumes tokenizer.json. Mirror the additional vocabulary + # and special-token files listed by the original release's download plan. + (source / "umt5/tokenizer/spiece.model").write_bytes(b"unused-fast-tokenizer-fixture") + (source / "umt5/tokenizer/special_tokens_map.json").write_text( + json.dumps({"pad_token": "", "eos_token": "", "unk_token": ""}) + ) + proxy_dir = source / "vision-encoder" + components["proxy_encoder"].config.save_pretrained(proxy_dir / "backbone") + proxy_state = { + f"module.encoder.backbone.{key}": value for key, value in components["proxy_encoder"].state_dict().items() + } + proxy_state.update( + {f"module.encoder.{key}": value for key, value in components["proxy_projection"].state_dict().items()} + ) + proxy_state.update( + { + "module.encoder.global_projection.weight": torch.zeros(2, 32), + "module.encoder.global_projection.bias": torch.zeros(2), + } + ) + save_file(proxy_state, proxy_dir / "projection.safetensors") + if size == "14b": + components["image_encoder"].save_pretrained(source / "clip/image_encoder") + components["image_processor"].save_pretrained(source / "clip/image_processor") + decoder = EdenLiveLightTAEDecoder(latent_channels=4, patch_size=1, block_out_channels=(4, 4, 4, 4)) + decoder_dir = source / f"decoder-{size}" + decoder_dir.mkdir() + decoder_state = { + key.replace("decoder.", "decoder.blocks.", 1): value for key, value in decoder.state_dict().items() + } + temporal_key = "decoder.blocks.7.conv.weight" + decoder_state[temporal_key] = torch.cat( + [torch.zeros_like(decoder_state[temporal_key]), decoder_state[temporal_key]] + ) + decoder_state["encoder.unused.weight"] = torch.ones(1) + torch.save(decoder_state, decoder_dir / ("lighttaew2_2.pth" if size == "5b" else "lighttaew2_1.pth")) + return SimpleNamespace( + source=source, + profile=profile, + pipeline_class=EdenLivePipeline, + components=components, + decoder=decoder, + hashes=file_digests(source), + transformer_dir=transformer_dir, + ) + + +@pytest.fixture +def original_hub(original_checkpoint, monkeypatch): + from diffusers.pipelines.edenlive import loading_utils as pipeline_utils + + fixture = original_checkpoint + calls = SimpleNamespace(metadata=[], snapshots=[]) + revision = "0123456789abcdef0123456789abcdef01234567" + repo_id = "test-org/edenlive-original-fixture" + + def metadata_download(repository, filename, **kwargs): + assert repository == repo_id and filename == "model_index.json" + assert kwargs["revision"] == revision + calls.metadata.append(kwargs) + return str(fixture.source / filename) + + def snapshot_download(repository, **kwargs): + assert repository == repo_id + calls.snapshots.append(kwargs) + return str(fixture.source) + + monkeypatch.setattr(pipeline_utils, "_resolve_revision", lambda *args, **kwargs: revision) + monkeypatch.setattr(pipeline_utils, "hf_hub_download", metadata_download) + monkeypatch.setattr(pipeline_utils, "snapshot_download", snapshot_download) + return SimpleNamespace(fixture=fixture, calls=calls, revision=revision, repo_id=repo_id, module=pipeline_utils) + + +class EdenLivePipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = EdenLivePipeline + profile = "5b" + causal = False + required_input_params_in_call_signature = frozenset( + ["prompt", "negative_prompt", "image", "video", "height", "width", "prompt_embeds", "negative_prompt_embeds"] + ) + optional_input_params = frozenset( + ["num_frames", "num_inference_steps", "guidance_scale", "generator", "latents", "output_type", "return_dict"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + 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) + text_encoder = UMT5EncoderModel( + UMT5Config(vocab_size=1103, d_model=32, d_ff=32, d_kv=8, num_layers=1, num_heads=4, dropout_rate=0.0) + ).eval() + tokenizer = make_tokenizer() + 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) + proxy_projection = EdenLiveProxyProjection(in_channels=32, out_channels=8) + 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_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, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "scheduler": scheduler, + "proxy_encoder": proxy_encoder, + "proxy_projection": proxy_projection, + "image_encoder": image_encoder, + "image_processor": image_processor, + } + + def get_pipeline(self, **components): + causal = self.causal + pipe = self.pipeline_class( + **(components or self.get_dummy_components()), + profile=f"{self.profile}-{'causal' if causal else 'bidir'}", + proxy_image_size=(32, 32), + proxy_resize_mode=("bicubic_center_crop" if causal else "bicubic") + if self.profile == "14b" + else "bilinear", + proxy_unconditional="zero" if self.profile == "14b" and causal else "same", + ) + for component in pipe.components.values(): + if hasattr(component, "eval"): + component.eval() + pipe.set_progress_bar_config(disable=True) + return pipe + + def test_inference_and_standard_save_load(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + expected = pipe(**self.get_dummy_inputs()).frames + assert expected.shape == (1, *self.output_shape) + assert torch.isfinite(expected).all() + pipe.save_pretrained(tmp_path) + for loader in (self.pipeline_class,): + restored = loader.from_pretrained(tmp_path, local_files_only=True).to(torch_device) + restored.set_progress_bar_config(disable=True) + actual = restored(**self.get_dummy_inputs()).frames + assert type(restored) is self.pipeline_class + assert restored.config.profile == pipe.config.profile + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + 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 TestEdenLivePipeline(EdenLivePipelineTesterConfig): + @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 + + def test_timeline_boundary(self): + from diffusers.pipelines.edenlive.pipeline_utils import encode_timeline + + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["text_timeline"] = { + "num_frames": 9, + "fps": 4, + "segments": [ + {"start_sec": 0, "end_sec": 1, "text": "snow"}, + {"start_sec": 1, "end_sec": 2, "text": "fire"}, + ], + } + output = pipe(**inputs).frames + assert output.shape == (1, *self.output_shape) + _, _, layout = encode_timeline( + pipe, inputs["text_timeline"], inputs["negative_prompt"], 9, (1, 4, 3, 2, 2), 16, True + ) + assert layout["lengths"] == ((2, 1),) + + +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="14B profile requires"): + restored(**self.get_dummy_inputs()) + + +class TestEdenLive14BPipeline(EdenLivePipelineTesterConfig, EdenLive14BRequiredComponentsTesterMixin): + profile = "14b" + + def test_timeline_is_rejected(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["text_timeline"] = {} + with pytest.raises(ValueError, match="5B"): + pipe(**inputs) + + +class EdenLiveChunkCausalPipelineTesterConfig(EdenLivePipelineTesterConfig): + causal = True + + def get_dummy_components(self): + components = super().get_dummy_components() + components["scheduler"] = EdenLiveFlowMatchEulerDiscreteScheduler( + stochastic_sampling=True, default_timesteps=[1000, 500] + ) + return components + + +class TestEdenLiveChunkCausalPipeline(EdenLiveChunkCausalPipelineTesterConfig): + 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()) + + def test_chunk_prompt_changes_tail(self): + 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") + 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"] + with pytest.raises(ValueError, match="chunk_prompts"): + pipe(**inputs) + + 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) + + +class TestEdenLive14BChunkCausalPipeline( + EdenLiveChunkCausalPipelineTesterConfig, EdenLive14BRequiredComponentsTesterMixin +): + profile = "14b" + + def test_zero_proxy_changes_unconditional_prediction(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + zero_proxy = pipe(**inputs).frames + pipe.register_to_config(proxy_unconditional="same") + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + same_proxy = pipe(**inputs).frames + assert (zero_proxy - same_proxy).abs().max() > 1e-6 + + +class TestEdenLiveUnifiedInterface: + @pytest.mark.parametrize( + "profile,steps,guidance", + [("5b-bidir", 48, 3), ("14b-bidir", 30, 3), ("5b-causal", 4, 1), ("14b-causal", 4, 1)], + ) + def test_profile_defaults(self, profile, steps, guidance): + fixture = ( + EdenLiveChunkCausalPipelineTesterConfig() if profile.endswith("causal") else EdenLivePipelineTesterConfig() + ) + fixture.profile = profile.split("-")[0] + components = fixture.get_dummy_components() + if fixture.causal: + components["scheduler"].register_to_config(default_timesteps=[1000, 750, 500, 250]) + pipe = fixture.get_pipeline(**components).to(torch_device) + inputs = fixture.get_dummy_inputs() + inputs.pop("num_inference_steps") + inputs.pop("guidance_scale") + inputs.update(num_frames=1, output_type="latent") + indices = [] + + def callback(pipeline, index, timestep, tensors): + indices.append(index) + return {} + + output = pipe(**inputs, callback_on_step_end=callback).frames + assert type(pipe) is EdenLivePipeline and pipe.config.profile == profile + assert pipe.guidance_scale == guidance + assert pipe.num_timesteps == steps + assert indices == list(range(steps)) + assert output.shape == (1, 4, 1, 2, 2) + + @pytest.mark.parametrize( + "profile,argument", + [ + ("5b-bidir", "chunk_prompts"), + ("14b-bidir", "chunk_prompts"), + ("14b-bidir", "text_timeline"), + ("5b-causal", "text_timeline"), + ("14b-causal", "text_timeline"), + ], + ) + def test_incompatible_temporal_prompts_fail_before_inference(self, profile, argument): + fixture = EdenLivePipelineTesterConfig() + fixture.profile, mode = profile.split("-") + fixture.causal = mode == "causal" + pipe = fixture.get_pipeline() + with pytest.raises(ValueError, match=argument): + pipe(**{argument: []}) + + @pytest.mark.parametrize("size", ["5b", "14b"]) + def test_causal_callback_exception_cleans_history(self, size): + fixture = EdenLiveChunkCausalPipelineTesterConfig() + fixture.profile = size + pipe = fixture.get_pipeline().to(torch_device) + inputs = fixture.get_dummy_inputs() + inputs.update(num_frames=13, output_type="latent") + inputs["video"] = np.concatenate([inputs["video"], inputs["video"][:4]]) + states = [] + steps = [] + + def record_state(module, args, kwargs): + state = kwargs["causal_state"] + if all(state is not item for item in states): + states.append(state) + + def fail_on_tail(pipeline, index, timestep, tensors): + steps.append((index, tensors["latents"].shape[2])) + if index == 2: + raise RuntimeError("callback failed on tail") + return {} + + handle = pipe.transformer.register_forward_pre_hook(record_state, with_kwargs=True) + try: + with pytest.raises(RuntimeError, match="callback failed on tail"): + pipe(**inputs, callback_on_step_end=fail_on_tail) + finally: + handle.remove() + assert steps == [(0, 3), (1, 3), (2, 1)] + assert len(states) == 2 and all(state.next_frame == 0 for state in states) + assert not any("cache" in key for key in pipe.transformer.state_dict()) + + @pytest.mark.parametrize("profile", ["5b-bidir", "14b-bidir", "5b-causal", "14b-causal"]) + def test_standard_profile_selector_conflict(self, tmp_path, profile): + fixture = EdenLivePipelineTesterConfig() + fixture.profile, mode = profile.split("-") + fixture.causal = mode == "causal" + pipe = fixture.get_pipeline() + pipe.save_pretrained(tmp_path) + restored = EdenLivePipeline.from_pretrained(tmp_path, profile=profile, local_files_only=True) + assert restored.config.profile == profile + wrong_profile = f"{fixture.profile}-{'bidir' if fixture.causal else 'causal'}" + with pytest.raises(ValueError, match="conflicts with saved profile"): + EdenLivePipeline.from_pretrained(tmp_path, profile=wrong_profile, local_files_only=True) + + +class RecordingDecoder(torch.nn.Module): + scale_factor_temporal = 4 + scale_factor_spatial = 8 + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros((), dtype=torch.float64)) + self.config = SimpleNamespace(latent_channels=4) + self.inputs = [] + + def decode(self, latents, return_dict=False): + self.inputs.append(latents.detach().clone()) + batch, _, frames, height, width = latents.shape + return (latents.new_zeros((batch, 3, (frames - 1) * 4 + 1, height * 8, width * 8)),) + + +class TestEdenLiveNativeNumerics: + @pytest.mark.parametrize("profile", ["5b", "14b"]) + @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_dino_input_uses_fp32_profile_preprocessing(self, monkeypatch, profile, input_dtype): + from diffusers.pipelines.edenlive.pipeline_utils import prepare_conditions + + fixture = EdenLivePipelineTesterConfig() + fixture.profile = profile + pipe = fixture.get_pipeline().to(torch_device) + # A bilinear override must not give 14B the 5B normalization round trip. + pipe.register_to_config(proxy_resize_mode="bilinear") + rgb = torch.arange(256).reshape(1, 1, 16, 16).expand(1, 3, -1, -1) + frame = (rgb.float() / 255).to(device=torch_device, dtype=input_dtype) + seen = [] + preprocess = pipe.proxy_video_processor.preprocess_proxy_frames + + def capture(value, *args): + seen.append(value.clone()) + return preprocess(value, *args) + + monkeypatch.setattr(pipe.proxy_video_processor, "preprocess_proxy_frames", capture) + prepare_conditions( + pipe, + fixture.get_dummy_inputs()["image"], + frame.unsqueeze(2), + 16, + 16, + 1, + 1, + torch.device(torch_device), + torch.float32, + ) + assert seen[0].dtype == torch.float32 + if profile == "5b": + # Fixed IEEE-754 values for the 1/255 input after the FP32 round trip. + expected = { + torch.float32: 0.003921568393707275, + torch.float16: 0.0039215087890625, + torch.bfloat16: 0.003936767578125, + }[input_dtype] + assert seen[0][0, 0, 0, 1].item() == expected + else: + torch.testing.assert_close(seen[0], frame.float(), rtol=0, atol=0) + + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_chunk_prompts_encode_unique_texts_individually(self, guidance): + fixture = EdenLiveChunkCausalPipelineTesterConfig() + 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("profile", ["5b", "14b"]) + def test_clean_commit_timestep_shape_and_cfg(self, profile): + fixture = EdenLiveChunkCausalPipelineTesterConfig() + fixture.profile = profile + pipe = fixture.get_pipeline().to(torch_device) + 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 commits[0].shape == ((1,) if profile == "14b" else (1, 3)) + noisy = calls[0][0] + assert noisy.shape == (1, 3) and noisy[0, 0] == 0 and (noisy[0, 1:] > 0).all() + + @pytest.mark.parametrize("decoder_kind,expected", [("tiny", 4.0625), ("custom", 4.03125), ("vae", 4.03125)]) + def test_bf16_decoder_denormalization(self, monkeypatch, decoder_kind, expected): + from diffusers.pipelines.edenlive.pipeline_utils import decode_latents + + fixture = EdenLivePipelineTesterConfig() + decoder = ( + EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) + if decoder_kind == "tiny" + else RecordingDecoder() + if decoder_kind == "custom" + else None + ) + pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) + pipe.vae.to(dtype=torch.bfloat16) + pipe.vae.register_to_config(latents_mean=[0.0] * 4, latents_std=[3.0] * 4) + seen = [] + + def decode(value, **kwargs): + seen.append(value.clone()) + return (value.new_zeros(1, 3, 1, 16, 16),) + + monkeypatch.setattr(decoder if decoder is not None else pipe.vae, "decode", decode) + z = torch.full((1, 4, 1, 2, 2), 1.3515625, dtype=torch.bfloat16) + assert decode_latents(pipe, z, "latent") is z and not seen + decode_latents(pipe, z, "pt") + assert len(seen) == 1 and (seen[0] == expected).all() + + def test_example_fsdp_casts_only_time_and_preserves_cache_identity(self, monkeypatch): + from torch.distributed import device_mesh, fsdp + + path = Path(__file__).resolve().parents[3] / "examples/edenlive/inference_edenlive.py" + spec = importlib.util.spec_from_file_location("edenlive_example_test", path) + example = importlib.util.module_from_spec(spec) + spec.loader.exec_module(example) + monkeypatch.setattr(device_mesh, "init_device_mesh", lambda *args, **kwargs: {"ulysses": object()}) + policies = [] + monkeypatch.setattr(fsdp, "fully_shard", lambda module, **kwargs: policies.append(kwargs["mp_policy"])) + monkeypatch.setattr(example, "ContextParallelConfig", lambda **kwargs: kwargs) + + class Recorder(torch.nn.Module): + def __init__(self): + super().__init__() + self.blocks = torch.nn.ModuleList([torch.nn.Identity()]) + + def enable_parallelism(self, **kwargs): + pass + + def forward(self, hidden_states, timestep, *, causal_state): + causal_state.next_frame += 3 + return hidden_states, timestep, causal_state + + model = Recorder() + example.configure_parallelism(model, torch.device("cpu"), 2, True) + assert len(policies) == 2 and all(not p.cast_forward_inputs for p in policies) + state = SimpleNamespace(next_frame=0, layers=[object()]) + hidden = torch.ones(1) + for positional, time in [(True, torch.tensor([999])), (False, torch.tensor([[0.0, 375.0]]))]: + args = (hidden, time) if positional else (hidden,) + kwargs = {} if positional else {"timestep": time} + actual_hidden, actual_time, actual_state = model(*args, **kwargs, causal_state=state) + assert actual_hidden is hidden and actual_state is state + assert actual_time.dtype == torch.bfloat16 + assert actual_time.tolist() == ([1000.0] if positional else [[0.0, 376.0]]) + assert state.next_frame == 6 + + +class TestEdenLiveDecoder: + @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()) + + @pytest.mark.parametrize("causal", [False, True]) + def test_custom_decoder_receives_raw_latents_once_and_latent_output_bypasses_it(self, causal): + from diffusers.pipelines.edenlive.pipeline_utils import decode_latents + + fixture = EdenLivePipelineTesterConfig() + fixture.causal = causal + decoder = RecordingDecoder() + pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) + pipe.vae.register_to_config(latents_mean=[0.25] * 4, latents_std=[2.0] * 4) + latents = torch.full((1, 4, 4, 2, 2), 0.5) + assert decode_latents(pipe, latents, "latent") is latents + assert decoder.inputs == [] + output = decode_latents(pipe, latents, "pt") + assert pipe.decoder is decoder and decoder.weight.dtype == torch.float64 + assert decoder.inputs[0].dtype == torch.float64 + torch.testing.assert_close(decoder.inputs[0], torch.full_like(decoder.inputs[0], 1.25)) + assert output.shape == (1, 13, 3, 16, 16) + + @pytest.mark.parametrize("causal", [False, True]) + @pytest.mark.parametrize("saved_decoder", ["tiny", "none", "legacy"]) + @pytest.mark.parametrize("size", ["5b", "14b"]) + @pytest.mark.parametrize("legacy_class", [False, True]) + def test_standard_snapshot_preserves_decoder_choice(self, tmp_path, causal, saved_decoder, size, legacy_class): + fixture = EdenLivePipelineTesterConfig() + fixture.profile = size + fixture.causal = causal + decoder = ( + EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) + if saved_decoder == "tiny" + else None + ) + pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) + pipe.save_pretrained(tmp_path) + index_path = tmp_path / "model_index.json" + index = json.loads(index_path.read_text()) + assert "_edenlive_original" not in index + if legacy_class: + index["_class_name"] = "EdenLiveChunkCausalPipeline" if causal else "EdenLivePipeline" + index["profile"] = size + if saved_decoder == "legacy": + index.pop("decoder", None) + index["_class_name"] = "EdenLiveChunkCausalPipeline" if causal else "EdenLivePipeline" + index["profile"] = size if legacy_class else pipe.config.profile + index_path.write_text(json.dumps(index)) + hashes = file_digests(tmp_path) + restored = EdenLivePipeline.from_pretrained(tmp_path, local_files_only=True) + assert type(restored) is EdenLivePipeline + assert restored.config.profile == f"{size}-{'causal' if causal else 'bidir'}" + assert (restored.config.chunk_size, restored.config.max_context_frames, restored.config.sink_frames) == ( + 3, + 18, + 6, + ) + assert file_digests(tmp_path) == hashes + if legacy_class: + restored.save_pretrained(tmp_path / "new-snapshot") + updated = json.loads((tmp_path / "new-snapshot/model_index.json").read_text()) + assert updated["_class_name"] == "EdenLivePipeline" + assert updated["profile"] == restored.config.profile + if saved_decoder == "tiny": + assert isinstance(restored.decoder, EdenLiveLightTAEDecoder) + assert restored.decoder.config.latent_channels == 4 + assert tuple(restored.decoder.config.block_out_channels) == (4, 4, 4, 4) + for key, value in decoder.state_dict().items(): + torch.testing.assert_close(restored.decoder.state_dict()[key], value, atol=0, rtol=0) + else: + assert restored.decoder is None + + @pytest.mark.parametrize("causal", [False, True]) + @pytest.mark.parametrize("saved_tiny", [False, True]) + def test_standard_snapshot_accepts_plain_decoder_override(self, tmp_path, causal, saved_tiny): + fixture = EdenLivePipelineTesterConfig() + fixture.causal = causal + decoder = EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) if saved_tiny else None + pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) + pipe.save_pretrained(tmp_path) + if saved_tiny: + for path in (tmp_path / "decoder").glob("*.safetensors"): + path.unlink() + override = RecordingDecoder() + restored = EdenLivePipeline.from_pretrained( + tmp_path, decoder=override, dtype=torch.bfloat16, local_files_only=True + ) + assert restored.decoder is override and override.weight.dtype == torch.float64 + + +class TestEdenLiveOriginalLoading: + @pytest.mark.parametrize("original_checkpoint", ["5b-bidir"], indirect=True) + def test_inactive_loading_options_are_accepted(self, original_checkpoint): + pipe = EdenLivePipeline.from_pretrained( + original_checkpoint.source, + profile=original_checkpoint.profile, + local_files_only=True, + device_map=None, + variant=None, + use_onnx=False, + offload_state_dict=False, + ) + assert pipe.config.profile == original_checkpoint.profile + + def test_original_profiles_components_and_default_dtype(self, original_checkpoint): + fixture = original_checkpoint + loader = EdenLivePipeline + pipe = loader.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) + assert type(pipe) is fixture.pipeline_class + assert pipe.config.profile == fixture.profile + assert isinstance(pipe.decoder, EdenLiveLightTAEDecoder) + expected_components = {**fixture.components, "decoder": fixture.decoder} + for name, expected in expected_components.items(): + if not isinstance(expected, torch.nn.Module): + continue + actual = getattr(pipe, name) + dtype = torch.float32 if name in ("proxy_encoder", "proxy_projection") else torch.bfloat16 + assert actual.dtype == dtype + assert set(actual.state_dict()) == set(expected.state_dict()) + for key, tensor in expected.state_dict().items(): + restored = actual.state_dict()[key] + torch.testing.assert_close(restored, tensor.to(restored.dtype), atol=0, rtol=0) + if fixture.profile == "5b-bidir": + assert isinstance(pipe.scheduler, UniPCMultistepScheduler) + assert pipe.scheduler.config.flow_shift == 5 + elif fixture.profile.endswith("causal"): + points = [1000, 937.5, 833.3333333333, 625] if fixture.profile.startswith("5b") else [1000, 900, 750, 500] + assert pipe.scheduler.config.stochastic_sampling + assert list(pipe.scheduler.config.default_timesteps) == points + assert (pipe.config.chunk_size, pipe.config.max_context_frames, pipe.config.sink_frames) == (3, 18, 6) + else: + assert pipe.scheduler.config.shift == 3 and not pipe.scheduler.config.stochastic_sampling + assert pipe.config.proxy_unconditional == ("zero" if fixture.profile == "14b-causal" else "same") + expected_resize = ( + "bilinear" + if fixture.profile.startswith("5b") + else ("bicubic_center_crop" if fixture.profile.endswith("causal") else "bicubic") + ) + assert pipe.config.proxy_resize_mode == expected_resize + assert tuple(pipe.config.proxy_image_size) == (480, 832) + assert file_digests(fixture.source) == fixture.hashes + + def test_original_save_load_produces_self_contained_standard_snapshot(self, original_checkpoint, tmp_path): + fixture = original_checkpoint + pipe = EdenLivePipeline.from_pretrained( + fixture.source, profile=fixture.profile, dtype=torch.float32, local_files_only=True + ) + snapshot = tmp_path / "standard" + pipe.save_pretrained(snapshot) + index = json.loads((snapshot / "model_index.json").read_text()) + assert "_edenlive_original" not in index + assert index["decoder"] == ["diffusers", "EdenLiveLightTAEDecoder"] + restored = EdenLivePipeline.from_pretrained(snapshot, local_files_only=True) + assert type(restored) is fixture.pipeline_class + assert isinstance(restored.decoder, EdenLiveLightTAEDecoder) + for name in ("transformer", "vae", "decoder", "proxy_encoder", "proxy_projection"): + expected, actual = getattr(pipe, name).state_dict(), getattr(restored, name).state_dict() + for key, value in expected.items(): + torch.testing.assert_close(actual[key], value, atol=0, rtol=0) + assert file_digests(fixture.source) == fixture.hashes + + @pytest.mark.parametrize("dtype_kind", ["scalar", "mapping"]) + def test_original_dtype_override_and_component_identity(self, original_checkpoint, dtype_kind): + fixture = original_checkpoint + decoder = RecordingDecoder() + text_encoder = fixture.components["text_encoder"].to(dtype=torch.float64) + dtype = torch.float32 if dtype_kind == "scalar" else {"default": torch.float64, "proxy_encoder": torch.float32} + pipe = EdenLivePipeline.from_pretrained( + fixture.source, + profile=fixture.profile, + decoder=decoder, + text_encoder=text_encoder, + dtype=dtype, + local_files_only=True, + ) + assert pipe.decoder is decoder and decoder.weight.dtype == torch.float64 + assert pipe.text_encoder is text_encoder and text_encoder.dtype == torch.float64 + assert pipe.transformer.dtype == (torch.float32 if dtype_kind == "scalar" else torch.float64) + assert pipe.proxy_encoder.dtype == torch.float32 + assert pipe.proxy_projection.dtype == (torch.float32 if dtype_kind == "scalar" else torch.float64) + assert file_digests(fixture.source) == fixture.hashes + + def test_original_decoder_none_skips_missing_tiny_weights(self, original_checkpoint): + fixture = original_checkpoint + size = fixture.profile.split("-")[0] + for path in (fixture.source / f"decoder-{size}").iterdir(): + path.unlink() + with pytest.raises((OSError, ValueError), match="decoder|lighttae|No such file|does not appear"): + EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) + pipe = EdenLivePipeline.from_pretrained( + fixture.source, profile=fixture.profile, decoder=None, local_files_only=True + ) + assert pipe.decoder is None + assert isinstance(pipe.vae, AutoencoderKLWan) + + def test_original_custom_weight_names_and_shared_proxy(self, original_checkpoint): + fixture = original_checkpoint + transformer_path = fixture.transformer_dir / "diffusion_pytorch_model.safetensors" + custom_transformer = transformer_path.with_name("trained-transformer.safetensors") + transformer_path.rename(custom_transformer) + proxy_path = fixture.source / "vision-encoder/projection.safetensors" + custom_proxy = proxy_path.with_name("trained-backbone-and-projector.safetensors") + proxy_path.rename(custom_proxy) + size = fixture.profile.split("-")[0] + decoder_path = next((fixture.source / f"decoder-{size}").glob("*.pth")) + custom_decoder = decoder_path.with_name("trained-decoder.pth") + decoder_path.rename(custom_decoder) + index_path = fixture.source / "model_index.json" + index = json.loads(index_path.read_text()) + index["_edenlive_original"]["_path_overrides"] = { + "transformer": {fixture.profile: {"weights": str(custom_transformer.relative_to(fixture.source))}}, + "proxy_encoder": {"weights": str(custom_proxy.relative_to(fixture.source))}, + "proxy_projection": {"weights": str(custom_proxy.relative_to(fixture.source))}, + "decoder": {size: {"weights": str(custom_decoder.relative_to(fixture.source))}}, + } + index_path.write_text(json.dumps(index)) + hashes = file_digests(fixture.source) + pipe = EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) + torch.testing.assert_close( + pipe.proxy_projection.patch_projection.weight, + fixture.components["proxy_projection"].patch_projection.weight, + atol=0, + rtol=0, + ) + assert isinstance(pipe.decoder, EdenLiveLightTAEDecoder) + assert file_digests(fixture.source) == hashes + + def test_original_profile_validation(self, original_checkpoint): + fixture = original_checkpoint + with pytest.raises(ValueError, match="profile"): + EdenLivePipeline.from_pretrained(fixture.source, local_files_only=True) + + def test_original_download_selected_profile_files(self, original_hub, tmp_path): + hub = original_hub + pipe = EdenLivePipeline.from_pretrained( + hub.repo_id, + profile=hub.fixture.profile, + revision="release-test", + cache_dir=tmp_path / "cache", + token="fixture-token", + local_files_only=False, + force_download=True, + ) + assert type(pipe) is EdenLivePipeline and pipe.config.profile == hub.fixture.profile + assert len(hub.calls.metadata) == len(hub.calls.snapshots) == 1 + options = hub.calls.snapshots[0] + assert set(options["allow_patterns"]) == set(hub.fixture.hashes) + assert options["revision"] == hub.revision + assert options["cache_dir"] == tmp_path / "cache" + assert options["token"] == "fixture-token" and options["force_download"] + assert options["local_files_only"] is False + assert options["user_agent"]["pipeline_class"] == hub.fixture.pipeline_class.__name__ + + @pytest.mark.parametrize("original_checkpoint", ["5b-causal"], indirect=True) + @pytest.mark.parametrize("complete_cache", [False, True]) + def test_original_download_offline_cache(self, original_hub, monkeypatch, tmp_path, complete_cache): + hub = original_hub + + def unexpected_model_info(*args, **kwargs): + pytest.fail("Offline loading must not query the remote repository tree.") + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "model_info", unexpected_model_info) + monkeypatch.setattr(huggingface_hub, "get_cached_repo_tree", unexpected_model_info) + if not complete_cache: + metadata_dir = tmp_path / "metadata-only-cache" + metadata_dir.mkdir() + index_path = metadata_dir / "model_index.json" + index_path.write_text((hub.fixture.source / "model_index.json").read_text()) + + def metadata_download(repository, filename, **kwargs): + assert repository == hub.repo_id and filename == "model_index.json" + assert kwargs["local_files_only"] is True + return str(index_path) + + monkeypatch.setattr(hub.module, "hf_hub_download", metadata_download) + pipe = EdenLivePipeline.from_pretrained( + hub.repo_id, profile=hub.fixture.profile, revision="release-test", local_files_only=True + ) + assert type(pipe) is EdenLivePipeline and pipe.config.profile == hub.fixture.profile + if complete_cache: + assert hub.calls.snapshots == [] + assert hub.calls.metadata[0]["local_files_only"] is True + else: + assert len(hub.calls.snapshots) == 1 + assert hub.calls.snapshots[0]["local_files_only"] is True + assert set(hub.calls.snapshots[0]["allow_patterns"]) == set(hub.fixture.hashes) + + @pytest.mark.parametrize("original_checkpoint", ["14b-causal"], indirect=True) + @pytest.mark.parametrize("options", [{"variant": "fp16"}, {"device_map": "balanced"}, {"use_safetensors": True}]) + def test_original_download_rejects_unsupported_before_weights(self, original_hub, options): + hub = original_hub + with pytest.raises(ValueError, match="does not support|includes PyTorch"): + EdenLivePipeline.from_pretrained( + hub.repo_id, profile=hub.fixture.profile, local_files_only=False, force_download=True, **options + ) + assert len(hub.calls.metadata) == 1 + assert hub.calls.snapshots == [] + + @pytest.mark.parametrize("original_checkpoint", ["5b-bidir"], indirect=True) + def test_original_rejects_unknown_proxy_checkpoint_key(self, original_checkpoint): + fixture = original_checkpoint + path = fixture.source / "vision-encoder/projection.safetensors" + state = {key: value.clone() for key, value in load_file(path).items()} + state["module.encoder.unrecognized.weight"] = torch.ones(1) + save_file(state, path) + with pytest.raises(ValueError, match="Unrecognized proxy checkpoint key"): + EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) + + +class EdenLiveTransformerTesterConfig: + 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, + "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), + "proxy_features": randn_tensor((batch_size, 8, 3, 4, 4), generator=generator, device=torch_device), + "encoder_hidden_states": randn_tensor((batch_size, 12, 16), generator=generator, device=torch_device), + "timestep": torch.tensor([500.0], device=torch_device).expand(batch_size), + } + + +class TestEdenLiveTransformerModel(EdenLiveTransformerTesterConfig): + @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() + 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 * state.tokens_per_frame for layer in state.layers) + if index == 1: + sink_keys = [layer.key[:, : 6 * state.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 * state.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()) + + @torch.no_grad() + def test_text_timeline_uses_complete_segment_bank(self): + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + attn = model.blocks[0].attn2 + hidden = randn_tensor((1, 12, 48), generator=self.generator, device=torch_device) + context = randn_tensor((1, 10, 48), generator=self.generator, device=torch_device) + layout = {"lengths": ((3, 9),), "text_length": 5} + expected = torch.cat((attn(hidden[:, :3], context[:, :5]), attn(hidden[:, 3:], context[:, 5:])), dim=1) + actual = attn(hidden, context, text_timeline=layout) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + local = attn(hidden[:, 6:], context, text_timeline=layout, query_start=6) + torch.testing.assert_close(local, expected[:, 6:], atol=1e-6, rtol=1e-6) + + +class TestEdenLiveCausalCacheOrder: + def test_local_chunk_order_survives_eviction_and_short_tail(self, monkeypatch): + # Capture the real processor's dispatch boundary. Each rank keeps its + # own chunk slices; Ulysses performs the rank-major exchange downstream. + processor = transformer_edenlive.EdenLiveCausalAttnProcessor() + identity = torch.nn.Identity() + attn = SimpleNamespace( + fused_projections=False, + to_q=identity, + to_k=identity, + to_v=identity, + norm_q=identity, + norm_k=identity, + heads=1, + to_out=[identity, identity], + ) + cache = transformer_edenlive._EdenLiveLayerCache(torch.zeros(1, 6, 1, 2), torch.zeros(1, 6, 1, 2)) + seen = [] + + def dispatch(query, key, value, **kwargs): + seen.append((key.clone(), value.clone())) + return query + + monkeypatch.setattr(transformer_edenlive, "dispatch_attention_fn", dispatch) + + def call(ids, commit): + hidden = torch.tensor(ids, dtype=torch.float32).view(1, -1, 1).expand(-1, -1, 2) + rope = (torch.ones(1, 1, 1, 2), torch.zeros(1, 1, 1, 2)) + processor( + attn, + hidden, + rotary_emb=rope, + key_rotary_emb=rope, + layer_cache=cache, + cache_update=commit, + cache_window_tokens=6, + cache_sink_tokens=2, + ) + assert torch.equal(seen[-1][0], seen[-1][1]) + return seen[-1][0][0, :, 0, 0].tolist() + + assert call([0, 1], True) == [0, 1] + assert call([4, 5], True) == [0, 1, 4, 5] + assert call([8, 9], True) == [0, 1, 4, 5, 8, 9] + assert call([12, 13], True) == [0, 1, 8, 9, 12, 13] + stored = cache.key.clone() + assert call([16], False) == [0, 1, 9, 12, 13, 16] + assert cache.num_tokens == 6 and torch.equal(cache.key, stored) + assert call([16], True) == [0, 1, 9, 12, 13, 16] + + +def _parallel_outputs(model, inputs): + with torch.no_grad(): + output = model(**inputs, return_dict=False)[0] + state = model.create_causal_state( + batch_size=1, + height=8, + width=8, + chunk_size=3, + max_context_frames=18, + sink_frames=6, + device=inputs["hidden_states"].device, + dtype=inputs["hidden_states"].dtype, + ) + clean_inputs = {**inputs, "timestep": torch.zeros_like(inputs["timestep"])} + sink_keys = None + for index in range(8): + model(**clean_inputs, causal_state=state, cache_update=True) + assert state.next_frame == 3 * (index + 1) + assert all(layer.num_tokens <= 18 * state.tokens_per_frame // state.cp_size for layer in state.layers) + sink_tokens = 6 * state.tokens_per_frame // state.cp_size + if index == 1: + sink_keys = [layer.key[:, :sink_tokens].clone() for layer in state.layers] + elif index > 1: + for expected, layer in zip(sink_keys, state.layers): + torch.testing.assert_close(layer.key[:, :sink_tokens], expected, rtol=0, atol=0) + tail_inputs = { + **inputs, + "hidden_states": inputs["hidden_states"][:, :, :1], + "proxy_features": inputs["proxy_features"][:, :, :1], + "timestep": inputs["timestep"][:, :16] if inputs["timestep"].ndim == 2 else inputs["timestep"], + } + tail = model(**tail_inputs, causal_state=state, return_dict=False)[0] + assert state.next_frame == 24 + tail_clean = {**tail_inputs, "timestep": torch.zeros_like(tail_inputs["timestep"])} + model(**tail_clean, causal_state=state, cache_update=True) + assert state.next_frame == 25 + return [output.float().cpu(), tail.float().cpu()] + + +def _context_parallel_worker(rank, world_size, port, init_dict, state_dict, inputs, backend, use_fsdp, output_path): + # Spawned workers must establish the same settings as the single-device reference. + enable_full_determinism() + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), RANK=str(rank), WORLD_SIZE=str(world_size)) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 + try: + dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=180)) + model = EdenLiveTransformer3DModel(**init_dict).eval() + model.load_state_dict(state_dict) + model.to(device=device, dtype=dtype) + model.set_attention_backend(backend) + inputs = { + name: tensor.to(device=device, dtype=torch.float32 if name == "timestep" else dtype) + for name, tensor in inputs.items() + } + + def check_timestep(_module, _args, kwargs): + assert kwargs["timestep"].dtype == torch.float32 + assert torch.isin(kwargs["timestep"], torch.tensor([0.0, 375.0], device=device)).all() + + model.register_forward_pre_hook(check_timestep, with_kwargs=True) + from torch.distributed.device_mesh import init_device_mesh + + mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) + model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) + if use_fsdp: + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + + policy = MixedPrecisionPolicy( + param_dtype=dtype, reduce_dtype=torch.float32, output_dtype=dtype, cast_forward_inputs=False + ) + for block in model.blocks: + fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + fully_shard(model, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + outputs = _parallel_outputs(model, inputs) + if rank == 0: + torch.save(outputs, output_path) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +class TestEdenLiveContextParallel(EdenLiveTransformerTesterConfig): + @pytest.mark.context_parallel + @pytest.mark.attention + @pytest.mark.parametrize("world_size", [2, 4]) + @pytest.mark.parametrize( + "backend,use_fsdp", [("native", False), ("_native_flash", False), ("_native_flash", True)] + ) + @pytest.mark.parametrize("num_layers,token_time", [(2, False), (1, True)]) + def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsdp, num_layers, token_time): + if not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < world_size: + pytest.skip(f"This case requires {world_size} CUDA devices with NCCL.") + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + try: + enable_full_determinism() + config = {**self.get_init_dict(), "attention_head_dim": 16, "num_layers": num_layers} + model = EdenLiveTransformer3DModel(**config).eval() + state = {name: tensor.cpu().clone() for name, tensor in model.state_dict().items()} + inputs = {name: tensor.cpu() for name, tensor in self.get_dummy_inputs().items()} + # This value would be rounded to 376 by an indiscriminate BF16 input cast. + inputs["timestep"] = torch.tensor([375.0], dtype=torch.float32) + if token_time: + inputs["timestep"] = torch.full((1, 48), 375.0) + inputs["timestep"][:, :16] = 0 + dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 + model.to(device="cuda:0", dtype=dtype).set_attention_backend(backend) + reference_inputs = { + name: tensor.to(device="cuda:0", dtype=torch.float32 if name == "timestep" else dtype) + for name, tensor in inputs.items() + } + expected = _parallel_outputs(model, reference_inputs) + del model, reference_inputs + torch.cuda.empty_cache() + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + output_path = str(tmp_path / "context_parallel.pt") + mp.spawn( + _context_parallel_worker, + args=(world_size, port, config, state, inputs, backend, use_fsdp, output_path), + nprocs=world_size, + join=True, + ) + actual = torch.load(output_path, map_location="cpu", weights_only=True) + for reference, output in zip(expected, actual): + assert torch.isfinite(output).all() + # Keep the existing small-model CP backend tolerance. + torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2) + finally: + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + + +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) From 845d85372a213b69ffb0908f332fc5ec7c96d2ad Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Fri, 16 Oct 2026 23:25:18 +0800 Subject: [PATCH 2/6] feat: add streaming v2v pipeline --- docs/source/en/api/pipelines/edenlive.md | 327 ++++++++----- examples/edenlive/README.md | 93 ++-- examples/edenlive/inference_edenlive.py | 2 +- .../edenlive/inference_edenlive_streaming.py | 94 ++++ src/diffusers/__init__.py | 4 + src/diffusers/models/__init__.py | 3 +- src/diffusers/models/autoencoders/__init__.py | 2 +- .../autoencoders/autoencoder_edenlive.py | 41 +- .../transformers/transformer_edenlive.py | 2 +- src/diffusers/pipelines/__init__.py | 4 +- src/diffusers/pipelines/edenlive/__init__.py | 6 +- .../pipelines/edenlive/loading_utils.py | 2 +- .../pipelines/edenlive/pipeline_edenlive.py | 456 +++++++++++++++++- .../pipelines/edenlive/pipeline_output.py | 2 +- .../pipelines/edenlive/pipeline_utils.py | 68 ++- .../pipelines/edenlive/video_processor.py | 33 +- .../schedulers/scheduling_edenlive.py | 2 +- src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + tests/pipelines/edenlive/__init__.py | 1 + tests/pipelines/edenlive/test_edenlive.py | 310 +++++++++++- 21 files changed, 1263 insertions(+), 219 deletions(-) create mode 100644 examples/edenlive/inference_edenlive_streaming.py diff --git a/docs/source/en/api/pipelines/edenlive.md b/docs/source/en/api/pipelines/edenlive.md index b3f4bc648e35..8036ed90f86d 100644 --- a/docs/source/en/api/pipelines/edenlive.md +++ b/docs/source/en/api/pipelines/edenlive.md @@ -12,124 +12,108 @@ specific language governing permissions and limitations under the License. # EdenLive -[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) generates video from a reference image, text, and a -proxy video. It uses a trained DINOv3 encoder and projection to condition a Wan-based diffusion transformer. +[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) generates video from a reference image, prompt, and +RGB proxy video. [`EdenLivePipeline`] supports bidirectional generation and causal generation with persistent history. -[`EdenLivePipeline`] supports both bidirectional generation and successive causal chunks with clean latent history. -Select a released configuration with `profile` when loading the original checkpoint repository. Defaults apply when -`num_inference_steps` and `guidance_scale` are omitted: +## Basic usage -| Profile | Sampling | Default steps | Guidance scale | Default height x width | -|---|---|---|---|---| -| `5b-bidir` | UniPC over the full video | 48 | 3 | 704 x 1280 | -| `14b-bidir` | Euler over the full video | 30 | 3 | 480 x 832 | -| `5b-causal` | Stochastic Euler per chunk | 4 | 1 | 704 x 1280 | -| `14b-causal` | Stochastic Euler per chunk | 4 | 1 | 480 x 832 | +Prepare an RGB reference image and a proxy video with at least 49 frames sampled at 24 fps. The reference image +provides appearance; the proxy video provides scene structure and motion. For installation, see the +[example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive#setup). -The released causal checkpoints use fixed four-step schedules. Changing `num_inference_steps` alone to another -value raises an error; it does not select a new distilled schedule. Guidance greater than one enables the negative -prompt branch. - -## Generate a video - -```py +```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", profile="5b-causal" -).to("cuda") +pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-causal") +pipe.to("cuda") + +image = load_image("reference.png") +video = load_video("proxy-24fps.mp4") +prompt = "A character explores a snowy landscape." frames = pipe( - prompt="A character explores a snowy landscape.", - image=load_image("reference.png"), - video=load_video("proxy-24fps.mp4"), + image=image, + video=video, + prompt=prompt, num_frames=49, - generator=torch.Generator("cuda").manual_seed(42), + generator=torch.Generator().manual_seed(42), ).frames[0] export_to_video(frames, "edenlive.mp4", fps=24) ``` -## Load and save checkpoints - -Use `EdenLivePipeline.from_pretrained()` for both formats: +`from_pretrained()` loads the model, `pipe(...)` generates the video, and `.frames[0]` selects the first video in +the output batch. Each offline call starts a new video. Use the [streaming API](#streaming-causal-video) to continue +generation across calls. -| Format | Required metadata | Profile selection | -|---|---|---| -| Original release | Root `model_index.json` with the EdenLive original-format marker | Pass one of the four complete profile names | -| Saved Diffusers pipeline | `model_index.json` and saved component directories | Omit `profile` to use the saved value | +## Choose a profile -The original release loader reads the existing checkpoint files without writing converted weights. It fetches only -the files needed for the selected profile and components. The same entry accepts a Hub repository ID or local -directory, plus `revision`, `cache_dir`, `token`, and `local_files_only=True` for cached or offline loading. +Select `profile` when loading an original release. These defaults apply when generation options are omitted: -The original release's default precision is BF16 for the transformer, VAE, UMT5 text encoder, CLIP encoder, and -LightTAE, with FP32 parameters for the DINO encoder and projection. An explicit `dtype` (or `torch_dtype`) can be -a single dtype or a mapping by component name. Initialized component overrides retain their identity and dtype. -Active `device_map`, `variant`, and quantization options are not supported by the original-format loader. Load on CPU -and then place the pipeline on the desired device, as in the example above. +| Profile | Sampling | Steps | Guidance scale | Height × width | Example fps | +|---|---|---|---|---|---| +| `5b-bidir` | Bidirectional | 48 | 3 | 704 × 1280 | 24 | +| `14b-bidir` | Bidirectional | 30 | 3 | 480 × 832 | 16 | +| `5b-causal` | Causal | 4 per chunk | 1 | 704 × 1280 | 24 | +| `14b-causal` | Causal | 4 per chunk | 1 | 480 × 832 | 16 | -Save and reload a standard pipeline with: - -```py -pipe.save_pretrained("edenlive-saved") -pipe = EdenLivePipeline.from_pretrained("edenlive-saved").to("cuda") -``` - -The saved configuration retains the full profile and decoder choice. A conflicting explicit profile raises an error. -Older standard snapshots with size-only profiles or the former causal class name are normalized in memory; newly -saved snapshots use `EdenLivePipeline`. The generic `DiffusionPipeline.from_pretrained()` entry is not supported for -original releases. +Causal profiles use four fixed denoising steps per chunk. Bidirectional profiles support `num_inference_steps`. +Set `guidance_scale` above one to use a `negative_prompt`. ## Prepare the inputs -The reference image and proxy video should depict the same composition. Supply a full RGB proxy video already -sampled at the desired output frame rate. The Python pipeline does not resample timestamps; the command-line -example below does. The pipeline then selects proxy frames at the VAE temporal stride once. +Supply proxy frames sampled at the desired output frame rate. The pipeline selects frames at the VAE temporal +stride; timestamp resampling is handled by the [CLI example](#multiple-gpus-and-the-command-line-example). | Input | Accepted layout and values | |---|---| -| `image` | RGB PIL image, or floating-point NumPy/tensor image with values in [0, 1] | +| `image` | RGB PIL image, or floating-point NumPy/tensor image in [0, 1] | | `video` | List of RGB PIL frames; tensor `(T, 3, H, W)` or `(B, T, 3, H, W)`; NumPy `(T, H, W, 3)` or `(B, T, H, W, 3)` with floating-point values in [0, 1] | -| `prompt` | String or list of strings; alternatively provide `prompt_embeds` or a supported temporal prompt format | -| `height`, `width` | Output dimensions, both multiples of 32 for 5B or 16 for 14B | -| `num_frames` | Positive frame count of the form \\(4k+1\\), no greater than the input video length | +| `prompt` | String or list of strings; use `prompt_embeds` for precomputed text embeddings | +| `height`, `width` | Output dimensions, multiples of 32 for 5B or 16 for 14B | +| `num_frames` | Positive frame count of the form $4k+1$, within the input video length | -Here `B` is the video batch size, `T` the decoded frame count, and `H` and `W` the frame height and width. The reference -image is resized and center-cropped to the output dimensions. Proxy RGB frames are processed at `proxy_image_size` -(480 x 832 by default) with the profile's `proxy_resize_mode` before DINO encoding. The command-line example also -crops and resizes proxy frames to the output aspect ratio before passing them to the pipeline. Omitting `num_frames` -selects the largest valid input prefix; missing frames are not padded or repeated by the pipeline. +Here `B` is batch size, `T` is frame count, and `H` and `W` are image height and width. The reference image is +resized and center-cropped to the output dimensions. Omitting `num_frames` uses the largest valid input prefix. -Both model sizes use temporal compression by four. A video of \\(F\\) frames has -\\(T_{\mathrm{latent}}=(F-1)/4+1\\) latent frames. The 5B model uses 48 latent channels and spatial compression by 16; -the 14B model uses 16 channels and spatial compression by 8. +## Change prompts during a video -## Choose a decoder +### Causal chunk prompts -Original releases load the matching [`EdenLiveLightTAEDecoder`] by default. Use the full Wan VAE with: +Use `chunk_prompts` for either causal profile. This example reuses the pipeline, image, and video from +[Basic usage](#basic-usage). A 49-frame video has five chunks, so it takes five prompts: -```py -pipe = EdenLivePipeline.from_pretrained( - "VAST-AI/EdenLive", profile="14b-causal", decoder=None -) +```python +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.", +] + +frames = pipe( + image=image, + video=video, + chunk_prompts=chunk_prompts, + num_frames=49, + generator=torch.Generator().manual_seed(42), +).frames[0] +export_to_video(frames, "chunk-prompts.mp4", fps=24) ``` -Pass `decoder=my_decoder` to use an initialized compatible decoder. The full VAE remains available to encode the -reference image. A custom decoder must implement `decode(latents, return_dict=False)` and return a tuple whose -first item is an RGB tensor in `(B, 3, T, H, W)` layout and `[-1, 1]` range. Its `config.latent_channels`, `scale_factor_temporal`, and -`scale_factor_spatial` must match the VAE. Use a compatible `ModelMixin` decoder for pipeline-wide device placement -and standard serialization. Plain `torch.nn.Module` decoders support instance loading and decoding with explicit -component placement. A decoder receives raw, unnormalized VAE latents; the pipeline applies the inverse latent -normalization before calling it. Selecting `output_type="latent"` bypasses this conversion and decoding. +`chunk_prompts` supplies one text prompt per chunk for a single video. + +### Bidirectional text timeline -## Temporal prompts +For `5b-bidir`, use `text_timeline` to assign prompts to time ranges. Reuse the image and video from Basic usage +and load the bidirectional profile: -The `5b-bidir` profile accepts `text_timeline` instead of a single prompt. Each segment contains `start_sec`, `end_sec`, -and `text`. For example, this timeline describes 49 frames sampled at 24 fps: +```python +pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-bidir") +pipe.to("cuda") -```py timeline = { "fps": 24, "num_frames": 49, @@ -138,29 +122,127 @@ timeline = { {"start_sec": 1.0, "end_sec": 2.0, "text": "The character approaches a warm campfire."}, ], } + +frames = pipe( + image=image, + video=video, + text_timeline=timeline, + num_frames=49, + generator=torch.Generator().manual_seed(42), +).frames[0] +export_to_video(frames, "timeline.mp4", fps=24) ``` -Pass it as `text_timeline=timeline` with the image and video. Segments must start at zero, be contiguous, and cover -the timestamp of the last requested frame. Each segment must receive at least one latent frame. A latent frame -exactly at a segment end belongs to the segment on its left. The timeline frame count must match the requested -video, and its `fps` describes the supplied frames; it does not trigger resampling. For a batch, provide one timeline -dictionary per video. Timeline prompts cannot be combined with precomputed text embeddings. Other profiles reject -`text_timeline`. +Match the timeline's `fps` and `num_frames` to the input and requested output. Use contiguous segments from zero +through the last frame timestamp, with at least one latent frame per segment. -Causal profiles accept `chunk_prompts`, a list containing one complete prompt per latent chunk. For the released -three-frame chunk size, 49 decoded frames produce 13 latent frames and require five prompts; the last chunk has -one latent frame. This interface requires batch size one and cannot be combined with precomputed positive text -embeddings. A single `prompt` applies to all chunks when `chunk_prompts` is omitted. Bidirectional profiles reject -`chunk_prompts`. +## Streaming causal video -Each pipeline call starts with fresh causal history. The released 18-frame latent context includes six retained -sink frames and the current chunk. Within a call, completed clean chunks are committed to the cache before the -next chunk is generated. +Both causal profiles support successive calls with shared history. Initialize one state, then pass three already +sampled RGB proxy images per chunk. This example changes the prompt for the second chunk and saves the 21 output +frames as one video: + +```python +import numpy as np +import torch +from diffusers import EdenLivePipeline +from diffusers.utils import export_to_video, load_image + +pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-causal") +pipe.to("cuda") + +proxy_paths = [ + ["proxy/chunk-0-0.png", "proxy/chunk-0-1.png", "proxy/chunk-0-2.png"], + ["proxy/chunk-1-0.png", "proxy/chunk-1-1.png", "proxy/chunk-1-2.png"], +] +prompts = ["A character walks through snow.", "The character approaches a campfire."] + +state = pipe.create_stream_state( + image=load_image("reference.png"), + prompt=prompts[0], + generator=torch.Generator().manual_seed(42), + output_type="np", +) +chunks = [] +try: + for paths, current_prompt in zip(proxy_paths, prompts): + result = pipe( + stream_state=state, + proxy_frames=[load_image(path) for path in paths], + prompt=current_prompt, + ) + chunks.append(result.frames[0]) +finally: + pipe.close_stream(state) + +export_to_video(np.concatenate(chunks, axis=0), "stream.mp4", fps=24) +``` + +Each call returns its chunk immediately: 9 RGB frames initially, then 12 per call. For incremental video-file +writing, see [`inference_edenlive_streaming.py`](https://github.com/huggingface/diffusers/blob/main/examples/edenlive/inference_edenlive_streaming.py). + +Set resolution, guidance, negative prompt, generator, and output format in `create_stream_state()`. Keep these +settings and the pipeline's device and dtype fixed for the stream. Update `prompt` or `prompt_embeds` in each +generation call; omitting both retains the current prompt. Run calls sharing a pipeline serially. + +Streaming uses batch size one and three proxy images per call. Tensor inputs have shape `(3, 3, H, W)`; NumPy +inputs use `(3, H, W, 3)`, with finite floating-point RGB values in [0, 1]. RGB output uses LightTAE. +Set `output_type="latent"` at initialization to receive three latent frames per call. For 14B, keep VAE slicing +and tiling disabled. Close a state after a model execution failure and initialize a new one. + +## Load and save checkpoints + +`EdenLivePipeline.from_pretrained()` accepts a Hub repository or a local directory. Original releases use a root +`model_index.json` and the selected profile's original weights. Saved Diffusers pipelines retain their profile and +decoder choice: + +```python +pipe.save_pretrained("edenlive-saved") +pipe = EdenLivePipeline.from_pretrained("edenlive-saved").to("cuda") +``` + +For local or cached weights: + +```python +pipe = EdenLivePipeline.from_pretrained( + "/path/to/EdenLive", profile="14b-causal", local_files_only=True +).to("cuda") +``` + +Original releases load the transformer, VAE, text encoder, CLIP, and LightTAE in BF16, with DINO and its projection +in FP32. Use `dtype` for an explicit dtype or per-component mapping. Original-release loading supports +`revision`, `cache_dir`, and `token`; `device_map`, `variant`, and quantization options require a saved +Diffusers pipeline. + +### CPU offload + +Load the pipeline on CPU and enable model offload before generation: + +```python +pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-causal") +pipe.enable_model_cpu_offload() +``` + +The offline and streaming calls above use the same interface with CPU offload enabled. + +### Offline decoder selection + +Original releases use [`EdenLiveLightTAEDecoder`] by default. To decode offline videos with the full Wan VAE, +pass `decoder=None` when loading: + +```python +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive", profile="5b-causal", decoder=None +).to("cuda") +``` + +Pass an initialized decoder as `decoder=my_decoder` to use a compatible custom decoder. +`output_type="latent"` returns normalized diffusion latents before decoding. ## Multiple GPUs and the command-line example -The inference example resamples video timestamps and uses Ulysses context parallelism with PyTorch FSDP2 on the -same ranks. Run with Python for one GPU or `torchrun` for two or four GPUs: +The CLI resamples proxy video timestamps and supports Ulysses context parallelism with PyTorch FSDP2 on two or four +GPUs. From the Diffusers repository root: ```bash torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ @@ -170,58 +252,51 @@ torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py --num-frames 81 --output edenlive.mp4 ``` -Use `--full-vae` for full VAE decoding. Only rank zero exports the video; the example also records timing and peak -memory. Keep the parallel degree fixed when comparing reproducible outputs: BF16 matrix multiplication shapes -can change predictions, as also observed in the reference implementation. - -The example defaults to 24 fps for 5B and 16 fps for 14B. See the -[example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive) for installation, single-GPU -usage, temporal prompt JSON files, and the timing report format. +See the [example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive) for installation, +single-GPU commands, prompt files, and output reports. ## EdenLivePipeline -The constructor registers initialized components. `from_pretrained()` loads a release or saved pipeline; -`__call__()` prepares the image, proxy features, text, and initial noise before running the profile's sampling path. -Use `generator` for reproducible noise and `callback_on_step_end` to inspect or replace supported tensors after a -step. In causal mode callbacks receive the current chunk, while their step indices count across all chunks. +Load weights with `from_pretrained()`. Generate a complete video with `pipe(image=..., video=..., prompt=...)`, +or initialize a stream with `create_stream_state()` and generate chunks with +`pipe(stream_state=..., proxy_frames=..., prompt=...)`. Call `close_stream()` when the stream ends. [[autodoc]] EdenLivePipeline - all - __call__ + - create_stream_state + - close_stream ## EdenLiveLightTAEDecoder -This decoder implements the released Tiny decoding path without an encoder. The 5B variant uses 48 latent channels -and `patch_size=2`; the 14B variant uses 16 channels and `patch_size=1`. `decode()` accepts raw VAE latents in -`(B, C, T_latent, H_latent, W_latent)` layout and returns `DecoderOutput.sample` in `(B, 3, T, H, W)` layout with -values in `[-1, 1]`. Unlike pipeline frame output, its channel axis precedes the time axis. +`decode()` takes raw VAE latents in `(B, C, T_latent, H_latent, W_latent)` layout and returns RGB in +`(B, 3, T, H, W)`, with values in [-1, 1]. To decode the pipeline's `output_type="latent"` yourself, first invert +the VAE normalization: -To decode `output_type="latent"` yourself, invert the VAE normalization first: - -\\[ +$$ z_{\mathrm{raw}} = z_{\mathrm{diffusion}} \odot \sigma + \mu. -\\] +$$ + +Here $\mu$ and $\sigma$ are `vae.config.latents_mean` and `vae.config.latents_std`, broadcast over batch, time, +and spatial dimensions; $\odot$ is element-wise multiplication. Match the latents to the decoder's device and dtype. -Here \\(\mu\\) and \\(\sigma\\) are the channel-wise `vae.config.latents_mean` and `vae.config.latents_std`, and -\\(\odot\\) denotes element-wise multiplication with broadcasting. Move the unnormalized latents to the decoder's -device and dtype before decoding. The pipeline handles this when returning decoded frames. Every `decode()` call -creates fresh temporal state and supports a short final chunk; separate calls do not share decoder history. +Pass an `EdenLiveLightTAEDecodeCache` to retain decoder history across calls. Without a cache, each decode starts +a new video. The pipeline manages this cache for streaming generation. [[autodoc]] EdenLiveLightTAEDecoder - decode ## EdenLivePipelineOutput -With `return_dict=True`, read the result from `.frames`: +Read generated frames from `.frames`, or use `return_dict=False` to receive `(frames,)`: | `output_type` | `.frames` layout | Values | |---|---|---| | `"np"` (default) | NumPy `(B, T, H, W, 3)` | Floating-point RGB in [0, 1] | | `"pt"` | Tensor `(B, T, 3, H, W)` | Floating-point RGB in [0, 1] | | `"pil"` | One list of PIL RGB frames per video | PIL images | -| `"latent"` | Tensor `(B, C, T_latent, H_latent, W_latent)` | Normalized diffusion latents, without RGB bounds | +| `"latent"` | Tensor `(B, C, T_latent, H_latent, W_latent)` | Normalized diffusion latents | -With `return_dict=False`, the pipeline returns `(frames,)` with the same selected representation. Select -`result.frames[0]` to export the first decoded video with `export_to_video()`. +Use `result.frames[0]` to select the first video for `export_to_video()`. [[autodoc]] pipelines.edenlive.pipeline_output.EdenLivePipelineOutput diff --git a/examples/edenlive/README.md b/examples/edenlive/README.md index 62dd10b83a75..cbc3f1551701 100644 --- a/examples/edenlive/README.md +++ b/examples/edenlive/README.md @@ -1,22 +1,24 @@ # EdenLive inference -`inference_edenlive.py` runs offline video generation from a reference image, text, and a proxy video. It supports -the four released profiles through `EdenLivePipeline`, with single-GPU execution or Ulysses context parallelism -and PyTorch FSDP2 across two or four GPUs. See the [pipeline guide](../../docs/source/en/api/pipelines/edenlive.md) -for the Python API, component loading, decoder overrides, and output formats. +Render game from a reference image, prompt, and proxy video with `EdenLivePipeline`. +See the [pipeline guide](../../docs/source/en/api/pipelines/edenlive.md) for the Python API. ## Setup -Run these commands from the Diffusers repository root in an environment with a CUDA-compatible PyTorch build: +Install a CUDA-compatible [PyTorch](https://pytorch.org/get-started/locally/) build and matching torchvision. +From the Diffusers repository root, install the pipeline dependencies: ```bash -pip install -e . -pip install --upgrade -r examples/edenlive/requirements.txt +pip install -e . "transformers>=4.56.0" accelerate sentencepiece ftfy ``` -Transformers must provide `DINOv3ViTModel`, `UMT5EncoderModel`, and `CLIPVisionModel`. Multi-GPU execution also -requires PyTorch's composable FSDP2 API (`torch.distributed.fsdp.fully_shard`). The release directory or Hub repository -must contain EdenLive's root `model_index.json` and the original weights for the selected profile. +For the CLI examples, add video reading and writing: + +```bash +pip install av imageio imageio-ffmpeg +``` + +The [example requirements](requirements.txt) collect both groups, plus PyTorch and torchvision. ## Single GPU @@ -28,8 +30,8 @@ python examples/edenlive/inference_edenlive.py \ --num-frames 49 --output edenlive.mp4 ``` -`--model` also accepts a local release or saved Diffusers pipeline. Add `--local-files-only` to use only local or -cached weights. The CLI requires `--profile`; for a saved pipeline it must match the saved configuration. +`--model` accepts a Hub repository, local release, or saved Diffusers pipeline. Select its profile with `--profile`. +Use `--local-files-only` for local or cached weights. | Profile | Default height x width | Default fps | Default steps | Guidance scale | |---|---|---|---|---| @@ -38,18 +40,16 @@ cached weights. The CLI requires `--profile`; for a saved pipeline it must match | `5b-causal` | 704 x 1280 | 24 | 4 per chunk | 1 | | `14b-causal` | 480 x 832 | 16 | 4 per chunk | 1 | -The CLI resamples source timestamps to `--fps`, center-crops to the output aspect ratio, and resizes the frames. -`--num-frames` must be positive, have the form $4k+1$, and fit the clip at the selected frame rate. If omitted, the -largest valid prefix is used. Output height and width must be multiples of 32 for 5B or 16 for 14B. Unlike this CLI, -the Python pipeline expects frames that are already sampled at the desired output frame rate. +The CLI resamples to `--fps`, center-crops, and resizes the input. Set `--num-frames` to a positive $4k+1$ within +the resampled clip, or omit it to use the largest valid prefix. Custom height and width use multiples of 32 for 5B +and 16 for 14B. -The released causal schedules contain four fixed evaluation points. Keep `--num-inference-steps` at four or omit it -for causal profiles. The bidirectional profiles accept a different positive step count. `--guidance-scale` above -one enables the negative prompt branch; supply its text with `--negative-prompt`. +Causal profiles use four fixed steps per chunk. Bidirectional profiles support `--num-inference-steps`. +Set `--guidance-scale` above one to use `--negative-prompt`. ## Multiple GPUs -Launch one process per GPU, with two or four processes: +Use two or four GPUs with Ulysses context parallelism and PyTorch FSDP2 (`torch.distributed.fsdp.fully_shard`): ```bash torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ @@ -59,10 +59,7 @@ torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py --num-frames 81 --output edenlive.mp4 ``` -Set `--nproc_per_node=4` for four GPUs. The transformer uses Ulysses context parallelism and FSDP2 over the same -ranks; other components are replicated. Add `--no-fsdp` to retain context parallelism with replicated transformer -weights. Only rank zero writes the video and JSON report. When comparing seeded results, keep the parallel degree -and attention backend fixed because BF16 matrix multiplication shapes can affect predictions. +Set `--nproc_per_node=4` for four GPUs. Add `--no-fsdp` to use context parallelism with replicated transformer weights. ## Temporal prompts @@ -79,12 +76,11 @@ For `5b-bidir`, pass `--text-timeline timeline.json`. This example describes 49 } ``` -Match the JSON `fps` and `num_frames` to `--fps` and `--num-frames`. Segments start at zero, are contiguous, and -cover the last frame timestamp; each segment must include at least one latent frame. A latent frame at an exact -segment boundary uses the segment on its left. Other profiles do not accept timelines. +Match the JSON `fps` and `num_frames` to the CLI values. Use contiguous segments covering zero through the last +frame timestamp, with at least one latent frame per segment. -For either causal profile, pass `--chunk-prompts prompts.json` containing a JSON list of complete prompts. A 49-frame -video has 13 latent frames, so the released three-frame chunks require five prompts: +For either causal profile, pass `--chunk-prompts prompts.json` with one prompt per chunk. A 49-frame video uses +five chunks, including a shorter final chunk: ```json [ @@ -96,18 +92,33 @@ video has 13 latent frames, so the released three-frame chunks require five prom ] ``` -The final chunk may be shorter. The CLI still requires `--prompt` when either temporal format is supplied; the -timeline or chunk prompts determine the positive captions used for those segments or chunks. Each invocation -starts a fresh causal history. +Include `--prompt` in the command; the timeline or chunk prompts supply the text used during generation. + +## Streaming from proxy images + +`inference_edenlive_streaming.py` supports both causal profiles. Pass ordered RGB images with `--proxy-frames`, +already sampled at latent cadence: three images per chunk. + +```bash +python examples/edenlive/inference_edenlive_streaming.py \ + --model /path/to/EdenLive --profile 5b-causal \ + --image reference.png --prompt "A character walks through snow" \ + --proxy-frames proxy/000.png proxy/004.png proxy/008.png proxy/012.png proxy/016.png proxy/020.png \ + --chunk-prompts stream_prompts.json --output stream.mp4 --fps 24 --local-files-only +``` + +The optional `stream_prompts.json` contains one string per chunk, for example +`["A character walks through snow", "A character walks through fire"]`. Omit it to use `--prompt` throughout. + +The script writes video incrementally: 9 frames for the first chunk, then 12 per chunk. Set `--fps` to the desired +playback rate (24 by default). For application integration, see the +[streaming API](../../docs/source/en/api/pipelines/edenlive.md#streaming-causal-video). -## Decoding and reports +## Offline options and reports -LightTAE is the default decoder for original releases. Add `--full-vae` to use the full Wan VAE for output decoding. -The VAE always encodes the reference image. `--vae-tiling` enables tiling on that VAE; it does not enable LightTAE -tiling. Use `--seed` to select the noise seed and `--attention-backend` to select an available Diffusers attention -backend. Run `python examples/edenlive/inference_edenlive.py --help` for all options. +LightTAE is the default decoder for original releases. Use `--full-vae` for Wan VAE decoding and `--vae-tiling` for +tiling on the Wan VAE. Set the noise seed with `--seed` and the attention backend with `--attention-backend`. +Run `python examples/edenlive/inference_edenlive.py --help` for all options. -The output path selects the video filename. A JSON file with the same stem records the CLI arguments, resampled -source frame indices, output fps and frame count, world size, FSDP2 status, and per-rank elapsed seconds and peak -allocated/reserved GPU memory in bytes. Timing covers the pipeline call, including conditioning and decoding; -it excludes checkpoint loading, input video preprocessing, and video export. +The offline script writes a JSON report beside the video with arguments, sampled frame indices, timing, and peak GPU +memory. Timing covers the pipeline call, including conditioning and decoding. diff --git a/examples/edenlive/inference_edenlive.py b/examples/edenlive/inference_edenlive.py index c549ffb56640..f0b4a7a74e11 100644 --- a/examples/edenlive/inference_edenlive.py +++ b/examples/edenlive/inference_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. diff --git a/examples/edenlive/inference_edenlive_streaming.py b/examples/edenlive/inference_edenlive_streaming.py new file mode 100644 index 000000000000..a48662eb0bfe --- /dev/null +++ b/examples/edenlive/inference_edenlive_streaming.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. + +"""Generate successive EdenLive chunks from ordered, already sampled proxy images.""" + +import argparse +import json +from pathlib import Path + +import av +import numpy as np +import torch +from PIL import Image + +from diffusers import EdenLivePipeline + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="VAST-AI/EdenLive") + parser.add_argument("--profile", choices=("5b-causal", "14b-causal"), required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--proxy-frames", nargs="+", required=True, help="Ordered, already sampled RGB image paths.") + parser.add_argument("--chunk-prompts", type=Path, help="JSON file containing one prompt string per chunk.") + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--local-files-only", action="store_true") + args = parser.parse_args() + if len(args.proxy_frames) % 3: + parser.error("--proxy-frames must contain a multiple of three image paths.") + if args.fps <= 0: + parser.error("--fps must be positive.") + prompts = None + if args.chunk_prompts is not None: + prompts = json.loads(args.chunk_prompts.read_text()) + if ( + not isinstance(prompts, list) + or len(prompts) != len(args.proxy_frames) // 3 + or not all(isinstance(prompt, str) for prompt in prompts) + ): + parser.error("--chunk-prompts must contain one string per chunk.") + pipe = EdenLivePipeline.from_pretrained( + args.model, profile=args.profile, local_files_only=args.local_files_only + ).to("cuda") + with Image.open(args.image) as image: + state = pipe.create_stream_state( + image=image.convert("RGB"), + prompt=args.prompt, + height=args.height, + width=args.width, + generator=torch.Generator().manual_seed(args.seed), + output_type="np", + ) + try: + args.output.parent.mkdir(parents=True, exist_ok=True) + with av.open(str(args.output), mode="w") as container: + stream = container.add_stream("libx264", rate=args.fps) + stream.height, stream.width = state.height, state.width + stream.pix_fmt = "yuv420p" + for index in range(0, len(args.proxy_frames), 3): + frames = [] + for path in args.proxy_frames[index : index + 3]: + with Image.open(path) as frame: + frames.append(frame.convert("RGB")) + output = pipe( + stream_state=state, proxy_frames=frames, prompt=None if prompts is None else prompts[index // 3] + ) + for frame in output.frames[0]: + pixels = np.clip(frame * 255, 0, 255).round().astype(np.uint8) + for packet in stream.encode(av.VideoFrame.from_ndarray(pixels, format="rgb24")): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + finally: + pipe.close_stream(state) + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 896c324230ac..da7846e3adf2 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -288,6 +288,7 @@ "EasyAnimateTransformer3DModel", "EdenLiveCausalState", "EdenLiveLightTAEDecoder", + "EdenLiveLightTAEDecodeCache", "EdenLiveProxyProjection", "EdenLiveTransformer3DModel", "ErnieImageTransformer2DModel", @@ -666,6 +667,7 @@ "EasyAnimatePipeline", "EdenLivePipeline", "EdenLivePipelineOutput", + "EdenLiveStreamState", "ErnieImagePipeline", "Flux2KleinInpaintPipeline", "Flux2KleinKVPipeline", @@ -1192,6 +1194,7 @@ DreamLiteUNetModel, EasyAnimateTransformer3DModel, EdenLiveCausalState, + EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel, @@ -1546,6 +1549,7 @@ EasyAnimatePipeline, EdenLivePipeline, EdenLivePipelineOutput, + EdenLiveStreamState, ErnieImagePipeline, Flux2KleinInpaintPipeline, Flux2KleinKVPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index b89ad4f8aaad..cc0d266505b2 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -31,7 +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_edenlive"] = ["EdenLiveLightTAEDecodeCache", "EdenLiveLightTAEDecoder"] _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"] _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"] @@ -221,6 +221,7 @@ AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, + EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder, Kandinsky6SRVAE, LongCatAudioDiTVae, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index a9664df7b8c0..5fa81587e65a 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -1,7 +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_edenlive import EdenLiveLightTAEDecodeCache, 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 index 14fad3d18ed8..e0a4abaf26db 100644 --- a/src/diffusers/models/autoencoders/autoencoder_edenlive.py +++ b/src/diffusers/models/autoencoders/autoencoder_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 NVIDIA CORPORATION & AFFILIATES and The HuggingFace Team. All rights reserved. +# 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. @@ -12,6 +12,8 @@ # 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 @@ -22,6 +24,17 @@ from .vae import DecoderOutput +@dataclass +class EdenLiveLightTAEDecodeCache: + """Temporal memory for one LightTAE stream with a fixed model, batch, shape, device and dtype.""" + + dec_state: dict[int, torch.Tensor] = field(default_factory=dict) + + def reset(self): + """Discard temporal memory and restore first-chunk trimming.""" + self.dec_state.clear() + + class EdenLiveTAEClamp(nn.Module): def forward(self, hidden_states) -> torch.Tensor: return torch.tanh(hidden_states / 3) * 3 @@ -67,7 +80,7 @@ class EdenLiveLightTAEDecoder(ModelMixin, ConfigMixin): 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. + `decode` call unless an explicit cache is supplied, and is never retained on the model. Parameters: latent_channels (`int`, *optional*, defaults to 16): @@ -141,8 +154,10 @@ def _decode_chunk(self, latents, state): return frames.permute(0, 2, 1, 3, 4) * 2 - 1 @apply_forward_hook - def decode(self, latents: torch.Tensor, return_dict: bool = True): - """Decode raw VAE latents using fresh temporal state for each call. + def decode( + self, latents: torch.Tensor, return_dict: bool = True, *, cache: EdenLiveLightTAEDecodeCache | None = None + ): + """Decode raw VAE latents, optionally continuing an existing stream. Args: latents (`torch.Tensor`): @@ -152,31 +167,39 @@ def decode(self, latents: torch.Tensor, return_dict: bool = True): 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. + cache (`EdenLiveLightTAEDecodeCache`, *optional*): + Temporal memory shared across calls. Only the first chunk is trimmed; subsequent calls produce + `T_latent * scale_factor_temporal` frames. Without a cache, every call starts a new video. 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 only within this call, including a short final chunk. + scale_factor_temporal + 1` frames for a new video, or `T_latent * scale_factor_temporal` for an + existing cache, 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.") - state = {} + state = {} if cache is None else cache.dec_state 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]: + def forward( + self, latents: torch.Tensor, return_dict: bool = True, *, cache: EdenLiveLightTAEDecodeCache | None = None + ) -> 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. + cache (`EdenLiveLightTAEDecodeCache`, *optional*): + Temporal memory carried across decode calls. 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) + return self.decode(latents, return_dict=return_dict, cache=cache) diff --git a/src/diffusers/models/transformers/transformer_edenlive.py b/src/diffusers/models/transformers/transformer_edenlive.py index c05482a2b987..58a57b3b7dee 100644 --- a/src/diffusers/models/transformers/transformer_edenlive.py +++ b/src/diffusers/models/transformers/transformer_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 The EdenLive Authors and The HuggingFace Team. All rights reserved. +# 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. diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 90f792b5f946..ec5e94d32ddf 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -274,7 +274,7 @@ ] _import_structure["diffusion_gemma"] = ["DiffusionGemmaPipeline", "DiffusionGemmaPipelineOutput"] _import_structure["dreamlite"] = ["DreamLitePipeline", "DreamLiteMobilePipeline", "DreamLitePipelineOutput"] - _import_structure["edenlive"] = ["EdenLivePipeline", "EdenLivePipelineOutput"] + _import_structure["edenlive"] = ["EdenLivePipeline", "EdenLivePipelineOutput", "EdenLiveStreamState"] _import_structure["easyanimate"] = [ "EasyAnimatePipeline", "EasyAnimateInpaintPipeline", @@ -725,7 +725,7 @@ EasyAnimateInpaintPipeline, EasyAnimatePipeline, ) - from .edenlive import EdenLivePipeline, EdenLivePipelineOutput + from .edenlive import EdenLivePipeline, EdenLivePipelineOutput, EdenLiveStreamState 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 index 13360a497a1d..72a32b470383 100644 --- a/src/diffusers/pipelines/edenlive/__init__.py +++ b/src/diffusers/pipelines/edenlive/__init__.py @@ -1,3 +1,5 @@ +# Copyright 2026 The VAST-AI Team. All rights reserved. + from typing import TYPE_CHECKING from ...utils import ( @@ -21,7 +23,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["pipeline_edenlive"] = ["EdenLivePipeline"] + _import_structure["pipeline_edenlive"] = ["EdenLivePipeline", "EdenLiveStreamState"] _import_structure["pipeline_output"] = ["EdenLivePipelineOutput"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -31,7 +33,7 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * else: - from .pipeline_edenlive import EdenLivePipeline + from .pipeline_edenlive import EdenLivePipeline, EdenLiveStreamState from .pipeline_output import EdenLivePipelineOutput else: import sys diff --git a/src/diffusers/pipelines/edenlive/loading_utils.py b/src/diffusers/pipelines/edenlive/loading_utils.py index e61b4772cc85..9b6ba7f4b2d2 100644 --- a/src/diffusers/pipelines/edenlive/loading_utils.py +++ b/src/diffusers/pipelines/edenlive/loading_utils.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py index e758d24e9b9a..83af59dd042f 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. @@ -12,19 +12,136 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math +import weakref +from dataclasses import dataclass, field from typing import Callable import torch from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel from ...models import AutoencoderKLWan, EdenLiveProxyProjection, EdenLiveTransformer3DModel +from ...models.autoencoders.autoencoder_edenlive import EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder +from ...models.autoencoders.autoencoder_kl_wan import WanCausalConv3d +from ...models.autoencoders.vae import DiagonalGaussianDistribution +from ...models.transformers.transformer_edenlive import EdenLiveCausalState from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler +from ...utils import replace_example_docstring +from ...utils.accelerate_utils import apply_forward_hook from ..pipeline_utils import DiffusionPipeline from . import pipeline_utils as utils from .loading_utils import load_original_pipeline, normalize_saved_profile, read_checkpoint_config from .pipeline_output import EdenLivePipelineOutput +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", profile="5b-causal") + >>> pipe.to("cuda") + >>> image = load_image("reference.png") + >>> video = load_video("proxy-24fps.mp4") # At least 49 RGB frames sampled at 24 fps. + >>> prompt = "A character walks through snow." + >>> frames = pipe( + ... image=image, + ... video=video, + ... prompt=prompt, + ... num_frames=49, + ... generator=torch.Generator().manual_seed(42), + ... ).frames[0] + >>> export_to_video(frames, "edenlive.mp4", fps=24) + + >>> # Generate successive chunks from three sampled RGB images per call. + >>> proxy_paths = [ + ... ["proxy/chunk-0-0.png", "proxy/chunk-0-1.png", "proxy/chunk-0-2.png"], + ... ["proxy/chunk-1-0.png", "proxy/chunk-1-1.png", "proxy/chunk-1-2.png"], + ... ] + >>> prompts = [prompt, "The character approaches a campfire."] + >>> state = pipe.create_stream_state( + ... image=image, prompt=prompt, generator=torch.Generator().manual_seed(42), output_type="np" + ... ) + >>> try: + ... for index, (paths, current_prompt) in enumerate(zip(proxy_paths, prompts)): + ... chunk = pipe( + ... stream_state=state, + ... proxy_frames=[load_image(path) for path in paths], + ... prompt=current_prompt, + ... ).frames[0] + ... export_to_video(chunk, f"chunk-{index:03d}.mp4", fps=24) + ... finally: + ... pipe.close_stream(state) + ``` +""" + + +@dataclass +class _EdenLiveConditionEncoderCache: + feat_map: list[torch.Tensor | None] | None = None + + def reset(self): + self.feat_map = None + + +@apply_forward_hook +def _encode_condition_chunk(vae, pixels, cache): + if vae.use_slicing or vae.use_tiling: + raise ValueError("Streaming condition encoding does not support VAE slicing or tiling.") + if vae.config.patch_size is not None: + raise ValueError("Streaming 14B conditioning requires an unpatchified Wan VAE.") + first = cache.feat_map is None + if pixels.shape[2] < 1 or (pixels.shape[2] - int(first)) % 4: + raise ValueError("Condition encoding requires 1 + 4n frames initially and 4n frames thereafter.") + if first: + cache.feat_map = [None for module in vae.encoder.modules() if isinstance(module, WanCausalConv3d)] + encoded = [] + if first: + encoded.append(vae.encoder(pixels[:, :, :1], feat_cache=cache.feat_map, feat_idx=[0])) + pixels = pixels[:, :, 1:] + for chunk in pixels.split(4, dim=2): + if chunk.shape[2]: + encoded.append(vae.encoder(chunk, feat_cache=cache.feat_map, feat_idx=[0])) + moments = vae.quant_conv(torch.cat(encoded, dim=2)) + return DiagonalGaussianDistribution(moments).mode() + + +@dataclass +class EdenLiveStreamState: + """Mutable, single-video rollout owned by one pipeline. Create with `pipe.create_stream_state()`. + + The attention states own the only time cursor (`next_frame`). Encoder and decoder caches retain temporal history + independently of the attention window. Calls sharing a pipeline must be serial, including calls for different streams. + Close the state in a `finally` block to release its tensors. States are runtime objects and are not serialized. + """ + + causal_states: list[EdenLiveCausalState] + condition_encoder_cache: _EdenLiveConditionEncoderCache | None + decoder_cache: EdenLiveLightTAEDecodeCache | None + image: torch.Tensor | None + first_frame: torch.Tensor | None + image_embeds: torch.Tensor | None + encoded_prompt: str | None + prompt_embeds: torch.Tensor | None + negative_prompt_embeds: torch.Tensor | None + generator: torch.Generator | None + profile: str + height: int + width: int + guidance_scale: float + num_inference_steps: int + max_sequence_length: int + output_type: str + device: torch.device + dtype: torch.dtype + _pipeline: weakref.ReferenceType = field(repr=False) + _component_ids: tuple = field(repr=False) + valid: bool = True + closed: bool = False + + class EdenLivePipeline(DiffusionPipeline): """Bidirectional and causal video generation conditioned on a reference image, text, and an RGB proxy video. @@ -175,6 +292,268 @@ def from_pretrained(cls, pretrained_model_name_or_path, *, profile=None, **kwarg pipeline.register_modules(decoder=decoder) return pipeline + @torch.no_grad() + def create_stream_state( + self, + image, + prompt=None, + *, + prompt_embeds=None, + negative_prompt=None, + negative_prompt_embeds=None, + height=None, + width=None, + guidance_scale=1.0, + num_inference_steps=4, + generator=None, + max_sequence_length=512, + output_type="np", + ) -> EdenLiveStreamState: + """Initialize one causal rollout with fixed image, resolution, guidance, RNG and output format. + + Uses the corresponding offline call options. Generate each chunk with `self(stream_state=state, + proxy_frames=frames)`. RGB output requires LightTAE. Close the state in a `finally` block. + """ + if not self.is_causal or self.config.chunk_size != 3: + raise ValueError("Streaming requires a causal profile with chunk_size=3.") + if output_type != "latent" and not isinstance(self.decoder, EdenLiveLightTAEDecoder): + raise ValueError( + "Streaming RGB output requires EdenLiveLightTAEDecoder; otherwise use output_type='latent'." + ) + if self.model_size == "14b" and (self.vae.use_slicing or self.vae.use_tiling): + raise ValueError("Streaming 14B condition encoding does not support VAE slicing or tiling.") + height = (704 if self.model_size == "5b" else 480) if height is None else height + width = (1280 if self.model_size == "5b" else 832) if width is None else width + if not isinstance(height, int) or not isinstance(width, int): + raise ValueError("Height and width must be integers.") + if image is None: + raise ValueError("A reference image is required.") + 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 self.model_size == "14b" and self.vae.config.patch_size is not None: + raise ValueError("Streaming 14B conditioning requires an unpatchified Wan VAE.") + self._check_prompt(prompt, prompt_embeds, required=True) + self._check_prompt(negative_prompt, negative_prompt_embeds) + if not isinstance(max_sequence_length, int) or max_sequence_length < 1: + raise ValueError("`max_sequence_length` must be a positive integer.") + if not isinstance(num_inference_steps, int) or num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be a positive integer.") + if not math.isfinite(guidance_scale): + raise ValueError("`guidance_scale` must be finite.") + if generator is not None and not isinstance(generator, torch.Generator): + raise ValueError("Streaming requires one torch.Generator, not a generator list.") + device, dtype = self._execution_device, self.transformer.dtype + reference = self.proxy_video_processor.preprocess(image) + if ( + reference.ndim != 4 + or reference.shape[:2] != (1, 3) + or reference.numel() == 0 + or not torch.isfinite(reference).all() + or reference.min() < 0 + or reference.max() > 1 + ): + raise ValueError("The reference image must be one finite RGB image in [0, 1].") + self.scheduler.set_timesteps(num_inference_steps, device=device) + pixels = utils.prepare_reference_pixels(self, reference, height, width, 1, device) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + negative_prompt, + guidance_scale > 1, + prompt_embeds, + negative_prompt_embeds, + max_sequence_length, + device, + dtype, + ) + image_embeds = utils.encode_image_embeds(self, pixels, device, dtype) + first_frame = self.vae.encode(pixels.unsqueeze(2).to(self.vae.dtype)).latent_dist.mode() + mean = first_frame.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / first_frame.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + first_frame = ((first_frame - mean) * inverse_std).to(dtype) + states = [ + self.transformer.create_causal_state( + batch_size=1, + height=first_frame.shape[3], + width=first_frame.shape[4], + chunk_size=3, + max_context_frames=self.config.max_context_frames, + sink_frames=self.config.sink_frames, + device=device, + dtype=dtype, + ) + for _ in range(2 if guidance_scale > 1 else 1) + ] + if generator is None: + generator = torch.Generator(device="cpu") + generator.seed() + return EdenLiveStreamState( + causal_states=states, + condition_encoder_cache=_EdenLiveConditionEncoderCache() if self.model_size == "14b" else None, + decoder_cache=EdenLiveLightTAEDecodeCache() if output_type != "latent" else None, + image=pixels if self.model_size == "14b" else None, + first_frame=first_frame, + image_embeds=image_embeds, + encoded_prompt=prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + generator=generator, + profile=self.config.profile, + height=height, + width=width, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + max_sequence_length=max_sequence_length, + output_type=output_type, + device=device, + dtype=dtype, + _pipeline=weakref.ref(self), + _component_ids=tuple(id(component) for component in self.components.values()), + ) + + def close_stream(self, state): + """Release a stream's tensors and model hooks. Repeated closing has no effect.""" + self._check_state(state) + if state.closed: + return + for causal_state in state.causal_states: + causal_state.reset() + state.causal_states.clear() + for cache in (state.condition_encoder_cache, state.decoder_cache): + if cache is not None: + cache.reset() + state.condition_encoder_cache = state.decoder_cache = None + state.image = state.first_frame = state.image_embeds = None + state.prompt_embeds = state.negative_prompt_embeds = None + state.encoded_prompt = state.generator = None + state.valid, state.closed = False, True + self.maybe_free_model_hooks() + + def _check_prompt(self, prompt, embeds, *, required=False): + if prompt is not None and embeds is not None: + raise ValueError("Pass either `prompt` or `prompt_embeds`, not both.") + if prompt is not None and not isinstance(prompt, str): + raise ValueError("Streaming requires a single string prompt.") + if embeds is not None and ( + not isinstance(embeds, torch.Tensor) + or embeds.ndim != 3 + or embeds.shape[0] != 1 + or embeds.shape[1] < 1 + or embeds.shape[2] != self.transformer.config.text_dim + or not torch.is_floating_point(embeds) + or not torch.isfinite(embeds).all() + ): + raise ValueError("Prompt embeddings must be finite floating-point tensors of shape (1, tokens, text_dim).") + if required and prompt is None and embeds is None: + raise ValueError("Provide `prompt` or `prompt_embeds`.") + + def _prepare_condition_chunk(self, state): + if state.condition_encoder_cache is None: + return state.first_frame + first = state.causal_states[0].next_frame == 0 + pixels = state.first_frame.new_zeros(1, 3, 8 if first else 12, state.height, state.width) + if first: + pixels = torch.cat([state.image.unsqueeze(2), pixels], dim=2) + condition_chunk = _encode_condition_chunk(self.vae, pixels.to(self.vae.dtype), state.condition_encoder_cache) + mean = condition_chunk.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / condition_chunk.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) + condition_chunk = ((condition_chunk - mean) * inverse_std).to(state.dtype) + mask = condition_chunk.new_zeros(1, self.vae_scale_factor_temporal, *condition_chunk.shape[2:]) + if first: + condition_chunk[:, :, :1] = state.first_frame + mask[:, :, 0] = 1 + state.image = None + return torch.cat([mask, condition_chunk], dim=1) + + def _update_prompt_embeds(self, state, prompt, prompt_embeds): + if prompt_embeds is not None: + state.prompt_embeds = prompt_embeds.to(device=state.device, dtype=state.dtype) + state.encoded_prompt = None + elif prompt is not None and prompt != state.encoded_prompt: + state.prompt_embeds, _ = self.encode_prompt( + prompt, + do_classifier_free_guidance=False, + max_sequence_length=state.max_sequence_length, + device=state.device, + dtype=state.dtype, + ) + state.encoded_prompt = prompt + + def _check_state(self, state): + if not isinstance(state, EdenLiveStreamState) or state._pipeline() is not self: + raise ValueError("The stream must belong to this pipeline.") + + def _stream_step(self, state, proxy_frames, *, prompt=None, prompt_embeds=None, latents=None, return_dict=True): + self._check_state(state) + if state.closed or not state.valid: + raise ValueError("This stream is closed or invalid; create a new stream state.") + if tuple(id(component) for component in self.components.values()) != state._component_ids: + raise ValueError("Pipeline components must remain fixed throughout a stream.") + if ( + self.config.profile != state.profile + or self._execution_device != state.device + or self.transformer.dtype != state.dtype + ): + raise ValueError("A stream requires a fixed model, device, and dtype.") + self._check_prompt(prompt, prompt_embeds) + shape = ( + 1, + self.vae.config.z_dim, + 3, + state.height // self.vae_scale_factor_spatial, + state.width // self.vae_scale_factor_spatial, + ) + if latents is not None and ( + not isinstance(latents, torch.Tensor) + or tuple(latents.shape) != shape + or not torch.is_floating_point(latents) + or not torch.isfinite(latents).all() + ): + raise ValueError(f"Initial latents must be a finite floating-point tensor with shape {shape}.") + proxy_video = self.proxy_video_processor.preprocess_proxy_chunk(proxy_frames) + self._guidance_scale = state.guidance_scale + self._num_timesteps = state.num_inference_steps + try: + self._update_prompt_embeds(state, prompt, prompt_embeds) + proxy_features = utils.encode_proxy_features(self, proxy_video, 1, state.device) + condition_chunk = self._prepare_condition_chunk(state) + latents = ( + utils.sample_noise(self, shape, state.generator, state.device, state.dtype) + if latents is None + else latents.to(device=state.device, dtype=state.dtype).clone() + ) + latents, state.prompt_embeds, state.negative_prompt_embeds = self._denoise_chunk( + latents, + condition_chunk, + state.first_frame, + proxy_features, + state.image_embeds, + state.prompt_embeds, + state.negative_prompt_embeds, + state.causal_states, + state.num_inference_steps, + state.generator, + None, + None, + [], + 0, + condition_start=0, + ) + frames = utils.decode_latents(self, latents, state.output_type, decoder_cache=state.decoder_cache) + return EdenLivePipelineOutput(frames=frames) if return_dict else (frames,) + except Exception: + state.valid = False + raise + def encode_prompt( self, prompt=None, @@ -263,6 +642,7 @@ 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, @@ -280,18 +660,31 @@ def __call__( negative_prompt_embeds: torch.Tensor | None = None, text_timeline: dict | list[dict] | None = None, chunk_prompts: list[str] | None = None, - max_sequence_length: int = 512, - output_type: str = "np", + max_sequence_length: int | None = None, + output_type: str | None = None, return_dict: bool = True, attention_kwargs: dict | None = None, callback_on_step_end: Callable | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end_tensor_inputs: list[str] | None = None, + *, + stream_state: EdenLiveStreamState | None = None, + proxy_frames=None, ) -> EdenLivePipelineOutput | tuple: - """Generate a video from a reference image and a full-rate RGB proxy video. + """Generate an offline video, or advance a stream using three already sampled RGB proxy frames. + + With `stream_state`, only `proxy_frames`, `prompt`, `prompt_embeds`, `latents`, and `return_dict` may be + supplied. All fixed settings come from the state. Calls sharing a pipeline must run serially. Args: + stream_state (`EdenLiveStreamState`, *optional*): + State from `create_stream_state`. Returns one chunk: 9 RGB frames initially and 12 thereafter, or + 3 latent frames. Invalid inputs leave the state usable; model execution failures invalidate it. + proxy_frames (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`, *optional*): + Exactly three already sampled RGB frames, required with `stream_state`. No temporal subsampling. + Tensor layout is `(3, 3, H, W)`; NumPy layout is `(3, H, W, 3)`. Float values must be in [0, 1]. prompt (`str` or `list[str]`, *optional*): Text describing the output video. Supply this or `prompt_embeds` unless using temporal prompts. + In a stream, new text takes effect in this chunk; None retains the current text. 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. @@ -349,10 +742,50 @@ def __call__( 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"`. """ + if stream_state is not None: + fixed_options = { + "image": image, + "video": video, + "height": height, + "width": width, + "num_frames": num_frames, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "negative_prompt": negative_prompt, + "negative_prompt_embeds": negative_prompt_embeds, + "generator": generator, + "text_timeline": text_timeline, + "chunk_prompts": chunk_prompts, + "max_sequence_length": max_sequence_length, + "output_type": output_type, + "attention_kwargs": attention_kwargs, + "callback_on_step_end": callback_on_step_end, + "callback_on_step_end_tensor_inputs": callback_on_step_end_tensor_inputs, + } + supplied = [name for name, value in fixed_options.items() if value is not None] + if supplied: + raise ValueError(f"Streaming calls take fixed settings from the state; omit: {', '.join(supplied)}.") + return self._stream_step( + stream_state, + proxy_frames, + prompt=prompt, + prompt_embeds=prompt_embeds, + latents=latents, + return_dict=return_dict, + ) + if proxy_frames is not None: + raise ValueError("`proxy_frames` requires `stream_state`; use `video` for offline generation.") + max_sequence_length = 512 if max_sequence_length is None else max_sequence_length + output_type = "np" if output_type is None else output_type + callback_on_step_end_tensor_inputs = ( + ["latents"] if callback_on_step_end_tensor_inputs is None else callback_on_step_end_tensor_inputs + ) if text_timeline is not None and (self.is_causal or self.model_size != "5b"): raise ValueError("`text_timeline` requires the 5B bidirectional profile.") if chunk_prompts is not None and not self.is_causal: @@ -576,6 +1009,8 @@ def _denoise_chunk( callback_on_step_end, callback_on_step_end_tensor_inputs, callback_step_offset=0, + *, + condition_start=None, ): start = states[0].next_frame if start == 0: @@ -583,7 +1018,7 @@ def _denoise_chunk( self.scheduler.set_timesteps(num_inference_steps, device=latents.device) for step_index, timestep in enumerate(self.scheduler.timesteps): model_input, model_timestep = utils.prepare_model_input( - self, latents, condition, first_frame, timestep, start, causal=True + self, latents, condition, first_frame, timestep, start, causal=True, condition_start=condition_start ) prediction = self.transformer( model_input, @@ -636,7 +1071,14 @@ def _denoise_chunk( 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 = utils.prepare_model_input( - self, latents, condition, first_frame, latents.new_zeros(()), start, causal=True + self, + latents, + condition, + first_frame, + latents.new_zeros(()), + start, + causal=True, + condition_start=condition_start, ) if self.model_size == "14b": # A clean cache commit uses one zero timestep per sample, not per token. diff --git a/src/diffusers/pipelines/edenlive/pipeline_output.py b/src/diffusers/pipelines/edenlive/pipeline_output.py index 3d5956179460..fe61e1d3184b 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_output.py +++ b/src/diffusers/pipelines/edenlive/pipeline_output.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. diff --git a/src/diffusers/pipelines/edenlive/pipeline_utils.py b/src/diffusers/pipelines/edenlive/pipeline_utils.py index d03db5a66407..c2948f7b4a36 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_utils.py +++ b/src/diffusers/pipelines/edenlive/pipeline_utils.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. @@ -217,7 +217,7 @@ def sample_noise(pipe, shape, generator, device, dtype): return noise -def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batch_size, device, dtype): +def prepare_reference_pixels(pipe, image, height, width, batch_size, device): pixels = pipe.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)) @@ -228,8 +228,12 @@ def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batc pixels = pixels.to(device=device) if pixels.shape[0] == 1: pixels = pixels.expand(batch_size, -1, -1, -1) - if pixels.shape[0] != batch_size or proxy_video.shape[0] not in (1, batch_size): - raise ValueError("Image and proxy-video batch sizes must be one or match the prompt batch.") + if pixels.shape[0] != batch_size: + raise ValueError("Image batch size must be one or match the prompt batch.") + return pixels + + +def encode_image_embeds(pipe, pixels, device, dtype): clip_embeds = None if pipe.model_size == "14b": if pipe.image_encoder is None or pipe.image_processor is None: @@ -243,20 +247,10 @@ def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batc .hidden_states[-2] .to(dtype) ) - 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 = pipe.vae.encode(condition_video.to(pipe.vae.dtype)).latent_dist.mode() - mean = condition.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) - inverse_std = 1 / condition.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) - condition = ((condition - mean) * inverse_std).to(dtype) - first_frame = condition[:, :, :1] - if pipe.model_size == "14b": - mask = condition.new_zeros(batch_size, pipe.vae_scale_factor_temporal, *condition.shape[2:]) - mask[:, :, 0] = 1 - condition = torch.cat([mask, condition], dim=1) + return clip_embeds + + +def encode_proxy_features(pipe, proxy_video, batch_size, device): proxy_video = proxy_video.expand(batch_size, -1, -1, -1, -1) proxy_height, proxy_width = pipe.config.proxy_image_size proxy_features = [] @@ -279,17 +273,44 @@ def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batc ): projected = pipe.proxy_projection(hidden if pipe.config.proxy_projection_autocast else hidden.float()) proxy_features.append(projected.float().permute(0, 3, 1, 2)) - return condition, first_frame, clip_embeds, torch.stack(proxy_features, dim=2) + return torch.stack(proxy_features, dim=2) -def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0, causal=False): +def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batch_size, device, dtype): + pixels = prepare_reference_pixels(pipe, image, height, width, batch_size, device) + 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 = encode_image_embeds(pipe, pixels, device, dtype) + if pipe.model_size == "14b": + 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 = pipe.vae.encode(condition_video.to(pipe.vae.dtype)).latent_dist.mode() + mean = condition.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) + inverse_std = 1 / condition.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) + condition = ((condition - mean) * inverse_std).to(dtype) + first_frame = condition[:, :, :1] + if pipe.model_size == "14b": + mask = condition.new_zeros(batch_size, pipe.vae_scale_factor_temporal, *condition.shape[2:]) + mask[:, :, 0] = 1 + condition = torch.cat([mask, condition], dim=1) + proxy_features = encode_proxy_features(pipe, proxy_video, batch_size, device) + return condition, first_frame, clip_embeds, proxy_features + + +def prepare_model_input( + pipe, latents, condition, first_frame, timestep, start=0, causal=False, *, condition_start=None +): model_input = latents pin_first = start == 0 and (causal or pipe.model_size == "5b") if pin_first: model_input = model_input.clone() model_input[:, :, :1] = first_frame if pipe.model_size == "14b": - model_input = torch.cat([model_input, condition[:, :, start : start + latents.shape[2]]], dim=1) + offset = start if condition_start is None else condition_start + model_input = torch.cat([model_input, condition[:, :, offset : offset + latents.shape[2]]], dim=1) timestep = timestep.expand(latents.shape[0]) if pin_first: patch = pipe.transformer.config.patch_size @@ -299,7 +320,7 @@ def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0 return model_input, timestep -def decode_latents(pipe, latents, output_type): +def decode_latents(pipe, latents, output_type, *, decoder_cache=None): if output_type == "latent": return latents latents = latents.to(pipe.vae.dtype) @@ -314,7 +335,8 @@ def decode_latents(pipe, latents, output_type): if dtype is None: parameter = next(decoder.parameters(), None) dtype = parameter.dtype if parameter is not None else raw_latents.dtype - video = decoder.decode(raw_latents.to(dtype=dtype), return_dict=False)[0] + decode_kwargs = {} if decoder_cache is None else {"cache": decoder_cache} + video = decoder.decode(raw_latents.to(dtype=dtype), return_dict=False, **decode_kwargs)[0] return pipe.video_processor.postprocess_video(video, output_type=output_type) diff --git a/src/diffusers/pipelines/edenlive/video_processor.py b/src/diffusers/pipelines/edenlive/video_processor.py index f5901f4dc0e6..aa37a23dd6c8 100644 --- a/src/diffusers/pipelines/edenlive/video_processor.py +++ b/src/diffusers/pipelines/edenlive/video_processor.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. @@ -27,6 +27,37 @@ class EdenLiveVideoProcessor(VideoProcessor): 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_chunk(self, proxy_frames): + """Convert three already sampled RGB condition frames to BCTHW without temporal subsampling.""" + if isinstance(proxy_frames, (np.ndarray, torch.Tensor)): + if proxy_frames.ndim == 5 and proxy_frames.shape[0] == 1: + proxy_frames = proxy_frames[0] + if proxy_frames.ndim != 4: + raise ValueError("Pass a single proxy chunk with three RGB frames.") + elif not isinstance(proxy_frames, list): + raise ValueError("Pass three RGB proxy frames as a list, NumPy array, or tensor.") + if len(proxy_frames) != 3: + raise ValueError("Each streaming chunk requires exactly three already sampled proxy frames.") + frame_type = type(proxy_frames[0]) + if not all(isinstance(frame, frame_type) for frame in proxy_frames): + raise ValueError("All proxy frames must use the same image format.") + for frame in proxy_frames: + if isinstance(frame, PIL.Image.Image): + if frame.mode != "RGB": + raise ValueError("Proxy frames must be RGB images.") + continue + if not isinstance(frame, (np.ndarray, torch.Tensor)) or frame.ndim != 3: + raise ValueError("Each proxy frame must be an RGB image.") + channel_axis = 0 if isinstance(frame, torch.Tensor) else -1 + if frame.shape[channel_axis] != 3 or any(size == 0 for size in frame.shape): + raise ValueError("Proxy frames must have three RGB channels.") + values = torch.as_tensor(frame) + if not torch.is_floating_point(values) or not torch.isfinite(values).all(): + raise ValueError("NumPy and tensor proxy frames must contain finite floating-point RGB in [0, 1].") + if values.min() < 0 or values.max() > 1: + raise ValueError("Proxy frames must contain RGB values in [0, 1].") + return self.preprocess_video([proxy_frames]) + 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) diff --git a/src/diffusers/schedulers/scheduling_edenlive.py b/src/diffusers/schedulers/scheduling_edenlive.py index 4d16e4079818..1a1a8e6f1188 100644 --- a/src/diffusers/schedulers/scheduling_edenlive.py +++ b/src/diffusers/schedulers/scheduling_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 7ec38ee07677..ae4e321847b7 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1339,6 +1339,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class EdenLiveLightTAEDecodeCache(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"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 3bf80360258f..5377117b62a4 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1967,6 +1967,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class EdenLiveStreamState(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/pipelines/edenlive/__init__.py b/tests/pipelines/edenlive/__init__.py index e69de29bb2d1..b8ca19521280 100644 --- a/tests/pipelines/edenlive/__init__.py +++ 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_edenlive.py b/tests/pipelines/edenlive/test_edenlive.py index 0f8b128c3a9f..5ef25788d0fb 100644 --- a/tests/pipelines/edenlive/test_edenlive.py +++ b/tests/pipelines/edenlive/test_edenlive.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# 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. @@ -1497,3 +1497,311 @@ def test_invalid_schedule(self): ): with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): scheduler.set_timesteps(**kwargs) + + +class TestEdenLiveStreaming(EdenLiveChunkCausalPipelineTesterConfig): + def setup_stream(self, profile="5b", output_type="latent", guidance=1.0): + self.profile = profile + components = self.get_dummy_components() + components["decoder"] = EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) + pipe = self.get_pipeline(**components).to(torch_device) + inputs = self.get_dummy_inputs() + options = { + "image": inputs["image"], + "prompt": "snow", + "height": 16, + "width": 16, + "num_inference_steps": 2, + "max_sequence_length": 16, + "output_type": output_type, + "guidance_scale": guidance, + } + state = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) + return pipe, state, inputs["video"][:3], options + + @pytest.mark.parametrize("profile", ["5b", "14b"]) + @pytest.mark.parametrize("guidance", [1.0, 2.0]) + def test_eight_chunks_prompt_and_bounded_history(self, profile, guidance): + from unittest.mock import patch + + pipe, state, proxy, _ = self.setup_stream(profile, guidance=guidance) + cache = state.condition_encoder_cache + with patch.object(pipe, "encode_prompt", wraps=pipe.encode_prompt) as encode: + capacities = [] + for index in range(8): + result = pipe(stream_state=state, proxy_frames=proxy, prompt="snow" if index < 4 else "fire") + assert result.frames.shape == (1, 4, 3, 2, 2) + assert torch.isfinite(result.frames).all() + assert all(s.next_frame == (index + 1) * 3 for s in state.causal_states) + capacities.append( + sum(layer.key.numel() + layer.value.numel() for s in state.causal_states for layer in s.layers) + ) + assert encode.call_count == 1 + assert capacities[-1] == capacities[-2] + pipe(stream_state=state, proxy_frames=proxy, prompt_embeds=state.prompt_embeds) + assert state.encoded_prompt is None + pipe(stream_state=state, proxy_frames=proxy, prompt="fire") + assert encode.call_count == 2 + assert state.condition_encoder_cache is cache + assert state.image is None + pipe.close_stream(state) + pipe.close_stream(state) + assert state.closed and not state.valid and not state.causal_states + with pytest.raises(ValueError): + pipe(stream_state=state, proxy_frames=proxy) + + @pytest.mark.parametrize("profile", ["5b", "14b"]) + def test_interleaved_streams_and_dynamic_prompt(self, profile): + pipe, first, proxy, options = self.setup_stream(profile) + second = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) + for _ in range(2): + a = pipe(stream_state=first, proxy_frames=proxy).frames + b = pipe(stream_state=second, proxy_frames=proxy).frames + torch.testing.assert_close(a, b, rtol=1e-4, atol=1e-5) + a = pipe(stream_state=first, proxy_frames=proxy, prompt="fire").frames + b = pipe(stream_state=second, proxy_frames=proxy, prompt="ocean").frames + assert (a - b).abs().max() > 1e-6 + pipe.close_stream(first) + assert pipe(stream_state=second, proxy_frames=proxy).frames.shape[2] == 3 + pipe.close_stream(second) + + @pytest.mark.parametrize("profile", ["5b", "14b"]) + def test_rgb_frames_and_invalid_input_then_execution_failure(self, profile): + from unittest.mock import patch + + pipe, state, proxy, _ = self.setup_stream(profile, output_type="pt") + rng = state.generator.get_state().clone() + for kwargs in ( + {"proxy_frames": proxy[:2]}, + {"proxy_frames": proxy, "latents": torch.zeros(1)}, + {"proxy_frames": proxy, "prompt_embeds": torch.zeros(2, 4, 32)}, + ): + with pytest.raises(ValueError): + pipe(stream_state=state, **kwargs) + assert state.valid and state.causal_states[0].next_frame == 0 + assert torch.equal(rng, state.generator.get_state()) + assert pipe(stream_state=state, proxy_frames=proxy).frames.shape == (1, 9, 3, 16, 16) + assert pipe(stream_state=state, proxy_frames=proxy).frames.shape == (1, 12, 3, 16, 16) + with patch.object(pipe.proxy_encoder, "forward", side_effect=RuntimeError("injected failure")): + with pytest.raises(RuntimeError, match="injected failure"): + pipe(stream_state=state, proxy_frames=proxy) + assert not state.valid + with pytest.raises(ValueError, match="invalid"): + pipe(stream_state=state, proxy_frames=proxy) + pipe.close_stream(state) + + @pytest.mark.parametrize("profile", ["5b", "14b"]) + def test_offline_stream_numerics_with_fixed_noise(self, profile): + from unittest.mock import patch + + from diffusers.pipelines.edenlive import pipeline_utils + + pipe, state, proxy, options = self.setup_stream(profile, guidance=2.0) + full_proxy = np.concatenate( + [proxy[:1].repeat(4, axis=0), proxy[1:2].repeat(4, axis=0), proxy[2:3].repeat(13, axis=0)] + ) + + def noise(pipe, shape, generator, device, dtype): + return torch.full(shape, 0.25, device=device, dtype=dtype) + + with patch.object(pipeline_utils, "sample_noise", side_effect=noise): + expected = pipe(**options, video=full_proxy, num_frames=21).frames + chunks = [ + pipe(stream_state=state, proxy_frames=full_proxy[:9:4]).frames, + pipe(stream_state=state, proxy_frames=full_proxy[12:21:4]).frames, + ] + torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=1e-4, atol=1e-5) + pipe.close_stream(state) + + @pytest.mark.parametrize("patch_size", [1, 2]) + @torch.no_grad() + def test_lighttae_continuity_reset_and_isolation(self, patch_size): + from diffusers import EdenLiveLightTAEDecodeCache + + decoder = EdenLiveLightTAEDecoder( + latent_channels=4, patch_size=patch_size, block_out_channels=(4, 4, 4, 4) + ).to(torch_device) + latents = torch.randn(1, 4, 9, 2, 2, device=torch_device) + expected = decoder.decode(latents).sample + cache = EdenLiveLightTAEDecodeCache() + outputs = [] + for chunk in latents.split(3, dim=2): + outputs.append(decoder.decode(chunk, cache=cache).sample) + decoder.decode(chunk + 1, cache=EdenLiveLightTAEDecodeCache()) + assert [x.shape[2] for x in outputs] == [9, 12, 12] + torch.testing.assert_close(torch.cat(outputs, dim=2), expected, rtol=1e-4, atol=1e-5) + cache.reset() + torch.testing.assert_close(decoder.decode(latents, cache=cache).sample, expected, rtol=1e-4, atol=1e-5) + + def test_proxy_chunk_is_not_subsampled(self): + from diffusers.pipelines.edenlive.video_processor import EdenLiveVideoProcessor + + processor = EdenLiveVideoProcessor() + frames = np.stack([np.full((16, 16, 3), x, dtype=np.float32) for x in (0.0, 0.5, 1.0)]) + result = processor.preprocess_proxy_chunk(frames) + assert result.shape == (1, 3, 3, 16, 16) + torch.testing.assert_close(result[0, 0, :, 0, 0], torch.tensor([0.0, 0.5, 1.0])) + + def test_initialization_validation_and_owner_lifetime(self): + import gc + import weakref + from unittest.mock import patch + + pipe, state, proxy, options = self.setup_stream() + for image in (np.full((16, 16, 3), np.nan, dtype=np.float32), np.ones((16, 16, 3), dtype=np.float32) * 2): + with patch.object(pipe.vae, "encode", side_effect=AssertionError("must validate before encoding")): + with pytest.raises(ValueError, match="reference image"): + pipe.create_stream_state(**{**options, "image": image}) + with patch.object(pipe.decoder, "decode", side_effect=AssertionError("latent output must bypass decode")): + pipe(stream_state=state, proxy_frames=proxy) + pipe.close_stream(state) + owner = weakref.ref(pipe) + del pipe + gc.collect() + assert owner() is None + + def test_rejects_foreign_state_and_component_replacement(self): + pipe, state, proxy, _ = self.setup_stream() + other = self.get_pipeline().to(torch_device) + with pytest.raises(ValueError, match="belong"): + other(stream_state=state, proxy_frames=proxy) + original = pipe.scheduler + pipe.scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(default_timesteps=[1000, 500]) + with pytest.raises(ValueError, match="components"): + pipe(stream_state=state, proxy_frames=proxy) + assert state.valid and state.causal_states[0].next_frame == 0 + pipe.scheduler = original + pipe.close_stream(state) + + @torch.no_grad() + def test_condition_encoding_continuity_isolation_and_root_hook(self): + from unittest.mock import Mock, patch + + from diffusers.pipelines.edenlive.pipeline_edenlive import ( + _EdenLiveConditionEncoderCache, + _encode_condition_chunk, + ) + + self.profile = "14b" + vae = self.get_dummy_components()["vae"].to(torch_device).eval() + video = randn_tensor((1, 3, 33, 16, 16), generator=torch.Generator().manual_seed(12), device=torch_device) + expected = vae.encode(video).latent_dist.mode() + internal_map, internal_index = vae._enc_feat_map, vae._enc_conv_idx + map_values, index_values = list(internal_map), list(internal_index) + cache, other = _EdenLiveConditionEncoderCache(), _EdenLiveConditionEncoderCache() + hook = Mock() + with patch.object(vae, "_hf_hook", hook, create=True): + first = _encode_condition_chunk(vae, video[:, :, :9], cache) + cache_map = cache.feat_map + _encode_condition_chunk(vae, video[:, :, :9] + 1, other) + middle = _encode_condition_chunk(vae, video[:, :, 9:21], cache) + tail = _encode_condition_chunk(vae, video[:, :, 21:], cache) + assert cache.feat_map is cache_map + assert hook.pre_forward.call_count == 4 + hook.pre_forward.assert_called_with(vae) + torch.testing.assert_close(torch.cat([first, middle, tail], dim=2), expected, rtol=1e-4, atol=1e-5) + assert vae._enc_feat_map is internal_map and vae._enc_conv_idx is internal_index + assert vae._enc_feat_map == map_values and vae._enc_conv_idx == index_values + cache.reset() + seed = _encode_condition_chunk(vae, video[:, :, :1], cache) + rest = _encode_condition_chunk(vae, video[:, :, 1:], cache) + torch.testing.assert_close(torch.cat([seed, rest], dim=2), expected, rtol=1e-4, atol=1e-5) + + @torch.no_grad() + def test_condition_encoding_rejects_invalid_chunks_and_memory_modes(self): + from diffusers.pipelines.edenlive.pipeline_edenlive import ( + _EdenLiveConditionEncoderCache, + _encode_condition_chunk, + ) + + self.profile = "14b" + vae = self.get_dummy_components()["vae"].to(torch_device).eval() + video = torch.zeros(1, 3, 9, 16, 16, device=torch_device) + cache = _EdenLiveConditionEncoderCache() + with pytest.raises(ValueError): + _encode_condition_chunk(vae, video[:, :, :4], cache) + assert cache.feat_map is None + _encode_condition_chunk(vae, video[:, :, :1], cache) + with pytest.raises(ValueError): + _encode_condition_chunk(vae, video[:, :, :1], cache) + for mode in ("slicing", "tiling"): + getattr(vae, f"enable_{mode}")() + with pytest.raises(ValueError): + _encode_condition_chunk(vae, video, _EdenLiveConditionEncoderCache()) + getattr(vae, f"disable_{mode}")() + + def test_stream_dispatch_conflicts_are_preflight_and_tuple_output(self): + from unittest.mock import patch + + pipe, state, proxy, options = self.setup_stream("14b", output_type="pt") + rng = state.generator.get_state().clone() + fixed = { + "image": options["image"], + "video": proxy, + "height": 16, + "width": 16, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "negative_prompt": "", + "negative_prompt_embeds": state.prompt_embeds, + "generator": state.generator, + "text_timeline": {}, + "chunk_prompts": [], + "max_sequence_length": 512, + "output_type": "np", + "attention_kwargs": {}, + "callback_on_step_end": lambda *args: {}, + "callback_on_step_end_tensor_inputs": ["latents"], + } + with patch.object(pipe.proxy_encoder, "forward", side_effect=AssertionError("must validate before encoding")): + for name, value in fixed.items(): + with pytest.raises(ValueError, match=name): + pipe(stream_state=state, proxy_frames=proxy, **{name: value}) + with pytest.raises(ValueError, match="stream_state"): + pipe(proxy_frames=proxy) + with pytest.raises(ValueError): + pipe(stream_state=state) + assert state.valid and state.causal_states[0].next_frame == 0 + assert state.condition_encoder_cache.feat_map is None and not state.decoder_cache.dec_state + assert torch.equal(rng, state.generator.get_state()) + result = pipe(stream_state=state, proxy_frames=proxy, return_dict=False) + assert isinstance(result, tuple) and len(result) == 1 + assert result[0].shape == (1, 9, 3, 16, 16) + pipe.close_stream(state) + + def test_offline_effective_defaults(self): + from unittest.mock import patch + + pipe, state, _, options = self.setup_stream() + pipe.close_stream(state) + inputs = self.get_dummy_inputs() + inputs.pop("output_type") + inputs.pop("max_sequence_length") + with patch.object(pipe, "encode_prompt", wraps=pipe.encode_prompt) as encode: + actual = pipe(**inputs).frames + assert encode.call_args.args[-1] == 512 + inputs = self.get_dummy_inputs() + inputs.update(max_sequence_length=512, output_type="np", callback_on_step_end_tensor_inputs=["latents"]) + expected = pipe(**inputs).frames + assert isinstance(actual, np.ndarray) + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5) + + @pytest.mark.parametrize("profile", ["5b", "14b"]) + def test_model_cpu_offload_matches_device_placement(self, profile): + if torch.device(torch_device).type == "cpu": + pytest.skip("Model CPU offload requires an accelerator.") + pipe, state, proxy, options = self.setup_stream(profile, output_type="pt", guidance=2.0) + expected = [pipe(stream_state=state, proxy_frames=proxy).frames.cpu() for _ in range(2)] + pipe.close_stream(state) + pipe.enable_model_cpu_offload(device=torch_device) + state = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) + cache = state.condition_encoder_cache + try: + for chunk in expected: + actual = pipe(stream_state=state, proxy_frames=proxy).frames.cpu() + torch.testing.assert_close(actual, chunk, rtol=1e-4, atol=1e-5) + assert state.condition_encoder_cache is cache + finally: + pipe.close_stream(state) + assert pipe.vae.device.type == "cpu" and pipe.transformer.device.type == "cpu" From 6f56be292f154cad02f51f6d2b683be8468b70ef Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Fri, 16 Oct 2026 23:50:11 +0800 Subject: [PATCH 3/6] update readme --- docs/source/en/api/pipelines/edenlive.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/source/en/api/pipelines/edenlive.md b/docs/source/en/api/pipelines/edenlive.md index 8036ed90f86d..26d07ec9c163 100644 --- a/docs/source/en/api/pipelines/edenlive.md +++ b/docs/source/en/api/pipelines/edenlive.md @@ -12,8 +12,7 @@ specific language governing permissions and limitations under the License. # EdenLive -[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) generates video from a reference image, prompt, and -RGB proxy video. [`EdenLivePipeline`] supports bidirectional generation and causal generation with persistent history. +[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) renders a proxy or reference video stream, a target first frame, and a text prompt into refined or stylized video. References can come from 2D/3D scenes, existing engines, or recordings. [`EdenLivePipeline`] supports bidirectional generation and causal generation with persistent history. ## Basic usage From c9b2d0e6e92b104990e122cdf11279b85045c3b8 Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Sun, 18 Oct 2026 18:55:17 +0800 Subject: [PATCH 4/6] fix: align EdenLive defaults and cache commits --- examples/edenlive/inference_edenlive.py | 3 ++- .../edenlive/inference_edenlive_streaming.py | 3 ++- .../autoencoders/autoencoder_edenlive.py | 4 +-- .../pipelines/edenlive/loading_utils.py | 4 +-- .../pipelines/edenlive/pipeline_edenlive.py | 23 ++++++++-------- tests/pipelines/edenlive/test_edenlive.py | 26 +++++++++---------- 6 files changed, 33 insertions(+), 30 deletions(-) diff --git a/examples/edenlive/inference_edenlive.py b/examples/edenlive/inference_edenlive.py index f0b4a7a74e11..d630c1630c5b 100644 --- a/examples/edenlive/inference_edenlive.py +++ b/examples/edenlive/inference_edenlive.py @@ -122,7 +122,7 @@ def parse_args(): 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, default=0) + 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("--no-fsdp", action="store_true", help="Use replicated DiT weights with context parallelism.") parser.add_argument("--local-files-only", action="store_true") @@ -162,6 +162,7 @@ def main(): 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.profile == "14b-causal" else 0) video, video_info = load_video(args.video, args.fps, args.height, args.width, args.num_frames) inputs = { "prompt": args.prompt, diff --git a/examples/edenlive/inference_edenlive_streaming.py b/examples/edenlive/inference_edenlive_streaming.py index a48662eb0bfe..70e4be3c84d5 100644 --- a/examples/edenlive/inference_edenlive_streaming.py +++ b/examples/edenlive/inference_edenlive_streaming.py @@ -37,10 +37,11 @@ def main(): parser.add_argument("--height", type=int) parser.add_argument("--width", type=int) parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--fps", type=int, help="Playback rate; defaults to 16 for 14b-causal and 24 for 5b-causal.") parser.add_argument("--output", type=Path, required=True) parser.add_argument("--local-files-only", action="store_true") args = parser.parse_args() + args.fps = args.fps if args.fps is not None else (16 if args.profile == "14b-causal" else 24) if len(args.proxy_frames) % 3: parser.error("--proxy-frames must contain a multiple of three image paths.") if args.fps <= 0: diff --git a/src/diffusers/models/autoencoders/autoencoder_edenlive.py b/src/diffusers/models/autoencoders/autoencoder_edenlive.py index e0a4abaf26db..cba0b56f28a2 100644 --- a/src/diffusers/models/autoencoders/autoencoder_edenlive.py +++ b/src/diffusers/models/autoencoders/autoencoder_edenlive.py @@ -175,8 +175,8 @@ def decode( [`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 for a new video, or `T_latent * scale_factor_temporal` for an - existing cache, and spatial dimensions multiplied by `scale_factor_spatial`. Internal - chunks share temporal state, including a short final chunk. + existing cache, 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.") diff --git a/src/diffusers/pipelines/edenlive/loading_utils.py b/src/diffusers/pipelines/edenlive/loading_utils.py index 9b6ba7f4b2d2..e273e2b507ef 100644 --- a/src/diffusers/pipelines/edenlive/loading_utils.py +++ b/src/diffusers/pipelines/edenlive/loading_utils.py @@ -43,7 +43,7 @@ "5b-bidir": {"size": "5b", "causal": False, "transformer": "dit-5b-bidir", "resize": "bilinear"}, "14b-bidir": {"size": "14b", "causal": False, "transformer": "dit-14b-bidir", "resize": "bicubic"}, "5b-causal": {"size": "5b", "causal": True, "transformer": "dit-5b", "resize": "bilinear"}, - "14b-causal": {"size": "14b", "causal": True, "transformer": "dit-14b", "resize": "bicubic_center_crop"}, + "14b-causal": {"size": "14b", "causal": True, "transformer": "dit-14b", "resize": "bicubic"}, } _COMPONENT_CLASSES = { @@ -398,7 +398,7 @@ def resolve_original_checkpoint(config_dict, *, profile, passed_components, load _class_name="EdenLivePipeline", profile=profile, proxy_resize_mode=preset["resize"], - proxy_unconditional="zero" if profile == "14b-causal" else "same", + proxy_unconditional="same", proxy_projection_autocast=True, proxy_image_size=(480, 832), ) diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py index 83af59dd042f..b9647cf796c6 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -34,6 +34,7 @@ from .pipeline_output import EdenLivePipelineOutput +# docstyle-ignore: Keep the continuation prompt on the try/finally example. EXAMPLE_DOC_STRING = """ Examples: ```python @@ -113,8 +114,9 @@ class EdenLiveStreamState: """Mutable, single-video rollout owned by one pipeline. Create with `pipe.create_stream_state()`. The attention states own the only time cursor (`next_frame`). Encoder and decoder caches retain temporal history - independently of the attention window. Calls sharing a pipeline must be serial, including calls for different streams. - Close the state in a `finally` block to release its tensors. States are runtime objects and are not serialized. + independently of the attention window. Calls sharing a pipeline must be serial, including calls for different + streams. Close the state in a `finally` block to release its tensors. States are runtime objects and are not + serialized. """ causal_states: list[EdenLiveCausalState] @@ -677,14 +679,14 @@ def __call__( Args: stream_state (`EdenLiveStreamState`, *optional*): - State from `create_stream_state`. Returns one chunk: 9 RGB frames initially and 12 thereafter, or - 3 latent frames. Invalid inputs leave the state usable; model execution failures invalidate it. + State from `create_stream_state`. Returns one chunk: 9 RGB frames initially and 12 thereafter, or 3 + latent frames. Invalid inputs leave the state usable; model execution failures invalidate it. proxy_frames (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`, *optional*): - Exactly three already sampled RGB frames, required with `stream_state`. No temporal subsampling. - Tensor layout is `(3, 3, H, W)`; NumPy layout is `(3, H, W, 3)`. Float values must be in [0, 1]. + Exactly three already sampled RGB frames, required with `stream_state`. No temporal subsampling. Tensor + layout is `(3, 3, H, W)`; NumPy layout is `(3, H, W, 3)`. Float values must be in [0, 1]. prompt (`str` or `list[str]`, *optional*): - Text describing the output video. Supply this or `prompt_embeds` unless using temporal prompts. - In a stream, new text takes effect in this chunk; None retains the current text. + Text describing the output video. Supply this or `prompt_embeds` unless using temporal prompts. In a + stream, new text takes effect in this chunk; None retains the current text. 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. @@ -1080,9 +1082,8 @@ def _denoise_chunk( causal=True, condition_start=condition_start, ) - if self.model_size == "14b": - # A clean cache commit uses one zero timestep per sample, not per token. - model_timestep = torch.zeros(latents.shape[0], device=latents.device, dtype=torch.float32) + # A clean cache commit uses one zero timestep per sample, not per token. + model_timestep = torch.zeros(latents.shape[0], device=latents.device, dtype=torch.float32) self.transformer( model_input, model_timestep, diff --git a/tests/pipelines/edenlive/test_edenlive.py b/tests/pipelines/edenlive/test_edenlive.py index 5ef25788d0fb..847f9bee2998 100644 --- a/tests/pipelines/edenlive/test_edenlive.py +++ b/tests/pipelines/edenlive/test_edenlive.py @@ -315,10 +315,8 @@ def get_pipeline(self, **components): **(components or self.get_dummy_components()), profile=f"{self.profile}-{'causal' if causal else 'bidir'}", proxy_image_size=(32, 32), - proxy_resize_mode=("bicubic_center_crop" if causal else "bicubic") - if self.profile == "14b" - else "bilinear", - proxy_unconditional="zero" if self.profile == "14b" and causal else "same", + proxy_resize_mode="bicubic" if self.profile == "14b" else "bilinear", + proxy_unconditional="same", ) for component in pipe.components.values(): if hasattr(component, "eval"): @@ -513,6 +511,7 @@ class TestEdenLive14BChunkCausalPipeline( def test_zero_proxy_changes_unconditional_prediction(self): pipe = self.get_pipeline().to(torch_device) + pipe.register_to_config(proxy_unconditional="zero") inputs = self.get_dummy_inputs() inputs["output_type"] = "latent" zero_proxy = pipe(**inputs).frames @@ -721,10 +720,11 @@ def record_dit(module, args, kwargs): assert not torch.equal(conditional[0], conditional[1]) @pytest.mark.parametrize("profile", ["5b", "14b"]) - def test_clean_commit_timestep_shape_and_cfg(self, profile): + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_clean_commit_timestep_shape_and_cfg(self, profile, dtype): fixture = EdenLiveChunkCausalPipelineTesterConfig() fixture.profile = profile - pipe = fixture.get_pipeline().to(torch_device) + pipe = fixture.get_pipeline().to(torch_device, dtype=dtype) calls = [] def capture(module, args, kwargs): @@ -742,7 +742,7 @@ def capture(module, args, kwargs): 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 commits[0].shape == ((1,) if profile == "14b" else (1, 3)) + 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() @@ -964,12 +964,8 @@ def test_original_profiles_components_and_default_dtype(self, original_checkpoin assert (pipe.config.chunk_size, pipe.config.max_context_frames, pipe.config.sink_frames) == (3, 18, 6) else: assert pipe.scheduler.config.shift == 3 and not pipe.scheduler.config.stochastic_sampling - assert pipe.config.proxy_unconditional == ("zero" if fixture.profile == "14b-causal" else "same") - expected_resize = ( - "bilinear" - if fixture.profile.startswith("5b") - else ("bicubic_center_crop" if fixture.profile.endswith("causal") else "bicubic") - ) + assert pipe.config.proxy_unconditional == "same" + expected_resize = "bilinear" if fixture.profile.startswith("5b") else "bicubic" assert pipe.config.proxy_resize_mode == expected_resize assert tuple(pipe.config.proxy_image_size) == (480, 832) assert file_digests(fixture.source) == fixture.hashes @@ -1372,8 +1368,11 @@ class TestEdenLiveContextParallel(EdenLiveTransformerTesterConfig): ) @pytest.mark.parametrize("num_layers,token_time", [(2, False), (1, True)]) def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsdp, num_layers, token_time): + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + if not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < world_size: pytest.skip(f"This case requires {world_size} CUDA devices with NCCL.") + previous_backend, _ = _AttentionBackendRegistry.get_active_backend() deterministic = torch.are_deterministic_algorithms_enabled() warn_only = torch.is_deterministic_algorithms_warn_only_enabled() try: @@ -1412,6 +1411,7 @@ def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsd # Keep the existing small-model CP backend tolerance. torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2) finally: + _AttentionBackendRegistry.set_active_backend(previous_backend) torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) From ea77dec66f389643bcfd40ef6c80b6274708ad16 Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Tue, 20 Oct 2026 00:38:24 +0800 Subject: [PATCH 5/6] Refactor pipelines to follow conventions --- .../en/api/models/edenlive_transformer_3d.md | 79 +- docs/source/en/api/pipelines/edenlive.md | 289 +-- docs/source/en/api/schedulers/edenlive.md | 68 +- examples/edenlive/README.md | 103 +- examples/edenlive/inference_edenlive.py | 58 +- .../edenlive/inference_edenlive_streaming.py | 95 - scripts/convert_edenlive_to_diffusers.py | 271 +++ src/diffusers/__init__.py | 6 +- src/diffusers/models/__init__.py | 3 +- src/diffusers/models/autoencoders/__init__.py | 2 +- .../autoencoders/autoencoder_edenlive.py | 61 +- .../transformers/transformer_edenlive.py | 117 +- src/diffusers/pipelines/__init__.py | 4 +- src/diffusers/pipelines/edenlive/__init__.py | 6 +- .../pipelines/edenlive/loading_utils.py | 523 ----- .../pipelines/edenlive/pipeline_edenlive.py | 1237 +++-------- .../edenlive/pipeline_edenlive_causal.py | 643 ++++++ .../pipelines/edenlive/pipeline_utils.py | 358 +--- .../pipelines/edenlive/video_processor.py | 44 +- src/diffusers/utils/dummy_pt_objects.py | 15 - .../dummy_torch_and_transformers_objects.py | 6 +- .../test_models_autoencoder_edenlive.py | 94 + .../test_models_edenlive_proxy_projection.py | 91 + .../test_models_transformer_edenlive.py | 140 ++ ...test_models_transformer_edenlive_causal.py | 224 ++ tests/pipelines/edenlive/test_edenlive.py | 1807 ----------------- .../edenlive/test_pipeline_edenlive.py | 243 +++ .../edenlive/test_pipeline_edenlive_causal.py | 246 +++ tests/schedulers/test_scheduler_edenlive.py | 104 + 29 files changed, 2557 insertions(+), 4380 deletions(-) delete mode 100644 examples/edenlive/inference_edenlive_streaming.py create mode 100644 scripts/convert_edenlive_to_diffusers.py delete mode 100644 src/diffusers/pipelines/edenlive/loading_utils.py create mode 100644 src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_edenlive.py create mode 100644 tests/models/transformers/test_models_edenlive_proxy_projection.py create mode 100644 tests/models/transformers/test_models_transformer_edenlive.py create mode 100644 tests/models/transformers/test_models_transformer_edenlive_causal.py delete mode 100644 tests/pipelines/edenlive/test_edenlive.py create mode 100644 tests/pipelines/edenlive/test_pipeline_edenlive.py create mode 100644 tests/pipelines/edenlive/test_pipeline_edenlive_causal.py create mode 100644 tests/schedulers/test_scheduler_edenlive.py diff --git a/docs/source/en/api/models/edenlive_transformer_3d.md b/docs/source/en/api/models/edenlive_transformer_3d.md index ea8f29733647..aa345cf42a3c 100644 --- a/docs/source/en/api/models/edenlive_transformer_3d.md +++ b/docs/source/en/api/models/edenlive_transformer_3d.md @@ -12,51 +12,23 @@ specific language governing permissions and limitations under the License. # EdenLiveTransformer3DModel -The EdenLive diffusion transformer predicts flow velocities from video latents, text embeddings, and projected -DINO proxy features. The same class supports full-video bidirectional attention and causal attention over successive -chunks. The pipeline loads the configuration and weights for the selected profile; see the -[pipeline guide](../pipelines/edenlive) for original-checkpoint loading and video generation. +[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) uses a video diffusion transformer conditioned on +text, a reference image, and projected DINO proxy features. The same model class supports bidirectional and causal +attention. See the [pipeline guide](../pipelines/edenlive) for video generation. -## Inputs and output +```python +import torch +from diffusers import EdenLiveTransformer3DModel -Call `forward()` with the following inputs. Here `B` is batch size, `T` is the number of latent frames in the video or -current chunk, and `H` and `W` are latent spatial dimensions. +transformer = EdenLiveTransformer3DModel.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 +) +``` -| Argument | Shape or format | Meaning | -|---|---|---| -| `hidden_states` | `(B, in_channels, T, H, W)` | Noisy latents, plus image-conditioning channels for 14B | -| `timestep` | Scalar, `(B,)`, `(B, T)`, or `(B, num_patch_tokens)` | Model time; retain FP32 even with BF16 model weights | -| `encoder_hidden_states` | `(B, sequence_length, text_dim)` | UMT5 text embeddings | -| `encoder_hidden_states_image` | `(B, image_sequence_length, image_dim)` | CLIP reference-image embeddings, required for 14B and unused by 5B | -| `proxy_features` | `(B, proxy_channels, T, proxy_height, proxy_width)` | Required projected DINO features; spatial dimensions are interpolated to match the latents | +Use `subfolder="transformer_causal"` for the causal weights. Keep model timesteps in FP32 when calling the transformer +directly, including when its weights are BF16. -The temporal patch size is one; latent height and width must be divisible by the spatial patch size, which is two -for the released profiles. Proxy features must match the latent batch and frame counts. The 5B model uses 48 input -and output latent channels. The 14B model receives 36 input channels: 16 noisy latent channels, 16 image latent -channels, and four mask channels; it predicts 16 output channels. Proxy channels are appended internally and are -not included in `in_channels`. - -With `return_dict=True`, `forward()` returns `Transformer2DModelOutput.sample` with shape -`(B, out_channels, T, H, W)`. With `return_dict=False`, it returns the same prediction as the first tuple entry. -The prediction is a latent flow velocity, not a decoded video. - -The transformer's `text_timeline` is an internal token layout for 5B bidirectional inference. It contains `lengths` -(positive query-token counts per segment and sample) and `text_length` (the number of tokens per text prompt). -User-facing timelines in seconds are accepted and converted by `EdenLivePipeline`; they cannot be passed directly -to the transformer. - -## Causal attention state - -Use `create_causal_state(batch_size, height, width, ...)` to allocate state for a rollout, with height and width in -latent coordinates. The released defaults are three latent frames per chunk, a maximum context of 18 latent frames, -and six retained sink frames. Allocate a separate state for each classifier-free guidance branch. - -Denoising calls pass `causal_state=state` and `cache_update=False`. After a chunk has been denoised, a clean call -at timestep zero with `cache_update=True` commits its keys and values and advances the frame cursor. A short chunk -is supported only at the end of the rollout. The pipeline manages these calls and releases its states after every -invocation. For custom loops, call `state.reset()` before beginning an independent rollout. - -## EdenLiveTransformer3DModel API +## EdenLiveTransformer3DModel [[autodoc]] EdenLiveTransformer3DModel - forward @@ -64,32 +36,17 @@ invocation. For custom loops, call `state.reset()` before beginning an independe ## EdenLiveProxyProjection -The trained projection maps the last channel dimension of DINO features to `out_channels`, applies L2 normalization -on that dimension, and returns FP32 features. Inputs can retain arbitrary leading batch, frame, or spatial token -dimensions. The pipeline reshapes the result into the channel-first video layout expected by the transformer. - -Its `in_channels` must match the DINO encoder feature width and `out_channels` must match the transformer's -`proxy_channels`. Use the released trained projection together with its DINO encoder. +Projects DINO features into the transformer's conditioning space. Use the trained projection and DINO encoder from +the same checkpoint; the projection returns FP32 features. [[autodoc]] EdenLiveProxyProjection - forward ## EdenLiveCausalState -State belongs to one rollout and guidance branch and is not part of the model checkpoint. Create it through the -transformer so cache allocation matches the model's layers, dtype, device, and context-parallel layout. - -| Field | Meaning | -|---|---| -| `layers` | Per-layer unrotated keys, values, and valid-token counts | -| `batch_size`, `height`, `width` | Batch and latent spatial dimensions used to allocate the caches | -| `chunk_size`, `max_context_frames`, `sink_frames` | Causal context limits in latent frames | -| `tokens_per_frame` | Spatial patch-token count per latent frame | -| `cp_size`, `cp_rank` | Context-parallel degree and local rank | -| `next_frame` | Number of latent frames committed so far; also the next chunk's position offset | - -`reset()` clears valid-token counts and resets `next_frame` to zero, allowing the allocation to be reused for an -independent rollout with the same shape and parallel layout. +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 index 26d07ec9c163..830d9a0525ff 100644 --- a/docs/source/en/api/pipelines/edenlive.md +++ b/docs/source/en/api/pipelines/edenlive.md @@ -12,290 +12,129 @@ specific language governing permissions and limitations under the License. # EdenLive -[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) renders a proxy or reference video stream, a target first frame, and a text prompt into refined or stylized video. References can come from 2D/3D scenes, existing engines, or recordings. [`EdenLivePipeline`] supports bidirectional generation and causal generation with persistent history. +[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) 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 reference image -provides appearance; the proxy video provides scene structure and motion. For installation, see the -[example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive#setup). +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", profile="5b-causal") -pipe.to("cuda") +pipe = EdenLivePipeline.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, +).to("cuda") image = load_image("reference.png") video = load_video("proxy-24fps.mp4") -prompt = "A character explores a snowy landscape." frames = pipe( image=image, video=video, - prompt=prompt, + prompt="A character explores a snowy landscape.", num_frames=49, - generator=torch.Generator().manual_seed(42), ).frames[0] export_to_video(frames, "edenlive.mp4", fps=24) ``` -`from_pretrained()` loads the model, `pipe(...)` generates the video, and `.frames[0]` selects the first video in -the output batch. Each offline call starts a new video. Use the [streaming API](#streaming-causal-video) to continue -generation across calls. - -## Choose a profile - -Select `profile` when loading an original release. These defaults apply when generation options are omitted: - -| Profile | Sampling | Steps | Guidance scale | Height × width | Example fps | -|---|---|---|---|---|---| -| `5b-bidir` | Bidirectional | 48 | 3 | 704 × 1280 | 24 | -| `14b-bidir` | Bidirectional | 30 | 3 | 480 × 832 | 16 | -| `5b-causal` | Causal | 4 per chunk | 1 | 704 × 1280 | 24 | -| `14b-causal` | Causal | 4 per chunk | 1 | 480 × 832 | 16 | - -Causal profiles use four fixed denoising steps per chunk. Bidirectional profiles support `num_inference_steps`. -Set `guidance_scale` above one to use a `negative_prompt`. - -## Prepare the inputs - -Supply proxy frames sampled at the desired output frame rate. The pipeline selects frames at the VAE temporal -stride; timestamp resampling is handled by the [CLI example](#multiple-gpus-and-the-command-line-example). - -| Input | Accepted layout and values | -|---|---| -| `image` | RGB PIL image, or floating-point NumPy/tensor image in [0, 1] | -| `video` | List of RGB PIL frames; tensor `(T, 3, H, W)` or `(B, T, 3, H, W)`; NumPy `(T, H, W, 3)` or `(B, T, H, W, 3)` with floating-point values in [0, 1] | -| `prompt` | String or list of strings; use `prompt_embeds` for precomputed text embeddings | -| `height`, `width` | Output dimensions, multiples of 32 for 5B or 16 for 14B | -| `num_frames` | Positive frame count of the form $4k+1$, within the input video length | +Keep the proxy encoder and projection in FP32 when loading the other models in BF16, as shown above. -Here `B` is batch size, `T` is frame count, and `H` and `W` are image height and width. The reference image is -resized and center-cropped to the output dimensions. Omitting `num_frames` uses the largest valid input prefix. +## Causal generation -## Change prompts during a video - -### Causal chunk prompts - -Use `chunk_prompts` for either causal profile. This example reuses the pipeline, image, and video from -[Basic usage](#basic-usage). A 49-frame video has five chunks, so it takes five prompts: +Load the causal transformer and scheduler from the same repository. This works with either model size and uses +four fixed denoising steps per chunk. ```python -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.", -] - -frames = pipe( - image=image, - video=video, - chunk_prompts=chunk_prompts, - num_frames=49, - generator=torch.Generator().manual_seed(42), -).frames[0] -export_to_video(frames, "chunk-prompts.mp4", fps=24) +from diffusers import EdenLiveCausalPipeline, EdenLiveTransformer3DModel, EdenLiveFlowMatchEulerDiscreteScheduler + +causal_pipe = EdenLiveCausalPipeline.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", + transformer=EdenLiveTransformer3DModel.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer_causal", torch_dtype=torch.bfloat16 + ), + scheduler=EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( + "VAST-AI/EdenLive-5B-Diffusers", subfolder="scheduler_causal" + ), + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, +).to("cuda") ``` -`chunk_prompts` supplies one text prompt per chunk for a single video. - -### Bidirectional text timeline - -For `5b-bidir`, use `text_timeline` to assign prompts to time ranges. Reuse the image and video from Basic usage -and load the bidirectional profile: +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 -pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-bidir") -pipe.to("cuda") - -timeline = { - "fps": 24, - "num_frames": 49, - "segments": [ - {"start_sec": 0.0, "end_sec": 1.0, "text": "A character walks through snow."}, - {"start_sec": 1.0, "end_sec": 2.0, "text": "The character approaches a warm campfire."}, - ], -} - -frames = pipe( +frames = causal_pipe( image=image, video=video, - text_timeline=timeline, + 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, - generator=torch.Generator().manual_seed(42), ).frames[0] -export_to_video(frames, "timeline.mp4", fps=24) -``` - -Match the timeline's `fps` and `num_frames` to the input and requested output. Use contiguous segments from zero -through the last frame timestamp, with at least one latent frame per segment. - -## Streaming causal video - -Both causal profiles support successive calls with shared history. Initialize one state, then pass three already -sampled RGB proxy images per chunk. This example changes the prompt for the second chunk and saves the 21 output -frames as one video: - -```python -import numpy as np -import torch -from diffusers import EdenLivePipeline -from diffusers.utils import export_to_video, load_image - -pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-causal") -pipe.to("cuda") - -proxy_paths = [ - ["proxy/chunk-0-0.png", "proxy/chunk-0-1.png", "proxy/chunk-0-2.png"], - ["proxy/chunk-1-0.png", "proxy/chunk-1-1.png", "proxy/chunk-1-2.png"], -] -prompts = ["A character walks through snow.", "The character approaches a campfire."] - -state = pipe.create_stream_state( - image=load_image("reference.png"), - prompt=prompts[0], - generator=torch.Generator().manual_seed(42), - output_type="np", -) -chunks = [] -try: - for paths, current_prompt in zip(proxy_paths, prompts): - result = pipe( - stream_state=state, - proxy_frames=[load_image(path) for path in paths], - prompt=current_prompt, - ) - chunks.append(result.frames[0]) -finally: - pipe.close_stream(state) - -export_to_video(np.concatenate(chunks, axis=0), "stream.mp4", fps=24) +export_to_video(frames, "chunk-prompts.mp4", fps=24) ``` -Each call returns its chunk immediately: 9 RGB frames initially, then 12 per call. For incremental video-file -writing, see [`inference_edenlive_streaming.py`](https://github.com/huggingface/diffusers/blob/main/examples/edenlive/inference_edenlive_streaming.py). - -Set resolution, guidance, negative prompt, generator, and output format in `create_stream_state()`. Keep these -settings and the pipeline's device and dtype fixed for the stream. Update `prompt` or `prompt_embeds` in each -generation call; omitting both retains the current prompt. Run calls sharing a pipeline serially. - -Streaming uses batch size one and three proxy images per call. Tensor inputs have shape `(3, 3, H, W)`; NumPy -inputs use `(3, H, W, 3)`, with finite floating-point RGB values in [0, 1]. RGB output uses LightTAE. -Set `output_type="latent"` at initialization to receive three latent frames per call. For 14B, keep VAE slicing -and tiling disabled. Close a state after a model execution failure and initialize a new one. +`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. -## Load and save checkpoints +## LightTAE decoder -`EdenLivePipeline.from_pretrained()` accepts a Hub repository or a local directory. Original releases use a root -`model_index.json` and the selected profile's original weights. Saved Diffusers pipelines retain their profile and -decoder choice: +The full Wan VAE is the default. For faster decoding, pass the optional [`EdenLiveLightTAEDecoder`] to either pipeline: ```python -pipe.save_pretrained("edenlive-saved") -pipe = EdenLivePipeline.from_pretrained("edenlive-saved").to("cuda") -``` - -For local or cached weights: +from diffusers import EdenLiveLightTAEDecoder -```python pipe = EdenLivePipeline.from_pretrained( - "/path/to/EdenLive", profile="14b-causal", local_files_only=True + "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, "proxy_projection": torch.float32}, ).to("cuda") ``` -Original releases load the transformer, VAE, text encoder, CLIP, and LightTAE in BF16, with DINO and its projection -in FP32. Use `dtype` for an explicit dtype or per-component mapping. Original-release loading supports -`revision`, `cache_dir`, and `token`; `device_map`, `variant`, and quantization options require a saved -Diffusers pipeline. - -### CPU offload +## Notes -Load the pipeline on CPU and enable model offload before generation: +| Checkpoint | Default height × width | Example fps | Bidirectional steps | +|---|---|---|---| +| [5B](https://huggingface.co/VAST-AI/EdenLive-5B-Diffusers) | 704 × 1280 | 24 | 48 | +| [14B](https://huggingface.co/VAST-AI/EdenLive-14B-Diffusers) | 480 × 832 | 16 | 30 | -```python -pipe = EdenLivePipeline.from_pretrained("VAST-AI/EdenLive", profile="5b-causal") -pipe.enable_model_cpu_offload() -``` - -The offline and streaming calls above use the same interface with CPU offload enabled. - -### Offline decoder selection - -Original releases use [`EdenLiveLightTAEDecoder`] by default. To decode offline videos with the full Wan VAE, -pass `decoder=None` when loading: - -```python -pipe = EdenLivePipeline.from_pretrained( - "VAST-AI/EdenLive", profile="5b-causal", decoder=None -).to("cuda") -``` - -Pass an initialized decoder as `decoder=my_decoder` to use a compatible custom decoder. -`output_type="latent"` returns normalized diffusion latents before decoding. - -## Multiple GPUs and the command-line example - -The CLI resamples proxy video timestamps and supports Ulysses context parallelism with PyTorch FSDP2 on two or four -GPUs. From the Diffusers repository root: - -```bash -torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ - --model VAST-AI/EdenLive --profile 14b-causal \ - --video proxy.mp4 --image reference.png \ - --prompt "A character explores a snowy landscape." \ - --num-frames 81 --output edenlive.mp4 -``` - -See the [example README](https://github.com/huggingface/diffusers/tree/main/examples/edenlive) for installation, -single-GPU commands, prompt files, and output reports. +- 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, multi-GPU inference, and checkpoint conversion. ## EdenLivePipeline -Load weights with `from_pretrained()`. Generate a complete video with `pipe(image=..., video=..., prompt=...)`, -or initialize a stream with `create_stream_state()` and generate chunks with -`pipe(stream_state=..., proxy_frames=..., prompt=...)`. Call `close_stream()` when the stream ends. - [[autodoc]] EdenLivePipeline - all - __call__ - - create_stream_state - - close_stream - -## EdenLiveLightTAEDecoder - -`decode()` takes raw VAE latents in `(B, C, T_latent, H_latent, W_latent)` layout and returns RGB in -`(B, 3, T, H, W)`, with values in [-1, 1]. To decode the pipeline's `output_type="latent"` yourself, first invert -the VAE normalization: -$$ -z_{\mathrm{raw}} = z_{\mathrm{diffusion}} \odot \sigma + \mu. -$$ +## EdenLiveCausalPipeline -Here $\mu$ and $\sigma$ are `vae.config.latents_mean` and `vae.config.latents_std`, broadcast over batch, time, -and spatial dimensions; $\odot$ is element-wise multiplication. Match the latents to the decoder's device and dtype. +[[autodoc]] EdenLiveCausalPipeline + - all + - __call__ -Pass an `EdenLiveLightTAEDecodeCache` to retain decoder history across calls. Without a cache, each decode starts -a new video. The pipeline manages this cache for streaming generation. +## EdenLiveLightTAEDecoder [[autodoc]] EdenLiveLightTAEDecoder - decode ## EdenLivePipelineOutput -Read generated frames from `.frames`, or use `return_dict=False` to receive `(frames,)`: - -| `output_type` | `.frames` layout | Values | -|---|---|---| -| `"np"` (default) | NumPy `(B, T, H, W, 3)` | Floating-point RGB in [0, 1] | -| `"pt"` | Tensor `(B, T, 3, H, W)` | Floating-point RGB in [0, 1] | -| `"pil"` | One list of PIL RGB frames per video | PIL images | -| `"latent"` | Tensor `(B, C, T_latent, H_latent, W_latent)` | Normalized diffusion latents | - -Use `result.frames[0]` to select the first video for `export_to_video()`. - [[autodoc]] pipelines.edenlive.pipeline_output.EdenLivePipelineOutput diff --git a/docs/source/en/api/schedulers/edenlive.md b/docs/source/en/api/schedulers/edenlive.md index 8564bf33e377..f4a2c86165c8 100644 --- a/docs/source/en/api/schedulers/edenlive.md +++ b/docs/source/en/api/schedulers/edenlive.md @@ -12,69 +12,15 @@ specific language governing permissions and limitations under the License. # EdenLiveFlowMatchEulerDiscreteScheduler -This scheduler implements the deterministic flow Euler trajectory used by EdenLive 14B bidirectional inference and -the stochastic trajectory used by the distilled causal checkpoints. The pipeline selects the appropriate -configuration when loading a profile. The 5B bidirectional profile uses [`UniPCMultistepScheduler`] instead. +[`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`]. -## Released schedules +Use the scheduler saved with the checkpoint: `scheduler` for bidirectional generation and `scheduler_causal` for +causal generation. The released causal schedules have four fixed steps per chunk. See the +[pipeline guide](../pipelines/edenlive#causal-generation) for loading the causal components. -| Profile | Scheduler mode | Default evaluations | Model timesteps | -|---|---|---|---| -| `14b-bidir` | Deterministic Euler, `shift=3.0` | 30 per video | Generated from the requested step count | -| `5b-causal` | `stochastic_sampling=True` | 4 per chunk | `[1000, 937.5, 833.3333333333, 625]` | -| `14b-causal` | `stochastic_sampling=True` | 4 per chunk | `[1000, 900, 750, 500]` | - -The causal schedules are stored as `default_timesteps` and already contain the released time shift. Passing a -different `num_inference_steps` without replacing the schedule raises an error. Arbitrary step counts are not a -supported inference setting for these distilled checkpoints. - -## Set a schedule - -Call `set_timesteps()` before stepping through a video or a new causal chunk. It resets both the step index and the -begin index. Choose one schedule source: - -- `timesteps`: explicit model times in [0, `num_train_timesteps`], converted to noise levels by dividing by - `num_train_timesteps` (1000 by default). -- `sigmas`: explicit noise levels in [0, 1]. Do not supply both `timesteps` and `sigmas`. -- Neither argument: use `default_timesteps` if configured; otherwise generate a schedule from a positive - `num_inference_steps` and apply `shift`. - -Explicit values must be finite and strictly decreasing. They are not shifted again. An optional terminal zero is -removed from the model evaluation points, and the internal `sigmas` tensor always receives one terminal zero. -If `num_inference_steps` is provided with an explicit or default schedule, it must match the number of positive -evaluation points. `timesteps` and `sigmas` are stored in FP32 on the requested `device`. - -For example, the 14B causal schedule can be constructed independently for a custom inference loop: - -```py -from diffusers import EdenLiveFlowMatchEulerDiscreteScheduler - -scheduler = EdenLiveFlowMatchEulerDiscreteScheduler( - stochastic_sampling=True, - default_timesteps=[1000, 900, 750, 500], -) -scheduler.set_timesteps(device="cuda") -``` - -For a loaded pipeline, use its existing `pipe.scheduler` to preserve the profile configuration. - -## Advance the sample - -Pass the model's flow velocity, the current value from `scheduler.timesteps`, and the noisy latent sample to -`step(model_output, timestep, sample, ...)`. Deterministic mode takes an Euler step. Stochastic mode first predicts -a clean endpoint and then adds fresh noise at the next noise level. The terminal step returns the clean endpoint -without drawing noise. - -Use `generator` to control stochastic draws, or supply `variance_noise` for a nonterminal stochastic step. The two -arguments are mutually exclusive. Supplied noise must have the same shape and dtype as the sample and be on its -device. This allows a distributed caller to generate noise once and share it across ranks. Noise is unused in -deterministic mode and at the terminal stochastic step. - -`step()` returns `FlowMatchEulerDiscreteSchedulerOutput.prev_sample`, or `(prev_sample,)` when -`return_dict=False`. The result preserves the input sample's dtype. `scale_model_input()` returns its input -unchanged. Once the schedule is exhausted, call `set_timesteps()` before processing the next chunk. - -## API +## EdenLiveFlowMatchEulerDiscreteScheduler [[autodoc]] EdenLiveFlowMatchEulerDiscreteScheduler - set_timesteps diff --git a/examples/edenlive/README.md b/examples/edenlive/README.md index cbc3f1551701..9e360474e60e 100644 --- a/examples/edenlive/README.md +++ b/examples/edenlive/README.md @@ -1,86 +1,54 @@ # EdenLive inference -Render game from a reference image, prompt, and proxy video with `EdenLivePipeline`. -See the [pipeline guide](../../docs/source/en/api/pipelines/edenlive.md) for the Python API. +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/) build and matching torchvision. -From the Diffusers repository root, install the pipeline dependencies: +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 +pip install -e . "transformers>=4.56.0" accelerate sentencepiece ftfy av imageio imageio-ffmpeg ``` -For the CLI examples, add video reading and writing: - -```bash -pip install av imageio imageio-ffmpeg -``` - -The [example requirements](requirements.txt) collect both groups, plus PyTorch and torchvision. +The [requirements file](requirements.txt) also lists the example dependencies. ## Single GPU ```bash python examples/edenlive/inference_edenlive.py \ - --model VAST-AI/EdenLive --profile 5b-causal \ + --model VAST-AI/EdenLive-5B-Diffusers --causal \ --video proxy.mp4 --image reference.png \ --prompt "A character explores a snowy landscape." \ - --num-frames 49 --output edenlive.mp4 + --num-frames 49 ``` -`--model` accepts a Hub repository, local release, or saved Diffusers pipeline. Select its profile with `--profile`. -Use `--local-files-only` for local or cached weights. +Use `--causal` for the causal weights or omit it for bidirectional generation. `--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. -| Profile | Default height x width | Default fps | Default steps | Guidance scale | -|---|---|---|---|---| -| `5b-bidir` | 704 x 1280 | 24 | 48 | 3 | -| `14b-bidir` | 480 x 832 | 16 | 30 | 3 | -| `5b-causal` | 704 x 1280 | 24 | 4 per chunk | 1 | -| `14b-causal` | 480 x 832 | 16 | 4 per chunk | 1 | - -The CLI resamples to `--fps`, center-crops, and resizes the input. Set `--num-frames` to a positive $4k+1$ within -the resampled clip, or omit it to use the largest valid prefix. Custom height and width use multiples of 32 for 5B -and 16 for 14B. - -Causal profiles use four fixed steps per chunk. Bidirectional profiles support `--num-inference-steps`. -Set `--guidance-scale` above one to use `--negative-prompt`. +The default output is `edenlive.mp4`, with timing and input details in `edenlive.json`. Set `--output` to change it. +Use `--lighttae` for faster decoding or `--vae-tiling` to reduce VAE memory use. Run +`python examples/edenlive/inference_edenlive.py --help` for all options. ## Multiple GPUs -Use two or four GPUs with Ulysses context parallelism and PyTorch FSDP2 (`torch.distributed.fsdp.fully_shard`): +Use Ulysses context parallelism and FSDP2 on two or four GPUs: ```bash torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ - --model VAST-AI/EdenLive --profile 14b-causal \ + --model VAST-AI/EdenLive-14B-Diffusers --causal \ --video proxy.mp4 --image reference.png \ --prompt "A character explores a snowy landscape." \ - --num-frames 81 --output edenlive.mp4 + --num-frames 81 ``` -Set `--nproc_per_node=4` for four GPUs. Add `--no-fsdp` to use context parallelism with replicated transformer weights. +Set `--nproc_per_node=4` for four GPUs. Add `--no-fsdp` to keep transformer weights replicated across GPUs. -## Temporal prompts +## Changing prompts -For `5b-bidir`, pass `--text-timeline timeline.json`. This example describes 49 frames at 24 fps: - -```json -{ - "fps": 24, - "num_frames": 49, - "segments": [ - {"start_sec": 0.0, "end_sec": 1.0, "text": "A character walks through snow."}, - {"start_sec": 1.0, "end_sec": 2.0, "text": "The character approaches a warm campfire."} - ] -} -``` - -Match the JSON `fps` and `num_frames` to the CLI values. Use contiguous segments covering zero through the last -frame timestamp, with at least one latent frame per segment. - -For either causal profile, pass `--chunk-prompts prompts.json` with one prompt per chunk. A 49-frame video uses -five chunks, including a shorter final chunk: +With `--causal`, add `--chunk-prompts prompts.json`. For 49 frames, supply five prompts: ```json [ @@ -92,33 +60,16 @@ five chunks, including a shorter final chunk: ] ``` -Include `--prompt` in the command; the timeline or chunk prompts supply the text used during generation. +The list replaces `--prompt`, with one prompt per chunk. The negative prompt remains fixed. -## Streaming from proxy images +## Convert original checkpoints -`inference_edenlive_streaming.py` supports both causal profiles. Pass ordered RGB images with `--proxy-frames`, -already sampled at latent cadence: three images per chunk. +To convert a local original release instead of loading the published Diffusers weights: ```bash -python examples/edenlive/inference_edenlive_streaming.py \ - --model /path/to/EdenLive --profile 5b-causal \ - --image reference.png --prompt "A character walks through snow" \ - --proxy-frames proxy/000.png proxy/004.png proxy/008.png proxy/012.png proxy/016.png proxy/020.png \ - --chunk-prompts stream_prompts.json --output stream.mp4 --fps 24 --local-files-only +python scripts/convert_edenlive_to_diffusers.py \ + --source /path/to/EdenLive --size 5b --output /path/to/EdenLive-5B-Diffusers ``` -The optional `stream_prompts.json` contains one string per chunk, for example -`["A character walks through snow", "A character walks through fire"]`. Omit it to use `--prompt` throughout. - -The script writes video incrementally: 9 frames for the first chunk, then 12 per chunk. Set `--fps` to the desired -playback rate (24 by default). For application integration, see the -[streaming API](../../docs/source/en/api/pipelines/edenlive.md#streaming-causal-video). - -## Offline options and reports - -LightTAE is the default decoder for original releases. Use `--full-vae` for Wan VAE decoding and `--vae-tiling` for -tiling on the Wan VAE. Set the noise seed with `--seed` and the attention backend with `--attention-backend`. -Run `python examples/edenlive/inference_edenlive.py --help` for all options. - -The offline script writes a JSON report beside the video with arguments, sampled frame indices, timing, and peak GPU -memory. Timing covers the pipeline call, including conditioning and decoding. +Use `--size 14b` and a separate output directory for 14B. Each package includes both sampling variants and the +optional LightTAE decoder. diff --git a/examples/edenlive/inference_edenlive.py b/examples/edenlive/inference_edenlive.py index d630c1630c5b..e368eb0e4af3 100644 --- a/examples/edenlive/inference_edenlive.py +++ b/examples/edenlive/inference_edenlive.py @@ -26,7 +26,14 @@ import torch.distributed as dist from PIL import Image -from diffusers import ContextParallelConfig, EdenLivePipeline +from diffusers import ( + ContextParallelConfig, + EdenLiveCausalPipeline, + EdenLiveFlowMatchEulerDiscreteScheduler, + EdenLiveLightTAEDecoder, + EdenLivePipeline, + EdenLiveTransformer3DModel, +) from diffusers.utils import export_to_video @@ -106,14 +113,15 @@ def native_timestep(module, args, kwargs): def parse_args(): parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model", default="VAST-AI/EdenLive", help="EdenLive release directory or Hub repository.") - parser.add_argument("--profile", choices=("5b-bidir", "14b-bidir", "5b-causal", "14b-causal"), required=True) - parser.add_argument("--full-vae", action="store_true", help="Decode with the full Wan VAE instead of LightTAE.") + parser.add_argument("--model", required=True, help="Standard Diffusers model directory or Hub repository.") + parser.add_argument("--causal", action="store_true", help="Use the causal transformer and scheduler subfolders.") + 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", required=True) + parser.add_argument("--prompt") parser.add_argument("--negative-prompt", default="") - parser.add_argument("--text-timeline", type=Path, help="JSON timeline for 5B bidirectional inference.") 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) @@ -127,7 +135,12 @@ def parse_args(): parser.add_argument("--no-fsdp", action="store_true", help="Use replicated DiT weights with context parallelism.") parser.add_argument("--local-files-only", action="store_true") parser.add_argument("--vae-tiling", action="store_true") - return parser.parse_args() + 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(): @@ -142,11 +155,24 @@ def main(): dist.init_process_group("nccl", device_id=device) rank = dist.get_rank() if dist.is_initialized() else 0 try: - loading_kwargs = {"decoder": None} if args.full_vae else {} - pipeline = EdenLivePipeline.from_pretrained( + loading_kwargs = {"local_files_only": args.local_files_only} + components = {"decoder": None} + pipeline_class = EdenLiveCausalPipeline if args.causal else EdenLivePipeline + if args.causal: + components["transformer"] = EdenLiveTransformer3DModel.from_pretrained( + args.model, subfolder="transformer_causal", torch_dtype=torch.bfloat16, **loading_kwargs + ) + components["scheduler"] = EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( + args.model, subfolder="scheduler_causal", **loading_kwargs + ) + 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, - profile=args.profile, - local_files_only=args.local_files_only, + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, + **components, **loading_kwargs, ) if args.attention_backend is not None: @@ -158,11 +184,11 @@ def main(): if args.vae_tiling: pipeline.vae.enable_tiling() pipeline.set_progress_bar_config(disable=rank != 0) - is_5b = args.profile.startswith("5b-") + 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.profile == "14b-causal" else 0) + 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, @@ -178,10 +204,8 @@ def main(): for name in ("num_inference_steps", "guidance_scale"): if getattr(args, name) is not None: inputs[name] = getattr(args, name) - for name in ("text_timeline", "chunk_prompts"): - path = getattr(args, name) - if path is not None: - inputs[name] = json.loads(path.read_text()) + 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() diff --git a/examples/edenlive/inference_edenlive_streaming.py b/examples/edenlive/inference_edenlive_streaming.py deleted file mode 100644 index 70e4be3c84d5..000000000000 --- a/examples/edenlive/inference_edenlive_streaming.py +++ /dev/null @@ -1,95 +0,0 @@ -# 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. - -"""Generate successive EdenLive chunks from ordered, already sampled proxy images.""" - -import argparse -import json -from pathlib import Path - -import av -import numpy as np -import torch -from PIL import Image - -from diffusers import EdenLivePipeline - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model", default="VAST-AI/EdenLive") - parser.add_argument("--profile", choices=("5b-causal", "14b-causal"), required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--prompt", required=True) - parser.add_argument("--proxy-frames", nargs="+", required=True, help="Ordered, already sampled RGB image paths.") - parser.add_argument("--chunk-prompts", type=Path, help="JSON file containing one prompt string per chunk.") - parser.add_argument("--height", type=int) - parser.add_argument("--width", type=int) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--fps", type=int, help="Playback rate; defaults to 16 for 14b-causal and 24 for 5b-causal.") - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--local-files-only", action="store_true") - args = parser.parse_args() - args.fps = args.fps if args.fps is not None else (16 if args.profile == "14b-causal" else 24) - if len(args.proxy_frames) % 3: - parser.error("--proxy-frames must contain a multiple of three image paths.") - if args.fps <= 0: - parser.error("--fps must be positive.") - prompts = None - if args.chunk_prompts is not None: - prompts = json.loads(args.chunk_prompts.read_text()) - if ( - not isinstance(prompts, list) - or len(prompts) != len(args.proxy_frames) // 3 - or not all(isinstance(prompt, str) for prompt in prompts) - ): - parser.error("--chunk-prompts must contain one string per chunk.") - pipe = EdenLivePipeline.from_pretrained( - args.model, profile=args.profile, local_files_only=args.local_files_only - ).to("cuda") - with Image.open(args.image) as image: - state = pipe.create_stream_state( - image=image.convert("RGB"), - prompt=args.prompt, - height=args.height, - width=args.width, - generator=torch.Generator().manual_seed(args.seed), - output_type="np", - ) - try: - args.output.parent.mkdir(parents=True, exist_ok=True) - with av.open(str(args.output), mode="w") as container: - stream = container.add_stream("libx264", rate=args.fps) - stream.height, stream.width = state.height, state.width - stream.pix_fmt = "yuv420p" - for index in range(0, len(args.proxy_frames), 3): - frames = [] - for path in args.proxy_frames[index : index + 3]: - with Image.open(path) as frame: - frames.append(frame.convert("RGB")) - output = pipe( - stream_state=state, proxy_frames=frames, prompt=None if prompts is None else prompts[index // 3] - ) - for frame in output.frames[0]: - pixels = np.clip(frame * 255, 0, 255).round().astype(np.uint8) - for packet in stream.encode(av.VideoFrame.from_ndarray(pixels, format="rgb24")): - container.mux(packet) - for packet in stream.encode(): - container.mux(packet) - finally: - pipe.close_stream(state) - - -if __name__ == "__main__": - main() diff --git a/scripts/convert_edenlive_to_diffusers.py b/scripts/convert_edenlive_to_diffusers.py new file mode 100644 index 000000000000..6af78c18a7f5 --- /dev/null +++ b/scripts/convert_edenlive_to_diffusers.py @@ -0,0 +1,271 @@ +# 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, + EdenLiveProxyProjection, + 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): + output.mkdir(parents=True, exist_ok=True) + manifest = {} + for folder, destination in ((f"dit-{size}-bidir", "transformer"), (f"dit-{size}", "transformer_causal")): + 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_channels"] = state["patch_embedding.weight"].shape[1] - config["in_channels"] + config["condition_rmsnorm"] = "patch_embedding.condition_norm.weight" in state + 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 / destination, manifest) + del model, state + 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, backbone + projection = {key: value for key, value in state.items() if key.startswith("patch_projection.")} + weight = projection["patch_projection.weight"] + with torch.device("meta"): + model = EdenLiveProxyProjection(in_channels=weight.shape[1], out_channels=weight.shape[0]) + save_component(model, projection, output / "proxy_projection", manifest) + del model, state, projection, weight + 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] + scheduler.save_pretrained(output / "scheduler") + EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True, default_timesteps=timesteps).save_pretrained( + output / "scheduler_causal" + ) + index = { + "_class_name": "EdenLivePipeline", + "transformer": ["diffusers", "EdenLiveTransformer3DModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + "decoder": [None, None], + "text_encoder": ["transformers", "UMT5EncoderModel"], + "tokenizer": ["transformers", type(tokenizer).__name__], + "proxy_encoder": ["transformers", "DINOv3ViTModel"], + "proxy_projection": ["diffusers", "EdenLiveProxyProjection"], + "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, + } + (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", "proxy_projection")), + ("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) + args = parser.parse_args() + convert(args.source, args.output, args.size) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index bb3d16619fbf..6b81504cceff 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -300,7 +300,6 @@ "EasyAnimateTransformer3DModel", "EdenLiveCausalState", "EdenLiveLightTAEDecoder", - "EdenLiveLightTAEDecodeCache", "EdenLiveProxyProjection", "EdenLiveTransformer3DModel", "ErnieImageTransformer2DModel", @@ -686,9 +685,9 @@ "EasyAnimateControlPipeline", "EasyAnimateInpaintPipeline", "EasyAnimatePipeline", + "EdenLiveCausalPipeline", "EdenLivePipeline", "EdenLivePipelineOutput", - "EdenLiveStreamState", "ErnieImagePipeline", "Flux2KleinInpaintPipeline", "Flux2KleinKVPipeline", @@ -1223,7 +1222,6 @@ DreamLiteUNetModel, EasyAnimateTransformer3DModel, EdenLiveCausalState, - EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel, @@ -1585,9 +1583,9 @@ EasyAnimateControlPipeline, EasyAnimateInpaintPipeline, EasyAnimatePipeline, + EdenLiveCausalPipeline, EdenLivePipeline, EdenLivePipelineOutput, - EdenLiveStreamState, ErnieImagePipeline, Flux2KleinInpaintPipeline, Flux2KleinKVPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 64661d17491b..adaca45dcd5b 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -31,7 +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"] = ["EdenLiveLightTAEDecodeCache", "EdenLiveLightTAEDecoder"] + _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"] @@ -222,7 +222,6 @@ AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, - EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder, Kandinsky6SRVAE, LongCatAudioDiTVae, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 5fa81587e65a..a9664df7b8c0 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -1,7 +1,7 @@ from .autoencoder_asym_kl import AsymmetricAutoencoderKL from .autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from .autoencoder_dc import AutoencoderDC -from .autoencoder_edenlive import EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder +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 index cba0b56f28a2..a91ac9f64b99 100644 --- a/src/diffusers/models/autoencoders/autoencoder_edenlive.py +++ b/src/diffusers/models/autoencoders/autoencoder_edenlive.py @@ -25,14 +25,8 @@ @dataclass -class EdenLiveLightTAEDecodeCache: - """Temporal memory for one LightTAE stream with a fixed model, batch, shape, device and dtype.""" - - dec_state: dict[int, torch.Tensor] = field(default_factory=dict) - - def reset(self): - """Discard temporal memory and restore first-chunk trimming.""" - self.dec_state.clear() +class _EdenLiveTAEDecodeState: + memory: dict[int, torch.Tensor] = field(default_factory=dict) class EdenLiveTAEClamp(nn.Module): @@ -56,12 +50,12 @@ def __init__(self, channels): 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.get(id(self)) + 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[id(self)] = frames[:, -1:].clone() + state.memory[id(self)] = frames[:, -1:].clone() return self.act(self.conv(torch.cat([hidden_states, past], dim=1)) + self.skip(hidden_states)) @@ -80,7 +74,7 @@ class EdenLiveLightTAEDecoder(ModelMixin, ConfigMixin): 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 unless an explicit cache is supplied, and is never retained on the model. + `decode` call and is never retained on the model. Parameters: latent_channels (`int`, *optional*, defaults to 16): @@ -91,6 +85,8 @@ class EdenLiveLightTAEDecoder(ModelMixin, ConfigMixin): 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__() @@ -120,24 +116,8 @@ def scale_factor_temporal(self): def scale_factor_spatial(self): return 8 * self.config.patch_size - def _convert_original_state_dict(self, state_dict): - converted = {} - for key, value in state_dict.items(): - if key.startswith("encoder."): - continue - key = key.replace("decoder.blocks.", "decoder.", 1) - converted[key] = value - for index, layer in enumerate(self.decoder): - if isinstance(layer, EdenLiveTAETemporalUpsample): - key = f"decoder.{index}.conv.weight" - expected = layer.conv.weight.shape[0] - if key in converted and converted[key].shape[0] > expected: - # Original LightTAE checkpoints retain a stride-two tensor at the stride-one first stage. - converted[key] = converted[key][-expected:] - return converted - def _decode_chunk(self, latents, state): - first_chunk = not 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: @@ -154,10 +134,8 @@ def _decode_chunk(self, latents, state): return frames.permute(0, 2, 1, 3, 4) * 2 - 1 @apply_forward_hook - def decode( - self, latents: torch.Tensor, return_dict: bool = True, *, cache: EdenLiveLightTAEDecodeCache | None = None - ): - """Decode raw VAE latents, optionally continuing an existing stream. + 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`): @@ -167,39 +145,32 @@ def decode( 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. - cache (`EdenLiveLightTAEDecodeCache`, *optional*): - Temporal memory shared across calls. Only the first chunk is trimmed; subsequent calls produce - `T_latent * scale_factor_temporal` frames. Without a cache, every call starts a new video. 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 for a new video, or `T_latent * scale_factor_temporal` for an - existing cache, and spatial dimensions multiplied by `scale_factor_spatial`. Internal chunks share - temporal state, including a short final chunk. + 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.") - state = {} if cache is None else cache.dec_state + # 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, *, cache: EdenLiveLightTAEDecodeCache | None = None - ) -> DecoderOutput | tuple[torch.Tensor]: + 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. - cache (`EdenLiveLightTAEDecodeCache`, *optional*): - Temporal memory carried across decode calls. 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, cache=cache) + return self.decode(latents, return_dict=return_dict) diff --git a/src/diffusers/models/transformers/transformer_edenlive.py b/src/diffusers/models/transformers/transformer_edenlive.py index 58a57b3b7dee..3acc4a381e8c 100644 --- a/src/diffusers/models/transformers/transformer_edenlive.py +++ b/src/diffusers/models/transformers/transformer_edenlive.py @@ -105,16 +105,12 @@ def __call__( attention_mask=None, rotary_emb=None, *, - text_timeline=None, - query_start=0, 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.") - if text_timeline is not None: - raise ValueError("Temporal text conditioning is supported by the 5B model only.") 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.") @@ -149,48 +145,16 @@ def __call__( .type_as(query) ) - if text_timeline is None: - output = dispatch_attention_fn( - query, - key, - value, - attn_mask=attention_mask, - dropout_p=0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=self._parallel_config if encoder_hidden_states is None else None, - ) - else: - if encoder_hidden_states is None or attention_mask is not None: - raise ValueError("Temporal text requires cross-attention without an additional attention mask.") - lengths = text_timeline["lengths"] - text_length = int(text_timeline["text_length"]) - query_end = query_start + query.shape[1] - samples = [] - for sample_index, segments in enumerate(lengths): - pieces = [] - segment_start = 0 - for segment_index, segment_length in enumerate(segments): - segment_end = segment_start + segment_length - start, end = max(query_start, segment_start), min(query_end, segment_end) - if start < end: - context_start = segment_index * text_length - pieces.append( - dispatch_attention_fn( - query[sample_index : sample_index + 1, start - query_start : end - query_start], - key[sample_index : sample_index + 1, context_start : context_start + text_length], - value[sample_index : sample_index + 1, context_start : context_start + text_length], - attn_mask=None, - dropout_p=0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=None, - ) - ) - segment_start = segment_end - samples.append(torch.cat(pieces, dim=1)) - output = torch.cat(samples, dim=0) - + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config if encoder_hidden_states is None else None, + ) output = output.flatten(2, 3).type_as(query) if image_output is not None: output = output + image_output @@ -287,23 +251,32 @@ class EdenLiveProxyProjection(ModelMixin, ConfigMixin): Number of projected proxy feature channels. """ + _no_split_modules = [] + @register_to_config def __init__(self, in_channels: int = 1024, out_channels: int = 1024): super().__init__() self.patch_projection = nn.Linear(in_channels, out_channels) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward( + self, hidden_states: torch.Tensor, return_dict: bool = True + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: r""" Args: hidden_states (`torch.Tensor`): DINO patch features with `in_channels` in the final dimension. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output instead of a tuple. Returns: - `torch.Tensor`: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: Unit-normalized projected features in float32, with `out_channels` in the final dimension. """ hidden_states = self.patch_projection(hidden_states) - return F.normalize(hidden_states, p=2, dim=-1).float() + hidden_states = F.normalize(hidden_states, p=2, dim=-1).float() + if not return_dict: + return (hidden_states,) + return Transformer2DModelOutput(sample=hidden_states) class EdenLivePatchEmbedding(nn.Conv3d): @@ -335,7 +308,6 @@ def __init__( added_kv_proj_dim: int | None = None, cross_attention_dim_head: int | None = None, processor=None, - is_cross_attention=None, ): super().__init__() @@ -363,10 +335,7 @@ def __init__( 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) - if is_cross_attention is not None: - self.is_cross_attention = is_cross_attention - else: - self.is_cross_attention = cross_attention_dim_head is not None + self.is_cross_attention = cross_attention_dim_head is not None self.set_processor(processor or EdenLiveAttnProcessor()) @@ -480,9 +449,7 @@ def forward( if timestep_seq_len is not None: timestep = timestep.unflatten(0, (-1, timestep_seq_len)) - time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype - if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: - timestep = timestep.to(time_embedder_dtype) + 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)) @@ -567,7 +534,7 @@ def forward( @maybe_allow_in_graph class EdenLiveTransformerBlock(nn.Module): - def __init__(self, dim, ffn_dim, num_heads, cross_attn_norm=True, eps=1e-6, added_kv_proj_dim=None): + 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( @@ -582,7 +549,7 @@ def __init__(self, dim, ffn_dim, num_heads, cross_attn_norm=True, eps=1e-6, adde cross_attention_dim_head=dim // num_heads, processor=EdenLiveAttnProcessor(), ) - self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + 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) @@ -599,8 +566,6 @@ def forward( cache_update=False, cache_window_tokens=None, cache_sink_tokens=0, - text_timeline=None, - query_start=0, text_context_length=512, ) -> torch.Tensor: if temb.ndim == 4: @@ -628,8 +593,6 @@ def forward( attn_output = self.attn2( norm_hidden_states, encoder_hidden_states, - text_timeline=text_timeline, - query_start=query_start, text_context_length=text_context_length, ) hidden_states = hidden_states + attn_output @@ -664,8 +627,6 @@ class EdenLiveTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Atte Feed-forward intermediate dimension. num_layers (`int`, defaults to `40`): Number of transformer blocks. - cross_attn_norm (`bool`, defaults to `True`): - Whether to normalize the hidden states before cross-attention. 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`): @@ -683,7 +644,7 @@ class EdenLiveTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Atte """ _supports_gradient_checkpointing = True - _no_split_modules = ["EdenLiveTransformerBlock", "EdenLiveTimeTextImageEmbedding"] + _no_split_modules = ["EdenLiveTransformerBlock", "EdenLiveTimeTextImageEmbedding", "EdenLivePatchEmbedding"] _repeated_blocks = ["EdenLiveTransformerBlock"] _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] _skip_keys = ["causal_state"] @@ -711,7 +672,6 @@ def __init__( freq_dim: int = 256, ffn_dim: int = 13824, num_layers: int = 40, - cross_attn_norm: bool = True, qk_norm: str = "rms_norm_across_heads", eps: float = 1e-6, image_dim: int | None = None, @@ -740,9 +700,7 @@ def __init__( ) self.blocks = nn.ModuleList( [ - EdenLiveTransformerBlock( - inner_dim, ffn_dim, num_attention_heads, cross_attn_norm, eps, added_kv_proj_dim - ) + EdenLiveTransformerBlock(inner_dim, ffn_dim, num_attention_heads, eps, added_kv_proj_dim) for _ in range(num_layers) ] ) @@ -840,7 +798,6 @@ def forward( encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, proxy_features: torch.Tensor | None = None, - text_timeline: dict[str, Any] | None = None, causal_state: EdenLiveCausalState | None = None, cache_update: bool = False, return_dict: bool = True, @@ -863,9 +820,6 @@ def forward( proxy_features (`torch.Tensor`, *optional*): Required proxy features of shape `(batch_size, proxy_channels, num_frames, proxy_height, proxy_width)`. Spatial dimensions are interpolated to match the latent input. - text_timeline (`dict[str, Any]`, *optional*): - Bidirectional 5B prompt layout with `lengths`, a sequence of positive segment token counts per sample, - and `text_length`, the token count of each prompt in the concatenated text embedding bank. causal_state (`EdenLiveCausalState`, *optional*): KV storage created by `create_causal_state`. Use separate states for different guidance branches. cache_update (`bool`, defaults to `False`): @@ -906,20 +860,7 @@ def forward( num_tokens = num_frames * tokens_per_frame if num_tokens % cp_size: raise ValueError("The input token count must be divisible by the Ulysses degree.") - query_start = cp_rank * (num_tokens // cp_size) text_context_length = encoder_hidden_states.shape[1] - if text_timeline is not None: - if causal_state is not None or self.config.image_dim is not None: - raise ValueError("Temporal prompt banks are supported by bidirectional 5B inference only.") - lengths = text_timeline["lengths"] - text_length = int(text_timeline["text_length"]) - if len(lengths) != batch_size or text_length < 1: - raise ValueError("A timeline must contain one nonempty prompt bank per sample.") - if any(not sample or min(sample) < 1 or sum(sample) != num_tokens for sample in lengths): - raise ValueError("Each timeline must partition all query tokens into positive-length segments.") - if text_context_length != max(map(len, lengths)) * text_length: - raise ValueError("The text embedding bank does not match the timeline layout.") - key_rotary_emb = None cache_window_tokens = None cache_sink_tokens = 0 @@ -998,8 +939,6 @@ def forward( "cache_update": cache_update, "cache_window_tokens": cache_window_tokens, "cache_sink_tokens": cache_sink_tokens, - "text_timeline": text_timeline, - "query_start": query_start, "text_context_length": text_context_length, } if torch.is_grad_enabled() and self.gradient_checkpointing: diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ec5e94d32ddf..d0d8142573c1 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -274,7 +274,7 @@ ] _import_structure["diffusion_gemma"] = ["DiffusionGemmaPipeline", "DiffusionGemmaPipelineOutput"] _import_structure["dreamlite"] = ["DreamLitePipeline", "DreamLiteMobilePipeline", "DreamLitePipelineOutput"] - _import_structure["edenlive"] = ["EdenLivePipeline", "EdenLivePipelineOutput", "EdenLiveStreamState"] + _import_structure["edenlive"] = ["EdenLiveCausalPipeline", "EdenLivePipeline", "EdenLivePipelineOutput"] _import_structure["easyanimate"] = [ "EasyAnimatePipeline", "EasyAnimateInpaintPipeline", @@ -725,7 +725,7 @@ EasyAnimateInpaintPipeline, EasyAnimatePipeline, ) - from .edenlive import EdenLivePipeline, EdenLivePipelineOutput, EdenLiveStreamState + 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 index 72a32b470383..d54aa3d6ea93 100644 --- a/src/diffusers/pipelines/edenlive/__init__.py +++ b/src/diffusers/pipelines/edenlive/__init__.py @@ -23,7 +23,8 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["pipeline_edenlive"] = ["EdenLivePipeline", "EdenLiveStreamState"] + _import_structure["pipeline_edenlive"] = ["EdenLivePipeline"] + _import_structure["pipeline_edenlive_causal"] = ["EdenLiveCausalPipeline"] _import_structure["pipeline_output"] = ["EdenLivePipelineOutput"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -33,7 +34,8 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * else: - from .pipeline_edenlive import EdenLivePipeline, EdenLiveStreamState + from .pipeline_edenlive import EdenLivePipeline + from .pipeline_edenlive_causal import EdenLiveCausalPipeline from .pipeline_output import EdenLivePipelineOutput else: import sys diff --git a/src/diffusers/pipelines/edenlive/loading_utils.py b/src/diffusers/pipelines/edenlive/loading_utils.py deleted file mode 100644 index e273e2b507ef..000000000000 --- a/src/diffusers/pipelines/edenlive/loading_utils.py +++ /dev/null @@ -1,523 +0,0 @@ -# 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. - -"""Read the four original EdenLive releases without writing converted checkpoints.""" - -import copy -import json -from collections import Counter -from contextlib import nullcontext -from pathlib import Path, PurePosixPath - -import torch -from huggingface_hub import hf_hub_download, snapshot_download -from transformers import ( - AutoTokenizer, - CLIPImageProcessor, - CLIPVisionModel, - DINOv3ViTConfig, - DINOv3ViTModel, - UMT5EncoderModel, -) - -from ...loaders.single_file_utils import convert_wan_vae_to_diffusers -from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel -from ...models.model_loading_utils import load_state_dict -from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler -from ...utils import _resolve_dtype -from ...utils.hub_utils import _resolve_revision - - -_EDENLIVE_PROFILES = { - "5b-bidir": {"size": "5b", "causal": False, "transformer": "dit-5b-bidir", "resize": "bilinear"}, - "14b-bidir": {"size": "14b", "causal": False, "transformer": "dit-14b-bidir", "resize": "bicubic"}, - "5b-causal": {"size": "5b", "causal": True, "transformer": "dit-5b", "resize": "bilinear"}, - "14b-causal": {"size": "14b", "causal": True, "transformer": "dit-14b", "resize": "bicubic"}, -} - -_COMPONENT_CLASSES = { - "transformer": ("diffusers", "EdenLiveTransformer3DModel"), - "vae": ("diffusers", "AutoencoderKLWan"), - "decoder": ("diffusers", "EdenLiveLightTAEDecoder"), - "text_encoder": ("transformers", "UMT5EncoderModel"), - "tokenizer": ("transformers", "AutoTokenizer"), - "proxy_encoder": ("transformers", "DINOv3ViTModel"), - "proxy_projection": ("diffusers", "EdenLiveProxyProjection"), - "image_encoder": ("transformers", "CLIPVisionModel"), - "image_processor": ("transformers", "CLIPImageProcessor"), -} - -_HUB_OPTIONS = {"cache_dir", "force_download", "proxies", "local_files_only", "token", "revision"} - -_UNSUPPORTED_LOADING_OPTIONS = { - "device_map", - "max_memory", - "offload_folder", - "offload_state_dict", - "variant", - "quantization_config", - "custom_pipeline", - "custom_revision", - "use_onnx", - "use_flashpack", - "load_connected_pipeline", - "provider", - "sess_options", - "provider_options", -} - - -def read_checkpoint_config(source, loading_options): - options = loading_options.copy() - if Path(source).is_dir(): - config_file = Path(source) / "model_index.json" - else: - options["revision"] = _resolve_revision( - source, - **{ - key: value - for key, value in options.items() - if key in {"revision", "cache_dir", "local_files_only", "token"} - }, - ) - config_file = Path( - hf_hub_download( - str(source), - "model_index.json", - **{key: value for key, value in options.items() if key in _HUB_OPTIONS}, - ) - ) - with config_file.open() as handle: - config = json.load(handle) - if not isinstance(config, dict): - raise ValueError("EdenLive model_index.json must contain an object.") - return config, config_file, options - - -def normalize_saved_profile(config, requested_profile): - profile = config.get("profile") - if profile in ("5b", "14b"): - modes = {"EdenLivePipeline": "bidir", "EdenLiveChunkCausalPipeline": "causal"} - mode = modes.get(config.get("_class_name")) - if mode is None: - raise ValueError("Cannot infer an EdenLive profile from this legacy pipeline class.") - profile = f"{profile}-{mode}" - if profile not in _EDENLIVE_PROFILES: - raise ValueError(f"Saved EdenLive pipeline has an invalid profile: {profile!r}.") - if requested_profile is not None and requested_profile != profile: - raise ValueError(f"Requested profile {requested_profile!r} conflicts with saved profile {profile!r}.") - return profile - - -def load_original_pipeline(pipeline_class, source, config_file, config, profile, loading_options): - component_names, option_names = pipeline_class._get_signature_keys(pipeline_class) - passed_components = {name: loading_options[name] for name in component_names if name in loading_options} - context = resolve_original_checkpoint( - config, profile=profile, passed_components=passed_components, loading_options=loading_options - ) - supported_options = ( - _HUB_OPTIONS - | _UNSUPPORTED_LOADING_OPTIONS - | { - "dtype", - "torch_dtype", - "low_cpu_mem_usage", - "disable_mmap", - "use_safetensors", - "trust_remote_code", - } - ) - unknown = set(loading_options) - set(component_names) - set(option_names) - supported_options - if unknown: - raise ValueError(f"Original EdenLive loading does not support {', '.join(sorted(unknown))}.") - dtype = _resolve_dtype(loading_options.get("dtype"), loading_options.get("torch_dtype")) - if dtype is None: - dtype = context["default_dtype"] - if not isinstance(dtype, (dict, torch.dtype)) or ( - isinstance(dtype, dict) and any(not isinstance(value, torch.dtype) for value in dtype.values()) - ): - raise ValueError("`dtype` must be a torch.dtype or a mapping of component names to torch.dtype values.") - folder = config_file.parent - if not Path(source).is_dir() and ( - loading_options.get("force_download", False) - or not all((folder / filename).is_file() for filename in context["files"]) - ): - folder = Path( - snapshot_download( - str(source), - allow_patterns=context["files"], - user_agent={"pipeline_class": pipeline_class.__name__}, - **{key: value for key, value in loading_options.items() if key in _HUB_OPTIONS}, - ) - ) - missing = [filename for filename in context["files"] if not (folder / filename).is_file()] - if missing: - raise FileNotFoundError(f"Original checkpoint files are missing from {folder}: {missing}") - init_kwargs = {name: context["config"][name] for name in option_names if name in context["config"]} - init_kwargs.update({name: loading_options[name] for name in option_names if name in loading_options}) - try: - for name in component_names: - if name in passed_components: - init_kwargs[name] = passed_components[name] - elif context["config"][name][0] is None: - init_kwargs[name] = None - else: - component_dtype = ( - dtype.get(name, dtype.get("default", torch.float32)) if isinstance(dtype, dict) else dtype - ) - init_kwargs[name] = load_original_component( - name, cached_folder=folder, context=context, dtype=component_dtype, loading_options=loading_options - ) - pipeline = pipeline_class(**init_kwargs) - finally: - context["state_dicts"].clear() - pipeline.register_to_config(_name_or_path=str(source)) - return pipeline - - -def _relative_path(value): - if not isinstance(value, str) or not value or "\\" in value or any(c in value for c in "*?["): - raise ValueError("EdenLive path overrides must be literal repository-relative paths.") - path = PurePosixPath(value) - if path.is_absolute() or ".." in path.parts or not path.parts: - raise ValueError(f"EdenLive path must stay within the checkpoint repository: {value!r}.") - return path.as_posix() - - -def _component_sources(profile_name, path_overrides): - profile = _EDENLIVE_PROFILES[profile_name] - size, transformer = profile["size"], profile["transformer"] - sources = { - "transformer": { - "config": f"{transformer}/config.json", - "weights": f"{transformer}/diffusion_pytorch_model.safetensors", - }, - "vae": { - "config": f"vae-{size}/config.json", - "weights": f"vae-{size}/Wan2.{'2' if size == '5b' else '1'}_VAE.pth", - }, - "decoder": {"weights": f"decoder-{size}/lighttaew2_{'2' if size == '5b' else '1'}.pth"}, - "text_encoder": {"config": "umt5/text_encoder/config.json", "weights": "umt5/text_encoder/model.safetensors"}, - "tokenizer": {"directory": "umt5/tokenizer"}, - "proxy_encoder": { - "config": "vision-encoder/backbone/config.json", - "weights": "vision-encoder/projection.safetensors", - }, - "proxy_projection": {"weights": "vision-encoder/projection.safetensors"}, - "image_encoder": { - "config": "clip/image_encoder/config.json", - "weights": "clip/image_encoder/model.safetensors", - }, - "image_processor": {"directory": "clip/image_processor"}, - } - if not isinstance(path_overrides, dict) or set(path_overrides) - set(sources): - raise ValueError("`_path_overrides` must map EdenLive component names to paths.") - for name, override in path_overrides.items(): - if not isinstance(override, dict): - raise ValueError(f"Path overrides for {name} must be an object.") - # Only transformers vary by full profile; VAE and decoder files vary by size. - selectors = ( - set(_EDENLIVE_PROFILES) - if name == "transformer" - else {"5b", "14b"} - if name in ("vae", "decoder") - else set() - ) - if selectors and set(override) <= selectors: - selector = profile_name if name == "transformer" else size - for value in override.values(): - if not isinstance(value, dict) or set(value) - (set(sources[name]) | {"directory"}): - raise ValueError(f"Invalid path fields for {name}.") - for path in value.values(): - _relative_path(path) - override = override.get(selector, {}) - if set(override) - (set(sources[name]) | {"directory"}): - raise ValueError(f"Invalid path fields for {name}: {sorted(override)}.") - override = {key: _relative_path(value) for key, value in override.items()} - if "directory" in override and "directory" not in sources[name]: - directory = override.pop("directory") - sources[name] = {key: f"{directory}/{PurePosixPath(value).name}" for key, value in sources[name].items()} - sources[name].update(override) - if size == "5b": - sources.pop("image_encoder") - sources.pop("image_processor") - return sources - - -def _source_files(name, source): - if name == "tokenizer": - return [ - f"{source['directory']}/{filename}" - for filename in ("tokenizer_config.json", "special_tokens_map.json", "tokenizer.json", "spiece.model") - ] - if name == "image_processor": - return [f"{source['directory']}/preprocessor_config.json"] - return list(source.values()) - - -def _read_config(folder, source): - with open(Path(folder) / source["config"]) as handle: - config = json.load(handle) - if not isinstance(config, dict): - raise ValueError(f"Expected an object in {source['config']}.") - return config - - -def _read_component_state(folder, name, context, disable_mmap): - filename = context["sources"][name]["weights"] - state = context["state_dicts"].get(filename) - if state is None: - state = load_state_dict(str(Path(folder) / filename), disable_mmap=disable_mmap, map_location="cpu") - if not isinstance(state, dict) or any( - not isinstance(k, str) or not isinstance(v, torch.Tensor) for k, v in state.items() - ): - raise ValueError(f"Expected a tensor state dictionary in {filename}.") - context["source_users"][filename] -= 1 - if context["source_users"][filename] > 0: - context["state_dicts"][filename] = state - else: - context["state_dicts"].pop(filename, None) - return state - - -def _strict_assign(model, state, dtype): - model.load_state_dict(state, strict=True, assign=True) - return model.to(dtype=dtype).eval().requires_grad_(False) - - -def _empty_weights(loading_options): - return torch.device("meta") if loading_options.get("low_cpu_mem_usage", True) else nullcontext() - - -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 _load_transformers_component(model_class, config, state, dtype, loading_options, allowed_unused=()): - state = {key: value for key, value in state.items() if key not in allowed_unused} - model, info = model_class.from_pretrained( - None, - config=model_class.config_class.from_dict(config), - state_dict=state, - dtype=dtype, - local_files_only=True, - low_cpu_mem_usage=loading_options.get("low_cpu_mem_usage", True), - output_loading_info=True, - ) - if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): - raise ValueError(f"{model_class.__name__} checkpoint did not load strictly: {info}.") - return model.eval().requires_grad_(False) - - -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 resolve_original_checkpoint(config_dict, *, profile, passed_components, loading_options): - if "_edenlive_original" not in config_dict: - return None - original = config_dict["_edenlive_original"] - if not isinstance(original, dict) or set(original) - {"_path_overrides"}: - raise ValueError("`_edenlive_original` must be an object with optional `_path_overrides`.") - if profile not in _EDENLIVE_PROFILES: - raise ValueError( - f"Original EdenLive checkpoints require `profile` in {tuple(_EDENLIVE_PROFILES)}; got {profile!r}." - ) - preset = _EDENLIVE_PROFILES[profile] - unsupported = { - key: loading_options.get(key) - for key in _UNSUPPORTED_LOADING_OPTIONS - if loading_options.get(key) is not None and loading_options.get(key) is not False - } - if unsupported: - raise ValueError( - f"Original EdenLive loading does not support {', '.join(sorted(unsupported))}; load first and place the pipeline explicitly." - ) - sources = _component_sources(profile, original.get("_path_overrides", {})) - sources = {name: source for name, source in sources.items() if name not in passed_components} - safe_only = loading_options.get("use_safetensors") - weights = [source["weights"] for source in sources.values() if "weights" in source] - if safe_only is True and any(not path.endswith(".safetensors") for path in weights): - raise ValueError( - "Original EdenLive includes PyTorch .pth components; omit `use_safetensors=True` or provide those components explicitly." - ) - if safe_only is False and any(path.endswith(".safetensors") for path in weights): - raise ValueError( - "Original EdenLive includes Safetensors components; omit `use_safetensors=False` or provide those components explicitly." - ) - config = copy.deepcopy(config_dict) - config.pop("_edenlive_original") - config.update( - _class_name="EdenLivePipeline", - profile=profile, - proxy_resize_mode=preset["resize"], - proxy_unconditional="same", - proxy_projection_autocast=True, - proxy_image_size=(480, 832), - ) - config.update(_COMPONENT_CLASSES) - if preset["size"] == "5b": - config.update(image_encoder=(None, None), image_processor=(None, None)) - scheduler_name = "UniPCMultistepScheduler" if profile == "5b-bidir" else "EdenLiveFlowMatchEulerDiscreteScheduler" - config["scheduler"] = ("diffusers", scheduler_name) - if preset["causal"]: - config.update(chunk_size=3, max_context_frames=18, sink_frames=6) - return { - "config": config, - "profile": profile, - "preset": preset, - "sources": sources, - "files": sorted( - {"model_index.json"} - | {filename for name, source in sources.items() for filename in _source_files(name, source)} - ), - "default_dtype": { - "default": torch.bfloat16, - "proxy_encoder": torch.float32, - "proxy_projection": torch.float32, - }, - "state_dicts": {}, - "source_users": Counter(source["weights"] for source in sources.values() if "weights" in source), - } - - -def load_original_component(name, *, cached_folder, context, dtype, loading_options): - preset = context["preset"] - if name == "scheduler": - if context["profile"] == "5b-bidir": - return UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5.0) - if preset["causal"]: - timesteps = [1000, 937.5, 833.3333333333, 625] if preset["size"] == "5b" else [1000, 900, 750, 500] - return EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True, default_timesteps=timesteps) - return EdenLiveFlowMatchEulerDiscreteScheduler(shift=3.0) - source = context["sources"][name] - if name in ("tokenizer", "image_processor"): - model_class = AutoTokenizer if name == "tokenizer" else CLIPImageProcessor - return model_class.from_pretrained(Path(cached_folder) / source["directory"], local_files_only=True) - state = _read_component_state(cached_folder, name, context, loading_options.get("disable_mmap", False)) - config = _read_config(cached_folder, source) if "config" in source else None - if name == "transformer": - proxy_channels = state["patch_embedding.weight"].shape[1] - config["in_channels"] - expected_proxy_channels = config.get("proxy_channels", 1024) - condition_rmsnorm = "patch_embedding.condition_norm.weight" in state - if proxy_channels != expected_proxy_channels or condition_rmsnorm != config.get( - "condition_rmsnorm", preset["size"] == "14b" - ): - raise ValueError( - "Transformer proxy channels or normalization do not match its configuration and selected profile." - ) - with _empty_weights(loading_options): - model = EdenLiveTransformer3DModel.from_config( - config, proxy_channels=proxy_channels, condition_rmsnorm=condition_rmsnorm - ) - # These rotary buffers are derived, not serialized checkpoint tensors. - 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)) - return _strict_assign(model, state, dtype) - if name == "vae": - mapped = {} - for key, value in state.items(): - target = _remap_wan22_vae_key(key) - if target in mapped: - raise ValueError(f"Multiple VAE keys map to {target}.") - mapped[target] = value - mapped = convert_wan_vae_to_diffusers(mapped) - if len(mapped) != len(state): - raise ValueError("Multiple VAE checkpoint keys map to the same parameter.") - # Wan 2.2 stores folded RGB convolutions; the current native class - # takes their post-patchify/pre-unpatchify channel counts. - patch_size = config.get("patch_size") or 1 - for field, key, dimension in ( - ("in_channels", "encoder.conv_in.weight", 1), - ("out_channels", "decoder.conv_out.weight", 0), - ): - channels = mapped[key].shape[dimension] - declared = config.get(field, 3) - if channels not in (declared, declared * patch_size**2): - raise ValueError(f"VAE {field} does not match the checkpoint convolutions.") - config[field] = channels - with _empty_weights(loading_options): - model = AutoencoderKLWan.from_config(config) - return _strict_assign(model, mapped, dtype) - if name == "decoder": - with _empty_weights(loading_options): - model = EdenLiveLightTAEDecoder( - latent_channels=48 if preset["size"] == "5b" else 16, patch_size=2 if preset["size"] == "5b" else 1 - ) - return _strict_assign(model, model._convert_original_state_dict(state), dtype) - if name in ("proxy_encoder", "proxy_projection"): - state = _proxy_state(state) - if name == "proxy_encoder": - config = DINOv3ViTConfig.from_dict(config) - with _empty_weights(loading_options): - model = DINOv3ViTModel(config) - expected = set(model.state_dict()) - mapped = {} - for key, value in state.items(): - if key.startswith("backbone."): - target = key.removeprefix("backbone.") - if target not in expected and f"model.{target}" in expected: - target = f"model.{target}" - mapped[target] = value - model.rope_embeddings = type(model.rope_embeddings)(config) - else: - mapped = {key: value for key, value in state.items() if key.startswith("patch_projection.")} - weight = mapped["patch_projection.weight"] - with _empty_weights(loading_options): - model = EdenLiveProxyProjection(in_channels=weight.shape[1], out_channels=weight.shape[0]) - return _strict_assign(model, mapped, dtype) - if name == "text_encoder": - return _load_transformers_component(UMT5EncoderModel, config, state, dtype, loading_options) - if name == "image_encoder": - return _load_transformers_component( - CLIPVisionModel, config, state, dtype, loading_options, ("visual_projection.weight",) - ) - raise ValueError(f"Unknown EdenLive component: {name}.") diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py index b9647cf796c6..3c34cfab06d9 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -12,184 +12,77 @@ # See the License for the specific language governing permissions and # limitations under the License. +import html import math -import weakref -from dataclasses import dataclass, field +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, EdenLiveProxyProjection, EdenLiveTransformer3DModel -from ...models.autoencoders.autoencoder_edenlive import EdenLiveLightTAEDecodeCache, EdenLiveLightTAEDecoder -from ...models.autoencoders.autoencoder_kl_wan import WanCausalConv3d -from ...models.autoencoders.vae import DiagonalGaussianDistribution -from ...models.transformers.transformer_edenlive import EdenLiveCausalState +from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler -from ...utils import replace_example_docstring -from ...utils.accelerate_utils import apply_forward_hook +from ...utils import is_ftfy_available, replace_example_docstring +from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from . import pipeline_utils as utils -from .loading_utils import load_original_pipeline, normalize_saved_profile, read_checkpoint_config from .pipeline_output import EdenLivePipelineOutput +from .pipeline_utils import prepare_model_input, sample_noise +from .video_processor import EdenLiveVideoProcessor -# docstyle-ignore: Keep the continuation prompt on the try/finally example. +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", profile="5b-causal") - >>> pipe.to("cuda") - >>> image = load_image("reference.png") - >>> video = load_video("proxy-24fps.mp4") # At least 49 RGB frames sampled at 24 fps. - >>> prompt = "A character walks through snow." + >>> pipe = EdenLivePipeline.from_pretrained( + ... "VAST-AI/EdenLive-5B-Diffusers", + ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, + ... ).to("cuda") >>> frames = pipe( - ... image=image, - ... video=video, - ... prompt=prompt, - ... num_frames=49, - ... generator=torch.Generator().manual_seed(42), + ... 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) - - >>> # Generate successive chunks from three sampled RGB images per call. - >>> proxy_paths = [ - ... ["proxy/chunk-0-0.png", "proxy/chunk-0-1.png", "proxy/chunk-0-2.png"], - ... ["proxy/chunk-1-0.png", "proxy/chunk-1-1.png", "proxy/chunk-1-2.png"], - ... ] - >>> prompts = [prompt, "The character approaches a campfire."] - >>> state = pipe.create_stream_state( - ... image=image, prompt=prompt, generator=torch.Generator().manual_seed(42), output_type="np" - ... ) - >>> try: - ... for index, (paths, current_prompt) in enumerate(zip(proxy_paths, prompts)): - ... chunk = pipe( - ... stream_state=state, - ... proxy_frames=[load_image(path) for path in paths], - ... prompt=current_prompt, - ... ).frames[0] - ... export_to_video(chunk, f"chunk-{index:03d}.mp4", fps=24) - ... finally: - ... pipe.close_stream(state) ``` """ -@dataclass -class _EdenLiveConditionEncoderCache: - feat_map: list[torch.Tensor | None] | None = None - - def reset(self): - self.feat_map = None - - -@apply_forward_hook -def _encode_condition_chunk(vae, pixels, cache): - if vae.use_slicing or vae.use_tiling: - raise ValueError("Streaming condition encoding does not support VAE slicing or tiling.") - if vae.config.patch_size is not None: - raise ValueError("Streaming 14B conditioning requires an unpatchified Wan VAE.") - first = cache.feat_map is None - if pixels.shape[2] < 1 or (pixels.shape[2] - int(first)) % 4: - raise ValueError("Condition encoding requires 1 + 4n frames initially and 4n frames thereafter.") - if first: - cache.feat_map = [None for module in vae.encoder.modules() if isinstance(module, WanCausalConv3d)] - encoded = [] - if first: - encoded.append(vae.encoder(pixels[:, :, :1], feat_cache=cache.feat_map, feat_idx=[0])) - pixels = pixels[:, :, 1:] - for chunk in pixels.split(4, dim=2): - if chunk.shape[2]: - encoded.append(vae.encoder(chunk, feat_cache=cache.feat_map, feat_idx=[0])) - moments = vae.quant_conv(torch.cat(encoded, dim=2)) - return DiagonalGaussianDistribution(moments).mode() - - -@dataclass -class EdenLiveStreamState: - """Mutable, single-video rollout owned by one pipeline. Create with `pipe.create_stream_state()`. - - The attention states own the only time cursor (`next_frame`). Encoder and decoder caches retain temporal history - independently of the attention window. Calls sharing a pipeline must be serial, including calls for different - streams. Close the state in a `finally` block to release its tensors. States are runtime objects and are not - serialized. - """ - - causal_states: list[EdenLiveCausalState] - condition_encoder_cache: _EdenLiveConditionEncoderCache | None - decoder_cache: EdenLiveLightTAEDecodeCache | None - image: torch.Tensor | None - first_frame: torch.Tensor | None - image_embeds: torch.Tensor | None - encoded_prompt: str | None - prompt_embeds: torch.Tensor | None - negative_prompt_embeds: torch.Tensor | None - generator: torch.Generator | None - profile: str - height: int - width: int - guidance_scale: float - num_inference_steps: int - max_sequence_length: int - output_type: str - device: torch.device - dtype: torch.dtype - _pipeline: weakref.ReferenceType = field(repr=False) - _component_ids: tuple = field(repr=False) - valid: bool = True - closed: bool = False - - class EdenLivePipeline(DiffusionPipeline): - """Bidirectional and causal video generation conditioned on a reference image, text, and an RGB proxy video. - - Load released weights with [`EdenLivePipeline.from_pretrained`]. The constructor accepts initialized components and - does not read checkpoints. The full `profile` selects the model size and sampling path. + """Bidirectional video generation conditioned on a reference image, text, and an RGB proxy video. Args: vae ([`AutoencoderKLWan`]): - Wan VAE for reference-image encoding and, when `decoder=None`, video decoding. + Wan VAE for reference encoding and full-quality video decoding. text_encoder (`UMT5EncoderModel`): - UMT5 text encoder. + Text encoder. tokenizer (`AutoTokenizer`): Tokenizer corresponding to the text encoder. transformer ([`EdenLiveTransformer3DModel`]): - Video denoiser with the weights and configuration for the selected profile. + EdenLive bidirectional video denoiser. scheduler ([`EdenLiveFlowMatchEulerDiscreteScheduler`] or [`UniPCMultistepScheduler`]): - Sampling schedule compatible with the transformer. The 5B bidirectional release uses UniPC; the other - profiles use the EdenLive Euler scheduler, with fixed stochastic schedules for causal checkpoints. + Scheduler corresponding to the transformer checkpoint. proxy_encoder (`DINOv3ViTModel`): - Trained DINOv3 encoder for the RGB proxy video. + Trained RGB proxy encoder. proxy_projection ([`EdenLiveProxyProjection`]): - Projection from DINO features to transformer conditioning features. + Learned projection of DINO patch features. image_encoder (`CLIPVisionModel`, *optional*): - CLIP reference-image encoder required by 14B profiles. + Reference image encoder required by the 14B model. image_processor (`CLIPImageProcessor`, *optional*): - Image processor corresponding to the CLIP encoder. - profile (`str`, defaults to `"5b-bidir"`): - One of `"5b-bidir"`, `"14b-bidir"`, `"5b-causal"`, or `"14b-causal"`. - proxy_resize_mode (`str`, defaults to `"bilinear"`): - RGB proxy preprocessing before DINO encoding: `"bilinear"`, `"bicubic"`, or `"bicubic_center_crop"`. - proxy_unconditional (`str`, defaults to `"same"`): - Use the same proxy features (`"same"`) or zeros (`"zero"`) in the negative guidance branch. - proxy_projection_autocast (`bool`, defaults to `True`): - Allow autocast during proxy projection while retaining FP32 projection parameters. - proxy_image_size (`tuple[int, int]`, defaults to `(480, 832)`): - Height and width used to preprocess RGB frames for the proxy encoder. - chunk_size (`int`, defaults to 3): - Number of latent frames per causal chunk. A final chunk may be shorter. - max_context_frames (`int`, defaults to 18): - Maximum causal attention context in latent frames, including sink frames and the current chunk. - sink_frames (`int`, defaults to 6): - Initial latent frames retained as attention sinks when the context window advances. - decoder (`torch.nn.Module`, *optional*): - Optional video decoder compatible with the VAE latent channels and temporal/spatial scales. `None` uses the - VAE. Original-release loading supplies LightTAE by default; pass `decoder=None` when loading to select the - full VAE. Use a compatible `ModelMixin` decoder for pipeline device placement and serialization. + 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->proxy_projection->transformer->vae->decoder" @@ -207,129 +100,109 @@ def __init__( proxy_projection: EdenLiveProxyProjection, image_encoder: CLIPVisionModel | None = None, image_processor: CLIPImageProcessor | None = None, - profile: str = "5b-bidir", - proxy_resize_mode: str = "bilinear", - proxy_unconditional: str = "same", - proxy_projection_autocast: bool = True, - proxy_image_size: tuple | list = (480, 832), - chunk_size: int = 3, - max_context_frames: int = 18, - sink_frames: int = 6, - decoder: torch.nn.Module | None = None, + decoder: EdenLiveLightTAEDecoder | None = None, + proxy_image_height: int = 480, + proxy_image_width: int = 832, ): super().__init__() - 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.") - utils.initialize_pipeline( - self, - vae, - text_encoder, - tokenizer, - transformer, - scheduler, - proxy_encoder, - proxy_projection, - image_encoder, - image_processor, - decoder, - profile, - proxy_resize_mode, - proxy_unconditional, - proxy_projection_autocast, - proxy_image_size, + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + proxy_encoder=proxy_encoder, + proxy_projection=proxy_projection, + 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 ) - self.register_to_config(chunk_size=chunk_size, max_context_frames=max_context_frames, sink_frames=sink_frames) - - @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, *, profile=None, **kwargs): - """Load an original EdenLive release or a standard saved pipeline. - - Original releases are identified by the EdenLive metadata in their root `model_index.json`. Only files needed - by the selected profile and components are loaded; weights are mapped in memory without conversion copies. - Standard snapshots use the inherited Diffusers component loader after normalizing legacy profile metadata. - - Args: - pretrained_model_name_or_path (`str` or `os.PathLike`): - Hub repository ID or local directory containing an original release or saved Diffusers pipeline. - profile (`str`, *optional*): - Required for original releases: `"5b-bidir"`, `"14b-bidir"`, `"5b-causal"`, or `"14b-causal"`. Standard - snapshots retain their saved profile; a conflicting explicit value raises an error. - dtype (`torch.dtype` or `dict[str, torch.dtype]`, *optional*): - Loading dtype, or a component-name mapping with an optional `"default"` entry. Original releases - default to BF16 for the transformer, VAE, text encoder, CLIP, and LightTAE, and FP32 for DINO and its - projection. `torch_dtype` is also accepted. Supplied component instances retain their own dtype. - decoder (`torch.nn.Module`, *optional*): - Original releases use LightTAE when omitted. Explicit `None` selects the full VAE. A compatible - initialized decoder replaces the saved or released decoder without changing its identity or dtype. - Standard snapshots use their saved decoder when this argument is omitted. - revision (`str`, *optional*): - Hub branch, tag, or commit to resolve before selecting component files. - cache_dir (`str` or `os.PathLike`, *optional*): - Hub cache directory. - local_files_only (`bool`, defaults to `False`): - Load local files or cached Hub files without downloading. - token (`str` or `bool`, *optional*): - Authentication token or Hub token-selection option. - kwargs (`dict`, *optional*): - Initialized component overrides, such as `transformer`, `vae`, or `text_encoder`, and loading options - supported by the checkpoint format. Original loading does not support active `device_map`, `variant`, - or quantization options. Load on CPU and move the pipeline to the desired device afterward. - - Returns: - [`EdenLivePipeline`]: - Pipeline with a canonical full profile and the selected components. Older standard snapshots with - size-only profiles or the former causal class name are normalized in memory. - """ - config, config_file, options = read_checkpoint_config(pretrained_model_name_or_path, kwargs) - if "_edenlive_original" in config: - return load_original_pipeline(cls, pretrained_model_name_or_path, config_file, config, profile, options) - profile = normalize_saved_profile(config, profile) - decoder = options.get("decoder") - if decoder is not None: - options["decoder"] = None - pipeline = super().from_pretrained(pretrained_model_name_or_path, profile=profile, **options) - if decoder is not None: - utils.validate_decoder(decoder, pipeline.vae) - pipeline.register_modules(decoder=decoder) - return pipeline + 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) - @torch.no_grad() - def create_stream_state( + def encode_prompt( self, - image, prompt=None, - *, - prompt_embeds=None, negative_prompt=None, + do_classifier_free_guidance=True, + prompt_embeds=None, negative_prompt_embeds=None, - height=None, - width=None, - guidance_scale=1.0, - num_inference_steps=4, - generator=None, max_sequence_length=512, - output_type="np", - ) -> EdenLiveStreamState: - """Initialize one causal rollout with fixed image, resolution, guidance, RNG and output format. + 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 - Uses the corresponding offline call options. Generate each chunk with `self(stream_state=state, - proxy_frames=frames)`. RGB output requires LightTAE. Close the state in a `finally` block. - """ - if not self.is_causal or self.config.chunk_size != 3: - raise ValueError("Streaming requires a causal profile with chunk_size=3.") - if output_type != "latent" and not isinstance(self.decoder, EdenLiveLightTAEDecoder): - raise ValueError( - "Streaming RGB output requires EdenLiveLightTAEDecoder; otherwise use output_type='latent'." - ) - if self.model_size == "14b" and (self.vae.use_slicing or self.vae.use_tiling): - raise ValueError("Streaming 14B condition encoding does not support VAE slicing or tiling.") - height = (704 if self.model_size == "5b" else 480) if height is None else height - width = (1280 if self.model_size == "5b" else 832) if width is None else width - if not isinstance(height, int) or not isinstance(width, int): - raise ValueError("Height and width must be integers.") - if image is None: - raise ValueError("A reference image is required.") + 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 @@ -342,294 +215,107 @@ def create_stream_state( ) if output_type not in ("np", "pt", "pil", "latent"): raise ValueError("`output_type` must be 'np', 'pt', 'pil', or 'latent'.") - if self.model_size == "14b" and self.vae.config.patch_size is not None: - raise ValueError("Streaming 14B conditioning requires an unpatchified Wan VAE.") - self._check_prompt(prompt, prompt_embeds, required=True) - self._check_prompt(negative_prompt, negative_prompt_embeds) - if not isinstance(max_sequence_length, int) or max_sequence_length < 1: - raise ValueError("`max_sequence_length` must be a positive integer.") - if not isinstance(num_inference_steps, int) or num_inference_steps < 1: - raise ValueError("`num_inference_steps` must be a positive integer.") - if not math.isfinite(guidance_scale): - raise ValueError("`guidance_scale` must be finite.") - if generator is not None and not isinstance(generator, torch.Generator): - raise ValueError("Streaming requires one torch.Generator, not a generator list.") - device, dtype = self._execution_device, self.transformer.dtype - reference = self.proxy_video_processor.preprocess(image) - if ( - reference.ndim != 4 - or reference.shape[:2] != (1, 3) - or reference.numel() == 0 - or not torch.isfinite(reference).all() - or reference.min() < 0 - or reference.max() > 1 - ): - raise ValueError("The reference image must be one finite RGB image in [0, 1].") - self.scheduler.set_timesteps(num_inference_steps, device=device) - pixels = utils.prepare_reference_pixels(self, reference, height, width, 1, device) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_prompt, - guidance_scale > 1, - prompt_embeds, - negative_prompt_embeds, - max_sequence_length, - device, - dtype, - ) - image_embeds = utils.encode_image_embeds(self, pixels, device, dtype) - first_frame = self.vae.encode(pixels.unsqueeze(2).to(self.vae.dtype)).latent_dist.mode() - mean = first_frame.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) - inverse_std = 1 / first_frame.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) - first_frame = ((first_frame - mean) * inverse_std).to(dtype) - states = [ - self.transformer.create_causal_state( - batch_size=1, - height=first_frame.shape[3], - width=first_frame.shape[4], - chunk_size=3, - max_context_frames=self.config.max_context_frames, - sink_frames=self.config.sink_frames, - device=device, - dtype=dtype, - ) - for _ in range(2 if guidance_scale > 1 else 1) - ] - if generator is None: - generator = torch.Generator(device="cpu") - generator.seed() - return EdenLiveStreamState( - causal_states=states, - condition_encoder_cache=_EdenLiveConditionEncoderCache() if self.model_size == "14b" else None, - decoder_cache=EdenLiveLightTAEDecodeCache() if output_type != "latent" else None, - image=pixels if self.model_size == "14b" else None, - first_frame=first_frame, - image_embeds=image_embeds, - encoded_prompt=prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generator=generator, - profile=self.config.profile, - height=height, - width=width, - guidance_scale=guidance_scale, - num_inference_steps=num_inference_steps, - max_sequence_length=max_sequence_length, - output_type=output_type, - device=device, - dtype=dtype, - _pipeline=weakref.ref(self), - _component_ids=tuple(id(component) for component in self.components.values()), - ) - - def close_stream(self, state): - """Release a stream's tensors and model hooks. Repeated closing has no effect.""" - self._check_state(state) - if state.closed: - return - for causal_state in state.causal_states: - causal_state.reset() - state.causal_states.clear() - for cache in (state.condition_encoder_cache, state.decoder_cache): - if cache is not None: - cache.reset() - state.condition_encoder_cache = state.decoder_cache = None - state.image = state.first_frame = state.image_embeds = None - state.prompt_embeds = state.negative_prompt_embeds = None - state.encoded_prompt = state.generator = None - state.valid, state.closed = False, True - self.maybe_free_model_hooks() - - def _check_prompt(self, prompt, embeds, *, required=False): - if prompt is not None and embeds is not None: - raise ValueError("Pass either `prompt` or `prompt_embeds`, not both.") - if prompt is not None and not isinstance(prompt, str): - raise ValueError("Streaming requires a single string prompt.") - if embeds is not None and ( - not isinstance(embeds, torch.Tensor) - or embeds.ndim != 3 - or embeds.shape[0] != 1 - or embeds.shape[1] < 1 - or embeds.shape[2] != self.transformer.config.text_dim - or not torch.is_floating_point(embeds) - or not torch.isfinite(embeds).all() - ): - raise ValueError("Prompt embeddings must be finite floating-point tensors of shape (1, tokens, text_dim).") - if required and prompt is None and embeds is None: - raise ValueError("Provide `prompt` or `prompt_embeds`.") - - def _prepare_condition_chunk(self, state): - if state.condition_encoder_cache is None: - return state.first_frame - first = state.causal_states[0].next_frame == 0 - pixels = state.first_frame.new_zeros(1, 3, 8 if first else 12, state.height, state.width) - if first: - pixels = torch.cat([state.image.unsqueeze(2), pixels], dim=2) - condition_chunk = _encode_condition_chunk(self.vae, pixels.to(self.vae.dtype), state.condition_encoder_cache) - mean = condition_chunk.new_tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1) - inverse_std = 1 / condition_chunk.new_tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1) - condition_chunk = ((condition_chunk - mean) * inverse_std).to(state.dtype) - mask = condition_chunk.new_zeros(1, self.vae_scale_factor_temporal, *condition_chunk.shape[2:]) - if first: - condition_chunk[:, :, :1] = state.first_frame - mask[:, :, 0] = 1 - state.image = None - return torch.cat([mask, condition_chunk], dim=1) - - def _update_prompt_embeds(self, state, prompt, prompt_embeds): - if prompt_embeds is not None: - state.prompt_embeds = prompt_embeds.to(device=state.device, dtype=state.dtype) - state.encoded_prompt = None - elif prompt is not None and prompt != state.encoded_prompt: - state.prompt_embeds, _ = self.encode_prompt( - prompt, - do_classifier_free_guidance=False, - max_sequence_length=state.max_sequence_length, - device=state.device, - dtype=state.dtype, - ) - state.encoded_prompt = prompt + if not set(callback_inputs).issubset(self._callback_tensor_inputs): + raise ValueError(f"Callback tensor inputs must belong to {self._callback_tensor_inputs}.") - def _check_state(self, state): - if not isinstance(state, EdenLiveStreamState) or state._pipeline() is not self: - raise ValueError("The stream must belong to this pipeline.") - - def _stream_step(self, state, proxy_frames, *, prompt=None, prompt_embeds=None, latents=None, return_dict=True): - self._check_state(state) - if state.closed or not state.valid: - raise ValueError("This stream is closed or invalid; create a new stream state.") - if tuple(id(component) for component in self.components.values()) != state._component_ids: - raise ValueError("Pipeline components must remain fixed throughout a stream.") - if ( - self.config.profile != state.profile - or self._execution_device != state.device - or self.transformer.dtype != state.dtype - ): - raise ValueError("A stream requires a fixed model, device, and dtype.") - self._check_prompt(prompt, prompt_embeds) + 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 = ( - 1, + batch_size, self.vae.config.z_dim, - 3, - state.height // self.vae_scale_factor_spatial, - state.width // self.vae_scale_factor_spatial, + (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 and ( - not isinstance(latents, torch.Tensor) - or tuple(latents.shape) != shape - or not torch.is_floating_point(latents) - or not torch.isfinite(latents).all() - ): - raise ValueError(f"Initial latents must be a finite floating-point tensor with shape {shape}.") - proxy_video = self.proxy_video_processor.preprocess_proxy_chunk(proxy_frames) - self._guidance_scale = state.guidance_scale - self._num_timesteps = state.num_inference_steps - try: - self._update_prompt_embeds(state, prompt, prompt_embeds) - proxy_features = utils.encode_proxy_features(self, proxy_video, 1, state.device) - condition_chunk = self._prepare_condition_chunk(state) - latents = ( - utils.sample_noise(self, shape, state.generator, state.device, state.dtype) - if latents is None - else latents.to(device=state.device, dtype=state.dtype).clone() + 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 sample_noise(self, shape, generator, device, 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) ) - latents, state.prompt_embeds, state.negative_prompt_embeds = self._denoise_chunk( - latents, - condition_chunk, - state.first_frame, - proxy_features, - state.image_embeds, - state.prompt_embeds, - state.negative_prompt_embeds, - state.causal_states, - state.num_inference_steps, - state.generator, - None, - None, - [], - 0, - condition_start=0, + if self._image_conditioned: + condition_video = torch.cat( + [pixels.unsqueeze(2), pixels.new_zeros(batch_size, 3, num_frames - 1, height, width)], dim=2 ) - frames = utils.decode_latents(self, latents, state.output_type, decoder_cache=state.decoder_cache) - return EdenLivePipelineOutput(frames=frames) if return_dict else (frames,) - except Exception: - state.valid = False - raise - - 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 prompts with the native UMT5 encoder. - - Args: - prompt (`str` or `list[str]`, *optional*): - Positive text, unless `prompt_embeds` is supplied. - negative_prompt (`str` or `list[str]`, *optional*): - Negative text for classifier-free guidance; defaults to empty strings. - do_classifier_free_guidance (`bool`, defaults to `True`): - Whether to produce negative embeddings. - prompt_embeds (`torch.Tensor`, *optional*): - Precomputed positive embeddings. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Precomputed negative embeddings. - max_sequence_length (`int`, defaults to 512): - Fixed token count per prompt. - device (`torch.device`, *optional*): - Device for encoding and returned embeddings. - dtype (`torch.dtype`, *optional*): - Output dtype; defaults to the transformer dtype, or the text encoder dtype when used alone. - - Returns: - `tuple[torch.Tensor, torch.Tensor | None]`: - Positive and negative text embeddings, with padding tokens zeroed. - """ - prompt_embeds, negative_prompt_embeds = utils.encode_prompt( - self, - prompt, - negative_prompt, - do_classifier_free_guidance, - prompt_embeds, - negative_prompt_embeds, - max_sequence_length, - device, - dtype, - ) - return prompt_embeds, negative_prompt_embeds - - def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): - """Prepare full-video initial noise before encoding the image condition. - - Args: - batch_size (`int`): - Number of generated videos. - height (`int`): - Decoded frame height. - width (`int`): - Decoded frame width. - num_frames (`int`): - Decoded frame count of the form `4k + 1`. - dtype (`torch.dtype`): - Latent dtype. - device (`torch.device`): - Latent device. - generator (`torch.Generator` or `list[torch.Generator]`, *optional*): - Random generator, or one per video. - latents (`torch.Tensor`, *optional*): - Existing initial noise to validate and move to the requested device and dtype. - - Returns: - `torch.Tensor`: - Noise with shape `(B, C, T_latent, H_latent, W_latent)`. - """ - return utils.prepare_latents(self, batch_size, height, width, num_frames, dtype, device, generator, latents) + 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 + ) + with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + projected = self.proxy_projection(hidden, return_dict=False)[0] + proxy_features.append(projected.float().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): @@ -660,33 +346,18 @@ def __call__( latents: torch.Tensor | None = None, prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, - text_timeline: dict | list[dict] | None = None, - chunk_prompts: list[str] | None = None, - max_sequence_length: int | None = None, - output_type: 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, - *, - stream_state: EdenLiveStreamState | None = None, - proxy_frames=None, ) -> EdenLivePipelineOutput | tuple: - """Generate an offline video, or advance a stream using three already sampled RGB proxy frames. - - With `stream_state`, only `proxy_frames`, `prompt`, `prompt_embeds`, `latents`, and `return_dict` may be - supplied. All fixed settings come from the state. Calls sharing a pipeline must run serially. + """Generate a complete bidirectional video. Args: - stream_state (`EdenLiveStreamState`, *optional*): - State from `create_stream_state`. Returns one chunk: 9 RGB frames initially and 12 thereafter, or 3 - latent frames. Invalid inputs leave the state usable; model execution failures invalidate it. - proxy_frames (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`, *optional*): - Exactly three already sampled RGB frames, required with `stream_state`. No temporal subsampling. Tensor - layout is `(3, 3, H, W)`; NumPy layout is `(3, H, W, 3)`. Float values must be in [0, 1]. prompt (`str` or `list[str]`, *optional*): - Text describing the output video. Supply this or `prompt_embeds` unless using temporal prompts. In a - stream, new text takes effect in this chunk; None retains the current text. + 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. @@ -702,11 +373,10 @@ def __call__( 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 bidirectional, 30 for 14B bidirectional, or 4 per causal chunk. - Released causal schedules contain four fixed evaluation points and reject a different step count. + Denoising steps; defaults to 48 for 5B or 30 for 14B. guidance_scale (`float`, *optional*): - Classifier-free guidance strength; defaults to 3 for bidirectional or 1 for causal generation. Values - greater than one enable the negative prompt branch. + 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*): @@ -719,13 +389,6 @@ def __call__( Precomputed positive text embeddings in place of `prompt`. negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative text embeddings in place of `negative_prompt`. - text_timeline (`dict` or `list[dict]`, *optional*): - 5B bidirectional temporal captions, one dictionary per video, with `fps`, `num_frames`, and contiguous - `segments`. Each segment has `start_sec`, `end_sec`, and `text`. Exact boundaries belong to the segment - on their left. This option cannot be combined with precomputed embeddings. - chunk_prompts (`list[str]`, *optional*): - Causal-only prompts, one per latent chunk. Requires batch size one and no precomputed positive - embeddings. Each call starts with fresh attention caches. max_sequence_length (`int`, defaults to 512): Token count for each text prompt, including zero padding. output_type (`str`, defaults to `"np"`): @@ -735,11 +398,10 @@ def __call__( return_dict (`bool`, defaults to `True`): Return `EdenLivePipelineOutput`; otherwise return a one-element tuple containing the frames. attention_kwargs (`dict`, *optional*): - Keyword arguments passed to the transformer's attention processors. + 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. + 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`. @@ -750,56 +412,16 @@ def __call__( [`EdenLivePipelineOutput`] or `tuple`: Generated frames, or latent videos when `output_type="latent"`. """ - if stream_state is not None: - fixed_options = { - "image": image, - "video": video, - "height": height, - "width": width, - "num_frames": num_frames, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - "negative_prompt": negative_prompt, - "negative_prompt_embeds": negative_prompt_embeds, - "generator": generator, - "text_timeline": text_timeline, - "chunk_prompts": chunk_prompts, - "max_sequence_length": max_sequence_length, - "output_type": output_type, - "attention_kwargs": attention_kwargs, - "callback_on_step_end": callback_on_step_end, - "callback_on_step_end_tensor_inputs": callback_on_step_end_tensor_inputs, - } - supplied = [name for name, value in fixed_options.items() if value is not None] - if supplied: - raise ValueError(f"Streaming calls take fixed settings from the state; omit: {', '.join(supplied)}.") - return self._stream_step( - stream_state, - proxy_frames, - prompt=prompt, - prompt_embeds=prompt_embeds, - latents=latents, - return_dict=return_dict, - ) - if proxy_frames is not None: - raise ValueError("`proxy_frames` requires `stream_state`; use `video` for offline generation.") - max_sequence_length = 512 if max_sequence_length is None else max_sequence_length - output_type = "np" if output_type is None else output_type callback_on_step_end_tensor_inputs = ( ["latents"] if callback_on_step_end_tensor_inputs is None else callback_on_step_end_tensor_inputs ) - if text_timeline is not None and (self.is_causal or self.model_size != "5b"): - raise ValueError("`text_timeline` requires the 5B bidirectional profile.") - if chunk_prompts is not None and not self.is_causal: - raise ValueError("`chunk_prompts` requires a causal profile.") if guidance_scale is None: - guidance_scale = 1.0 if self.is_causal else 3.0 + guidance_scale = 3.0 if num_inference_steps is None: - num_inference_steps = 4 if self.is_causal else (48 if self.model_size == "5b" else 30) - height = (704 if self.model_size == "5b" else 480) if height is None else height - width = (1280 if self.model_size == "5b" else 832) if width is None else width - utils.check_inputs( - self, + 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, @@ -816,48 +438,15 @@ def __call__( proxy_video, num_frames = self.proxy_video_processor.preprocess_proxy_video( video, num_frames, self.vae_scale_factor_temporal ) - if text_timeline is not None: - if self.model_size != "5b" or prompt_embeds is not None or negative_prompt_embeds is not None: - raise ValueError( - "`text_timeline` requires the 5B profile and text prompts instead of precomputed embeddings." - ) - batch_size = 1 if isinstance(text_timeline, dict) else len(text_timeline) - else: - batch_size = ( - prompt_embeds.shape[0] - if prompt_embeds is not None - else ( - 1 - if isinstance(prompt, str) or (prompt is None and chunk_prompts is not None) - else len(prompt or []) - ) - ) + 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 `prompt`, `prompt_embeds`, `text_timeline`, or `chunk_prompts`.") - latents = self.prepare_latents(batch_size, height, width, num_frames, dtype, device, generator, latents) - chunk_embeds = None - num_chunks = (latents.shape[2] + self.config.chunk_size - 1) // self.config.chunk_size - if chunk_prompts is not None: - 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(chunk_prompt, str) for chunk_prompt in chunk_prompts) - ): - raise ValueError( - "`chunk_prompts` requires batch size one, no prompt embeddings, and one string per chunk." - ) - unique_embeds = {} - for chunk_prompt in chunk_prompts: - if chunk_prompt not in unique_embeds: - unique_embeds[chunk_prompt], _ = self.encode_prompt( - chunk_prompt, do_classifier_free_guidance=False, max_sequence_length=max_sequence_length - ) - chunk_embeds = torch.cat([unique_embeds[chunk_prompt] for chunk_prompt in chunk_prompts]) - prompt = chunk_prompts[0] - prompt_embeds = chunk_embeds[:1] - if text_timeline is None: + raise ValueError("Provide a positive prompt or prompt embeddings.") + try: + 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, @@ -866,317 +455,53 @@ def __call__( negative_prompt_embeds, max_sequence_length, ) - layout = None - else: - prompt_embeds, negative_prompt_embeds, layout = utils.encode_timeline( - self, - text_timeline, - negative_prompt, - num_frames, - latents.shape, - max_sequence_length, - self.do_classifier_free_guidance, + condition, first_frame, image_embeds, proxy_features = self.prepare_conditions( + image, proxy_video, height, width, num_frames, batch_size, device, dtype ) - condition, first_frame, image_embeds, proxy_features = utils.prepare_conditions( - self, image, proxy_video, height, width, num_frames, batch_size, device, dtype - ) - if self.is_causal: - return self._denoise_causal( - latents, - condition, - first_frame, - proxy_features, - image_embeds, - prompt_embeds, - negative_prompt_embeds, - chunk_embeds, - num_chunks, - batch_size, - device, - dtype, - num_inference_steps, - generator, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - output_type, - return_dict, - ) - return self._denoise_bidir( - latents, - condition, - first_frame, - proxy_features, - image_embeds, - prompt_embeds, - negative_prompt_embeds, - layout, - device, - num_inference_steps, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - output_type, - return_dict, - ) - - def _denoise_bidir( - self, - latents, - condition, - first_frame, - proxy_features, - image_embeds, - prompt_embeds, - negative_prompt_embeds, - layout, - device, - num_inference_steps, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - output_type, - return_dict, - ): - if self.model_size == "5b": - latents[:, :, :1] = first_frame - self.scheduler.set_timesteps(num_inference_steps, device=device) - if self.model_size == "5b" 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 = utils.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, - text_timeline=layout, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - if self.do_classifier_free_guidance: - negative_proxy = ( - proxy_features - if self.config.proxy_unconditional == "same" - else torch.zeros_like(proxy_features) - ) - negative_prediction = self.transformer( + 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, - negative_prompt_embeds, + prompt_embeds, encoder_hidden_states_image=image_embeds, - proxy_features=negative_proxy, - text_timeline=layout, + 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 self.model_size == "5b": - 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 = utils.decode_latents(self, latents, output_type) - self.maybe_free_model_hooks() - if not return_dict: - return (frames,) - return EdenLivePipelineOutput(frames=frames) - - def _denoise_chunk( - self, - latents, - condition, - first_frame, - proxy_features, - image_embeds, - prompt_embeds, - negative_prompt_embeds, - states, - num_inference_steps, - generator, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - callback_step_offset=0, - *, - condition_start=None, - ): - start = states[0].next_frame - if start == 0: - latents[:, :, :1] = first_frame - self.scheduler.set_timesteps(num_inference_steps, device=latents.device) - for step_index, timestep in enumerate(self.scheduler.timesteps): - model_input, model_timestep = utils.prepare_model_input( - self, latents, condition, first_frame, timestep, start, causal=True, condition_start=condition_start - ) - prediction = self.transformer( - model_input, - model_timestep, - prompt_embeds, - encoder_hidden_states_image=image_embeds, - proxy_features=proxy_features, - causal_state=states[0], - cache_update=False, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - if self.do_classifier_free_guidance: - negative_proxy = ( - proxy_features if self.config.proxy_unconditional == "same" else torch.zeros_like(proxy_features) - ) - negative_prediction = self.transformer( - model_input, - model_timestep, - negative_prompt_embeds, - encoder_hidden_states_image=image_embeds, - proxy_features=negative_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 = utils.sample_noise(self, latents.shape, generator, latents.device, latents.dtype) - latents = self.scheduler.step( - prediction, - timestep, - latents, - generator=None if variance_noise is not None else generator, - variance_noise=variance_noise, - return_dict=False, - )[0] - if start == 0: - 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, callback_step_offset + 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) - model_input, model_timestep = utils.prepare_model_input( - self, - latents, - condition, - first_frame, - latents.new_zeros(()), - start, - causal=True, - condition_start=condition_start, - ) - # A clean cache commit uses one zero timestep per sample, not per token. - model_timestep = torch.zeros(latents.shape[0], device=latents.device, dtype=torch.float32) - self.transformer( - model_input, - model_timestep, - prompt_embeds, - encoder_hidden_states_image=image_embeds, - proxy_features=proxy_features, - causal_state=states[0], - cache_update=True, - attention_kwargs=attention_kwargs, - return_dict=False, - ) - if self.do_classifier_free_guidance: - negative_proxy = ( - proxy_features if self.config.proxy_unconditional == "same" else torch.zeros_like(proxy_features) - ) - self.transformer( - model_input, - model_timestep, - negative_prompt_embeds, - encoder_hidden_states_image=image_embeds, - proxy_features=negative_proxy, - causal_state=states[1], - cache_update=True, - attention_kwargs=attention_kwargs, - return_dict=False, - ) - return latents, prompt_embeds, negative_prompt_embeds - - def _denoise_causal( - self, - latents, - condition, - first_frame, - proxy_features, - image_embeds, - prompt_embeds, - negative_prompt_embeds, - chunk_embeds, - num_chunks, - batch_size, - device, - dtype, - num_inference_steps, - generator, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - output_type, - return_dict, - ): - 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 = [] - try: - 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, prompt_embeds, negative_prompt_embeds = self._denoise_chunk( - latents[:, :, start : start + self.config.chunk_size].clone(), - condition, - first_frame, - proxy_features[:, :, start : start + self.config.chunk_size], - image_embeds, - prompt_embeds, - negative_prompt_embeds, - states, - num_inference_steps, - generator, - attention_kwargs, - callback_on_step_end, - callback_on_step_end_tensor_inputs, - chunk_index * num_inference_steps, - ) - generated.append(current) + 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() - latents = torch.cat(generated, dim=2) - frames = utils.decode_latents(self, latents, output_type) + frames = self.decode_latents(latents, output_type) finally: - for state in states: - state.reset() self.maybe_free_model_hooks() if not return_dict: return (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..07ee8b9ca288 --- /dev/null +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py @@ -0,0 +1,643 @@ +# 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, EdenLiveProxyProjection, EdenLiveTransformer3DModel +from ...schedulers import EdenLiveFlowMatchEulerDiscreteScheduler +from ...utils import is_ftfy_available, replace_example_docstring +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import EdenLivePipelineOutput +from .pipeline_utils import prepare_model_input, sample_noise +from .video_processor import EdenLiveVideoProcessor + + +if is_ftfy_available(): + import ftfy + + +# docstyle-ignore +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import EdenLiveCausalPipeline, EdenLiveTransformer3DModel, EdenLiveFlowMatchEulerDiscreteScheduler + >>> from diffusers.utils import export_to_video, load_image, load_video + >>> transformer = EdenLiveTransformer3DModel.from_pretrained( + ... "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer_causal", torch_dtype=torch.bfloat16 + ... ) + >>> scheduler = EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( + ... "VAST-AI/EdenLive-5B-Diffusers", subfolder="scheduler_causal" + ... ) + >>> pipe = EdenLiveCausalPipeline.from_pretrained( + ... "VAST-AI/EdenLive-5B-Diffusers", transformer=transformer, scheduler=scheduler, + ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": 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. + proxy_projection ([`EdenLiveProxyProjection`]): + Learned projection of DINO patch features. + 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->proxy_projection->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, + proxy_projection: EdenLiveProxyProjection, + 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, + proxy_projection=proxy_projection, + 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 sample_noise(self, shape, generator, device, 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 + ) + with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + projected = self.proxy_projection(hidden, return_dict=False)[0] + proxy_features.append(projected.float().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. Under context parallelism, the + first rank generates and broadcasts each complete noise tensor. + 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.") + states = [] + try: + 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 = sample_noise( + self, current.shape, generator, current.device, 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) + finally: + for state in states: + state.reset() + self.maybe_free_model_hooks() + if not return_dict: + return (frames,) + return EdenLivePipelineOutput(frames=frames) diff --git a/src/diffusers/pipelines/edenlive/pipeline_utils.py b/src/diffusers/pipelines/edenlive/pipeline_utils.py index c2948f7b4a36..3e70bd9547c0 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_utils.py +++ b/src/diffusers/pipelines/edenlive/pipeline_utils.py @@ -12,194 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import bisect -import html -import math -import re - import torch -import torch.nn.functional as F -from ...models.autoencoders.autoencoder_edenlive import EdenLiveLightTAEDecoder -from ...utils import is_ftfy_available from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from .video_processor import EdenLiveVideoProcessor - - -if is_ftfy_available(): - import ftfy - - -def initialize_pipeline( - pipe, - vae, - text_encoder, - tokenizer, - transformer, - scheduler, - proxy_encoder, - proxy_projection, - image_encoder, - image_processor, - decoder, - profile, - proxy_resize_mode, - proxy_unconditional, - proxy_projection_autocast, - proxy_image_size, -): - if profile not in ("5b-bidir", "14b-bidir", "5b-causal", "14b-causal"): - raise ValueError("`profile` must be one of 5b-bidir, 14b-bidir, 5b-causal, or 14b-causal.") - pipe.model_size, mode = profile.split("-") - pipe.is_causal = mode == "causal" - if proxy_unconditional not in ("same", "zero"): - raise ValueError("`proxy_unconditional` must be 'same' or 'zero'.") - if decoder is not None: - validate_decoder(decoder, vae) - pipe.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - proxy_encoder=proxy_encoder, - proxy_projection=proxy_projection, - image_encoder=image_encoder, - image_processor=image_processor, - decoder=decoder, - ) - pipe.register_to_config( - profile=profile, - proxy_resize_mode=proxy_resize_mode, - proxy_unconditional=proxy_unconditional, - proxy_projection_autocast=proxy_projection_autocast, - proxy_image_size=proxy_image_size, - ) - pipe.vae_scale_factor_temporal = vae.config.scale_factor_temporal if vae is not None else 4 - pipe.vae_scale_factor_spatial = vae.config.scale_factor_spatial if vae is not None else 8 - pipe.video_processor = VideoProcessor(vae_scale_factor=pipe.vae_scale_factor_spatial) - pipe.proxy_video_processor = EdenLiveVideoProcessor(vae_scale_factor=1, do_normalize=False, do_convert_rgb=True) - - -def validate_decoder(decoder, vae): - if not isinstance(decoder, torch.nn.Module) or not callable(getattr(decoder, "decode", None)): - raise ValueError("`decoder` must be a module with a `decode(latents, return_dict=False)` method.") - config = getattr(decoder, "config", None) - latent_channels = getattr(config, "latent_channels", getattr(config, "z_dim", None)) - temporal = getattr(decoder, "scale_factor_temporal", getattr(config, "scale_factor_temporal", None)) - spatial = getattr(decoder, "scale_factor_spatial", getattr(config, "scale_factor_spatial", None)) - if vae is None or (latent_channels, temporal, spatial) != ( - vae.config.z_dim, - vae.config.scale_factor_temporal, - vae.config.scale_factor_spatial, - ): - raise ValueError("`decoder` latent channels and temporal/spatial scale factors must match the Wan VAE.") - - -def encode_prompt( - pipe, - prompt, - negative_prompt=None, - do_classifier_free_guidance=True, - prompt_embeds=None, - negative_prompt_embeds=None, - max_sequence_length=512, - device=None, - dtype=None, -): - device = device or pipe._execution_device - dtype = dtype or (pipe.transformer.dtype if pipe.transformer is not None else pipe.text_encoder.dtype) - - def encode(texts): - 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 = pipe.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 = pipe.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) - - 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 = encode(prompt) - 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 = encode(negative_prompt) - 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( - pipe, - prompt, - prompt_embeds, - negative_prompt, - negative_prompt_embeds, - image, - video, - height, - width, - output_type, - callback_inputs, -): - 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 = pipe.transformer.config.patch_size - if ( - height <= 0 - or width <= 0 - or height % (pipe.vae_scale_factor_spatial * patch[1]) - or width % (pipe.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(pipe._callback_tensor_inputs): - raise ValueError(f"Callback tensor inputs must belong to {pipe._callback_tensor_inputs}.") - - -def prepare_latents(pipe, batch_size, height, width, num_frames, dtype, device, generator=None, latents=None): - shape = ( - batch_size, - pipe.vae.config.z_dim, - (num_frames - 1) // pipe.vae_scale_factor_temporal + 1, - height // pipe.vae_scale_factor_spatial, - width // pipe.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 sample_noise(pipe, shape, generator, device, dtype) def sample_noise(pipe, shape, generator, device, dtype): @@ -217,100 +32,14 @@ def sample_noise(pipe, shape, generator, device, dtype): return noise -def prepare_reference_pixels(pipe, image, height, width, batch_size, device): - pixels = pipe.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.") - return pixels - - -def encode_image_embeds(pipe, pixels, device, dtype): - clip_embeds = None - if pipe.model_size == "14b": - if pipe.image_encoder is None or pipe.image_processor is None: - raise ValueError("The 14B profile 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 = pipe.image_processor(images=list(rgb), return_tensors="pt").pixel_values - clip_embeds = ( - pipe.image_encoder( - clip_inputs.to(device=device, dtype=pipe.image_encoder.dtype), output_hidden_states=True - ) - .hidden_states[-2] - .to(dtype) - ) - return clip_embeds - - -def encode_proxy_features(pipe, proxy_video, batch_size, device): - proxy_video = proxy_video.expand(batch_size, -1, -1, -1, -1) - proxy_height, proxy_width = pipe.config.proxy_image_size - proxy_features = [] - for frame in proxy_video.unbind(2): - frame = frame.to(device=device, dtype=torch.float32) - if pipe.model_size == "5b": - # Match the original FP32 [-1, 1] -> [0, 1] round trip before DINO. - frame = (frame * 2 - 1 + 1) * 0.5 - inputs = pipe.proxy_video_processor.preprocess_proxy_frames( - frame, proxy_height, proxy_width, pipe.config.proxy_resize_mode - ) - with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): - hidden = pipe.proxy_encoder(pixel_values=inputs).last_hidden_state - patch_size = pipe.proxy_encoder.config.patch_size - grid_height, grid_width = proxy_height // patch_size, proxy_width // patch_size - start = 1 + pipe.proxy_encoder.config.num_register_tokens - hidden = hidden[:, start : start + grid_height * grid_width].reshape(batch_size, grid_height, grid_width, -1) - with torch.autocast( - device.type, dtype=torch.bfloat16, enabled=device.type == "cuda" and pipe.config.proxy_projection_autocast - ): - projected = pipe.proxy_projection(hidden if pipe.config.proxy_projection_autocast else hidden.float()) - proxy_features.append(projected.float().permute(0, 3, 1, 2)) - return torch.stack(proxy_features, dim=2) - - -def prepare_conditions(pipe, image, proxy_video, height, width, num_frames, batch_size, device, dtype): - pixels = prepare_reference_pixels(pipe, image, height, width, batch_size, device) - 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 = encode_image_embeds(pipe, pixels, device, dtype) - if pipe.model_size == "14b": - 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 = pipe.vae.encode(condition_video.to(pipe.vae.dtype)).latent_dist.mode() - mean = condition.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) - inverse_std = 1 / condition.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) - condition = ((condition - mean) * inverse_std).to(dtype) - first_frame = condition[:, :, :1] - if pipe.model_size == "14b": - mask = condition.new_zeros(batch_size, pipe.vae_scale_factor_temporal, *condition.shape[2:]) - mask[:, :, 0] = 1 - condition = torch.cat([mask, condition], dim=1) - proxy_features = encode_proxy_features(pipe, proxy_video, batch_size, device) - return condition, first_frame, clip_embeds, proxy_features - - -def prepare_model_input( - pipe, latents, condition, first_frame, timestep, start=0, causal=False, *, condition_start=None -): +def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0, causal=False): model_input = latents - pin_first = start == 0 and (causal or pipe.model_size == "5b") + 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.model_size == "14b": - offset = start if condition_start is None else condition_start - model_input = torch.cat([model_input, condition[:, :, offset : offset + latents.shape[2]]], dim=1) + 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 @@ -318,82 +47,3 @@ def prepare_model_input( timestep = timestep[:, None].expand(-1, latents.shape[2] * tokens_per_frame).clone() timestep[:, :tokens_per_frame] = 0 return model_input, timestep - - -def decode_latents(pipe, latents, output_type, *, decoder_cache=None): - if output_type == "latent": - return latents - latents = latents.to(pipe.vae.dtype) - mean = latents.new_tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1, 1) - std = latents.new_tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1, 1) - if isinstance(pipe.decoder, EdenLiveLightTAEDecoder): - raw_latents = latents * std + mean - else: - raw_latents = latents / (1 / std) + mean - decoder = pipe.decoder if pipe.decoder is not None else pipe.vae - dtype = getattr(decoder, "dtype", None) - if dtype is None: - parameter = next(decoder.parameters(), None) - dtype = parameter.dtype if parameter is not None else raw_latents.dtype - decode_kwargs = {} if decoder_cache is None else {"cache": decoder_cache} - video = decoder.decode(raw_latents.to(dtype=dtype), return_dict=False, **decode_kwargs)[0] - return pipe.video_processor.postprocess_video(video, output_type=output_type) - - -def encode_timeline(pipe, timelines, negative_prompt, num_frames, latent_shape, max_sequence_length, guidance): - timelines = [timelines] if isinstance(timelines, dict) else timelines - batch_size = len(timelines) - negatives = ( - [negative_prompt or ""] * batch_size if isinstance(negative_prompt, (str, type(None))) else negative_prompt - ) - if len(negatives) != batch_size: - raise ValueError("Negative prompt batch must match the timeline batch.") - texts, negative_texts, lengths, counts = [], [], [], [] - tokens_per_frame = (latent_shape[-2] // pipe.transformer.config.patch_size[1]) * ( - latent_shape[-1] // pipe.transformer.config.patch_size[2] - ) - for sample, negative in zip(timelines, negatives): - fps = float(sample["fps"]) - if not math.isfinite(fps) or fps <= 0 or sample["num_frames"] != num_frames: - raise ValueError("Each timeline needs positive finite fps and the requested decoded num_frames.") - duration, previous = (num_frames - 1) / fps, 0.0 - segments = [] - for segment in sample["segments"]: - start, end = float(segment["start_sec"]), float(segment["end_sec"]) - if ( - not math.isfinite(start) - or not math.isfinite(end) - or end <= start - or not math.isclose(start, previous, abs_tol=1e-6) - ): - raise ValueError("Timeline segments must be finite, ordered, contiguous, and start at zero.") - previous = end - if start < duration: - segments.append({"text": segment["text"], "end_sec": min(end, duration)}) - if not segments or previous < duration: - raise ValueError("Timeline segments must cover the requested video.") - frame_counts = [0] * len(segments) - for frame in range(latent_shape[2]): - index = bisect.bisect_left( - [segment["end_sec"] for segment in segments], frame * pipe.vae_scale_factor_temporal / fps - 1e-7 - ) - frame_counts[index] += 1 - if 0 in frame_counts: - raise ValueError("A timeline segment is shorter than the latent time resolution.") - lengths.append(tuple(count * tokens_per_frame for count in frame_counts)) - counts.append(len(segments)) - texts.extend(segment["text"] for segment in segments) - negative_texts.extend([negative] * len(segments)) - positive, negative = encode_prompt(pipe, texts, negative_texts, guidance, max_sequence_length=max_sequence_length) - - def pack(embeddings): - if embeddings is None: - return None - result = embeddings.new_zeros(batch_size, max(counts) * max_sequence_length, embeddings.shape[-1]) - offset = 0 - for index, count in enumerate(counts): - result[index, : count * max_sequence_length] = embeddings[offset : offset + count].flatten(0, 1) - offset += count - return result - - return pack(positive), pack(negative), {"lengths": tuple(lengths), "text_length": max_sequence_length} diff --git a/src/diffusers/pipelines/edenlive/video_processor.py b/src/diffusers/pipelines/edenlive/video_processor.py index aa37a23dd6c8..8822346e46a0 100644 --- a/src/diffusers/pipelines/edenlive/video_processor.py +++ b/src/diffusers/pipelines/edenlive/video_processor.py @@ -27,37 +27,6 @@ class EdenLiveVideoProcessor(VideoProcessor): 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_chunk(self, proxy_frames): - """Convert three already sampled RGB condition frames to BCTHW without temporal subsampling.""" - if isinstance(proxy_frames, (np.ndarray, torch.Tensor)): - if proxy_frames.ndim == 5 and proxy_frames.shape[0] == 1: - proxy_frames = proxy_frames[0] - if proxy_frames.ndim != 4: - raise ValueError("Pass a single proxy chunk with three RGB frames.") - elif not isinstance(proxy_frames, list): - raise ValueError("Pass three RGB proxy frames as a list, NumPy array, or tensor.") - if len(proxy_frames) != 3: - raise ValueError("Each streaming chunk requires exactly three already sampled proxy frames.") - frame_type = type(proxy_frames[0]) - if not all(isinstance(frame, frame_type) for frame in proxy_frames): - raise ValueError("All proxy frames must use the same image format.") - for frame in proxy_frames: - if isinstance(frame, PIL.Image.Image): - if frame.mode != "RGB": - raise ValueError("Proxy frames must be RGB images.") - continue - if not isinstance(frame, (np.ndarray, torch.Tensor)) or frame.ndim != 3: - raise ValueError("Each proxy frame must be an RGB image.") - channel_axis = 0 if isinstance(frame, torch.Tensor) else -1 - if frame.shape[channel_axis] != 3 or any(size == 0 for size in frame.shape): - raise ValueError("Proxy frames must have three RGB channels.") - values = torch.as_tensor(frame) - if not torch.is_floating_point(values) or not torch.isfinite(values).all(): - raise ValueError("NumPy and tensor proxy frames must contain finite floating-point RGB in [0, 1].") - if values.min() < 0 or values.max() > 1: - raise ValueError("Proxy frames must contain RGB values in [0, 1].") - return self.preprocess_video([proxy_frames]) - 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) @@ -86,21 +55,12 @@ def preprocess_proxy_frames(self, images, height, width, resize_mode="bilinear") pixels = F.interpolate( images.float(), (height, width), mode="bilinear", align_corners=False, antialias=True ) - elif resize_mode in ("bicubic", "bicubic_center_crop"): + 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) - if resize_mode == "bicubic": - image = image.resize((width, height), PIL.Image.Resampling.BICUBIC) - else: - ratio = image.width / image.height - resized_height, resized_width = ( - (height, round(height * ratio)) if ratio > width / height else (round(width / ratio), width) - ) - image = image.resize((resized_width, resized_height), PIL.Image.Resampling.BICUBIC) - left, top = (resized_width - width) // 2, (resized_height - height) // 2 - image = image.crop((left, top, left + width, top + height)) + 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: diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index a2b5920d2be5..f37686781430 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1339,21 +1339,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class EdenLiveLightTAEDecodeCache(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"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 9ffd7c8c5db7..54d6ac8b76ba 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2057,7 +2057,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class EdenLivePipeline(metaclass=DummyObject): +class EdenLiveCausalPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -2072,7 +2072,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class EdenLivePipelineOutput(metaclass=DummyObject): +class EdenLivePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -2087,7 +2087,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class EdenLiveStreamState(metaclass=DummyObject): +class EdenLivePipelineOutput(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): 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_edenlive_proxy_projection.py b/tests/models/transformers/test_models_edenlive_proxy_projection.py new file mode 100644 index 000000000000..9cf8d6bbf743 --- /dev/null +++ b/tests/models/transformers/test_models_edenlive_proxy_projection.py @@ -0,0 +1,91 @@ +# 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 EdenLiveProxyProjection +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 EdenLiveProxyProjectionTesterConfig(BaseModelTesterConfig): + model_class = EdenLiveProxyProjection + main_input_name = "hidden_states" + input_shape = (4, 4, 16) + output_shape = (4, 4, 8) + + def get_init_dict(self): + return {"in_channels": 16, "out_channels": 8} + + def get_dummy_inputs(self, height=4, width=4): + return { + "hidden_states": randn_tensor( + (1, height, width, 16), + generator=torch.Generator().manual_seed(0), + device=torch_device, + dtype=self.torch_dtype, + ) + } + + +class TestEdenLiveProxyProjectionModel(EdenLiveProxyProjectionTesterConfig, ModelTesterMixin): + pass + + +class TestEdenLiveProxyProjectionMemory(EdenLiveProxyProjectionTesterConfig, 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] + + @torch.no_grad() + def test_cpu_offload(self, base_model_output): + from accelerate import cpu_offload + + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval() + cpu_offload(model, execution_device=torch.device(torch_device)) + actual = model(**self.get_dummy_inputs()).sample + torch.testing.assert_close(actual, base_model_output, atol=1e-5, rtol=0) + + @pytest.mark.skip(reason="One linear layer cannot be partitioned between an execution device and disk.") + def test_disk_offload_without_safetensors(self): + pass + + @pytest.mark.skip(reason="One linear layer cannot be partitioned between an execution device and disk.") + def test_disk_offload_with_safetensors(self): + pass + + +class TestEdenLiveProxyProjectionTorchCompile(EdenLiveProxyProjectionTesterConfig, 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..26d0eb2c6537 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_edenlive.py @@ -0,0 +1,140 @@ +# 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 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, + ContextParallelTesterMixin, + 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, + "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, 8, 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): + @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() + 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 * state.tokens_per_frame for layer in state.layers) + if index == 1: + sink_keys = [layer.key[:, : 6 * state.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 * state.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 TestEdenLiveTransformerContextParallel(EdenLiveTransformerTesterConfig, ContextParallelTesterMixin): + pass + + +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..77b8dbe09f68 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_edenlive_causal.py @@ -0,0 +1,224 @@ +# 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 os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from diffusers import ContextParallelConfig, EdenLiveTransformer3DModel +from diffusers.models.transformers import transformer_edenlive + +from ...testing_utils import enable_full_determinism +from .test_models_transformer_edenlive import EdenLiveTransformerTesterConfig + + +class TestEdenLiveCausalCacheOrder: + @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) + + +def _parallel_outputs(model, inputs): + with torch.no_grad(): + output = model(**inputs, return_dict=False)[0] + state = model.create_causal_state( + batch_size=1, + height=8, + width=8, + chunk_size=3, + max_context_frames=18, + sink_frames=6, + device=inputs["hidden_states"].device, + dtype=inputs["hidden_states"].dtype, + ) + clean_inputs = {**inputs, "timestep": torch.zeros_like(inputs["timestep"])} + sink_keys = None + for index in range(8): + model(**clean_inputs, causal_state=state, cache_update=True) + assert state.next_frame == 3 * (index + 1) + assert all(layer.num_tokens <= 18 * state.tokens_per_frame // state.cp_size for layer in state.layers) + sink_tokens = 6 * state.tokens_per_frame // state.cp_size + if index == 1: + sink_keys = [layer.key[:, :sink_tokens].clone() for layer in state.layers] + elif index > 1: + for expected, layer in zip(sink_keys, state.layers): + torch.testing.assert_close(layer.key[:, :sink_tokens], expected, rtol=0, atol=0) + tail_inputs = { + **inputs, + "hidden_states": inputs["hidden_states"][:, :, :1], + "proxy_features": inputs["proxy_features"][:, :, :1], + "timestep": inputs["timestep"][:, :16] if inputs["timestep"].ndim == 2 else inputs["timestep"], + } + tail = model(**tail_inputs, causal_state=state, return_dict=False)[0] + assert state.next_frame == 24 + tail_clean = {**tail_inputs, "timestep": torch.zeros_like(tail_inputs["timestep"])} + model(**tail_clean, causal_state=state, cache_update=True) + assert state.next_frame == 25 + return [output.float().cpu(), tail.float().cpu()] + + +def _context_parallel_worker(rank, world_size, port, init_dict, state_dict, inputs, backend, use_fsdp, output_path): + # Spawned workers must establish the same settings as the single-device reference. + enable_full_determinism() + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), RANK=str(rank), WORLD_SIZE=str(world_size)) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 + try: + dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=180)) + model = EdenLiveTransformer3DModel(**init_dict).eval() + model.load_state_dict(state_dict) + model.to(device=device, dtype=dtype) + model.set_attention_backend(backend) + inputs = { + name: tensor.to(device=device, dtype=torch.float32 if name == "timestep" else dtype) + for name, tensor in inputs.items() + } + + def check_timestep(_module, _args, kwargs): + assert kwargs["timestep"].dtype == torch.float32 + assert torch.isin(kwargs["timestep"], torch.tensor([0.0, 375.0], device=device)).all() + + model.register_forward_pre_hook(check_timestep, with_kwargs=True) + from torch.distributed.device_mesh import init_device_mesh + + mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) + model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) + if use_fsdp: + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + + policy = MixedPrecisionPolicy( + param_dtype=dtype, reduce_dtype=torch.float32, output_dtype=dtype, cast_forward_inputs=False + ) + for block in model.blocks: + fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + fully_shard(model, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) + outputs = _parallel_outputs(model, inputs) + if rank == 0: + torch.save(outputs, output_path) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +class TestEdenLiveContextParallel(EdenLiveTransformerTesterConfig): + @pytest.mark.context_parallel + @pytest.mark.attention + @pytest.mark.parametrize("world_size", [2, 4]) + @pytest.mark.parametrize( + "backend,use_fsdp", [("native", False), ("_native_flash", False), ("_native_flash", True)] + ) + @pytest.mark.parametrize("num_layers,token_time", [(2, False), (1, True)]) + def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsdp, num_layers, token_time): + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + if not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < world_size: + pytest.skip(f"This case requires {world_size} CUDA devices with NCCL.") + previous_backend, _ = _AttentionBackendRegistry.get_active_backend() + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + try: + enable_full_determinism() + config = {**self.get_init_dict(), "attention_head_dim": 16, "num_layers": num_layers} + model = EdenLiveTransformer3DModel(**config).eval() + state = {name: tensor.cpu().clone() for name, tensor in model.state_dict().items()} + inputs = {name: tensor.cpu() for name, tensor in self.get_dummy_inputs().items()} + # This value would be rounded to 376 by an indiscriminate BF16 input cast. + inputs["timestep"] = torch.tensor([375.0], dtype=torch.float32) + if token_time: + inputs["timestep"] = torch.full((1, 48), 375.0) + inputs["timestep"][:, :16] = 0 + dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 + model.to(device="cuda:0", dtype=dtype).set_attention_backend(backend) + reference_inputs = { + name: tensor.to(device="cuda:0", dtype=torch.float32 if name == "timestep" else dtype) + for name, tensor in inputs.items() + } + expected = _parallel_outputs(model, reference_inputs) + del model, reference_inputs + torch.cuda.empty_cache() + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + output_path = str(tmp_path / "context_parallel.pt") + mp.spawn( + _context_parallel_worker, + args=(world_size, port, config, state, inputs, backend, use_fsdp, output_path), + nprocs=world_size, + join=True, + ) + actual = torch.load(output_path, map_location="cpu", weights_only=True) + for reference, output in zip(expected, actual): + assert torch.isfinite(output).all() + # Keep the existing small-model CP backend tolerance. + torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2) + finally: + _AttentionBackendRegistry.set_active_backend(previous_backend) + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) diff --git a/tests/pipelines/edenlive/test_edenlive.py b/tests/pipelines/edenlive/test_edenlive.py deleted file mode 100644 index 847f9bee2998..000000000000 --- a/tests/pipelines/edenlive/test_edenlive.py +++ /dev/null @@ -1,1807 +0,0 @@ -# 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 hashlib -import importlib.util -import json -import os -import socket -import tempfile -import unittest -from datetime import timedelta -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from PIL import Image -from safetensors.torch import load_file, save_file -from tokenizers import Tokenizer, models, pre_tokenizers -from transformers import ( - CLIPImageProcessor, - CLIPVisionConfig, - CLIPVisionModel, - DINOv3ViTConfig, - DINOv3ViTModel, - PreTrainedTokenizerFast, - UMT5Config, - UMT5EncoderModel, -) - -from diffusers import ( - AutoencoderKLWan, - ContextParallelConfig, - EdenLiveFlowMatchEulerDiscreteScheduler, - EdenLiveLightTAEDecoder, - EdenLivePipeline, - EdenLiveProxyProjection, - EdenLiveTransformer3DModel, - UniPCMultistepScheduler, -) -from diffusers.models.transformers import transformer_edenlive -from diffusers.utils.torch_utils import randn_tensor - -from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BasePipelineTesterConfig - - -@pytest.fixture(autouse=True, scope="class") -def full_precision_matmul(): - precision = torch.get_float32_matmul_precision() - matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 - cudnn_tf32 = torch.backends.cudnn.allow_tf32 - torch.set_float32_matmul_precision("highest") - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - try: - yield - finally: - torch.set_float32_matmul_precision(precision) - torch.backends.cuda.matmul.allow_tf32 = matmul_tf32 - torch.backends.cudnn.allow_tf32 = cudnn_tf32 - - -def make_tokenizer(): - words = "a character walks through snow blurry racing car on sunny road fire ocean one too many".split() - vocabulary = {word: index for index, word in enumerate(["", "", "", *words])} - backend = Tokenizer(models.WordLevel(vocabulary, unk_token="")) - backend.pre_tokenizer = pre_tokenizers.Whitespace() - return PreTrainedTokenizerFast( - tokenizer_object=backend, pad_token="", eos_token="", unk_token="", model_max_length=512 - ) - - -def file_digests(directory): - return { - str(path.relative_to(directory)): hashlib.sha256(path.read_bytes()).hexdigest() - for path in directory.rglob("*") - if path.is_file() - } - - -@pytest.fixture(params=["5b-bidir", "14b-bidir", "5b-causal", "14b-causal"]) -def original_checkpoint(tmp_path, request, monkeypatch): - import httpx - import requests - - from diffusers.pipelines.edenlive import loading_utils - - def unexpected_request(*args, **kwargs): - pytest.fail("A local original checkpoint must not request any external Hub configuration.") - - monkeypatch.setattr(httpx.Client, "send", unexpected_request) - monkeypatch.setattr(requests.Session, "send", unexpected_request) - - def tiny_decoder(**_kwargs): - # Production LightTAE dimensions stay fixed by the release profile. Only - # its construction is reduced for this weight-loading regression fixture. - return EdenLiveLightTAEDecoder(latent_channels=4, patch_size=1, block_out_channels=(4, 4, 4, 4)) - - monkeypatch.setattr(loading_utils, "EdenLiveLightTAEDecoder", tiny_decoder) - profile = request.param - size, mode = profile.split("-") - fixture = EdenLivePipelineTesterConfig() - fixture.profile = size - components = fixture.get_dummy_components() - source = tmp_path / "original" - source.mkdir() - (source / "model_index.json").write_text( - json.dumps({"_class_name": "EdenLivePipeline", "_edenlive_original": {"_path_overrides": {}}}) - ) - transformer_dir = source / (f"dit-{size}-bidir" if mode == "bidir" else f"dit-{size}") - transformer_dir.mkdir() - components["transformer"].save_config(transformer_dir) - save_file(components["transformer"].state_dict(), transformer_dir / "diffusion_pytorch_model.safetensors") - vae_dir = source / f"vae-{size}" - vae_dir.mkdir() - components["vae"].save_config(vae_dir) - # Include original Wan parameter names and the original .pth container. - vae_state = {} - replacements = { - "encoder.conv_in.": "encoder.conv1.", - "decoder.conv_out.": "decoder.head.2.", - "quant_conv.": "conv1.", - "post_quant_conv.": "conv2.", - } - for key, value in components["vae"].state_dict().items(): - for prefix, original_prefix in replacements.items(): - if key.startswith(prefix): - key = original_prefix + key[len(prefix) :] - break - vae_state[key] = value - torch.save(vae_state, vae_dir / ("Wan2.2_VAE.pth" if size == "5b" else "Wan2.1_VAE.pth")) - components["text_encoder"].save_pretrained(source / "umt5/text_encoder") - components["tokenizer"].save_pretrained(source / "umt5/tokenizer") - # This fast fixture consumes tokenizer.json. Mirror the additional vocabulary - # and special-token files listed by the original release's download plan. - (source / "umt5/tokenizer/spiece.model").write_bytes(b"unused-fast-tokenizer-fixture") - (source / "umt5/tokenizer/special_tokens_map.json").write_text( - json.dumps({"pad_token": "", "eos_token": "", "unk_token": ""}) - ) - proxy_dir = source / "vision-encoder" - components["proxy_encoder"].config.save_pretrained(proxy_dir / "backbone") - proxy_state = { - f"module.encoder.backbone.{key}": value for key, value in components["proxy_encoder"].state_dict().items() - } - proxy_state.update( - {f"module.encoder.{key}": value for key, value in components["proxy_projection"].state_dict().items()} - ) - proxy_state.update( - { - "module.encoder.global_projection.weight": torch.zeros(2, 32), - "module.encoder.global_projection.bias": torch.zeros(2), - } - ) - save_file(proxy_state, proxy_dir / "projection.safetensors") - if size == "14b": - components["image_encoder"].save_pretrained(source / "clip/image_encoder") - components["image_processor"].save_pretrained(source / "clip/image_processor") - decoder = EdenLiveLightTAEDecoder(latent_channels=4, patch_size=1, block_out_channels=(4, 4, 4, 4)) - decoder_dir = source / f"decoder-{size}" - decoder_dir.mkdir() - decoder_state = { - key.replace("decoder.", "decoder.blocks.", 1): value for key, value in decoder.state_dict().items() - } - temporal_key = "decoder.blocks.7.conv.weight" - decoder_state[temporal_key] = torch.cat( - [torch.zeros_like(decoder_state[temporal_key]), decoder_state[temporal_key]] - ) - decoder_state["encoder.unused.weight"] = torch.ones(1) - torch.save(decoder_state, decoder_dir / ("lighttaew2_2.pth" if size == "5b" else "lighttaew2_1.pth")) - return SimpleNamespace( - source=source, - profile=profile, - pipeline_class=EdenLivePipeline, - components=components, - decoder=decoder, - hashes=file_digests(source), - transformer_dir=transformer_dir, - ) - - -@pytest.fixture -def original_hub(original_checkpoint, monkeypatch): - from diffusers.pipelines.edenlive import loading_utils as pipeline_utils - - fixture = original_checkpoint - calls = SimpleNamespace(metadata=[], snapshots=[]) - revision = "0123456789abcdef0123456789abcdef01234567" - repo_id = "test-org/edenlive-original-fixture" - - def metadata_download(repository, filename, **kwargs): - assert repository == repo_id and filename == "model_index.json" - assert kwargs["revision"] == revision - calls.metadata.append(kwargs) - return str(fixture.source / filename) - - def snapshot_download(repository, **kwargs): - assert repository == repo_id - calls.snapshots.append(kwargs) - return str(fixture.source) - - monkeypatch.setattr(pipeline_utils, "_resolve_revision", lambda *args, **kwargs: revision) - monkeypatch.setattr(pipeline_utils, "hf_hub_download", metadata_download) - monkeypatch.setattr(pipeline_utils, "snapshot_download", snapshot_download) - return SimpleNamespace(fixture=fixture, calls=calls, revision=revision, repo_id=repo_id, module=pipeline_utils) - - -class EdenLivePipelineTesterConfig(BasePipelineTesterConfig): - pipeline_class = EdenLivePipeline - profile = "5b" - causal = False - required_input_params_in_call_signature = frozenset( - ["prompt", "negative_prompt", "image", "video", "height", "width", "prompt_embeds", "negative_prompt_embeds"] - ) - optional_input_params = frozenset( - ["num_frames", "num_inference_steps", "guidance_scale", "generator", "latents", "output_type", "return_dict"] - ) - batch_input_params = frozenset(["prompt", "negative_prompt"]) - 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) - text_encoder = UMT5EncoderModel( - UMT5Config(vocab_size=1103, d_model=32, d_ff=32, d_kv=8, num_layers=1, num_heads=4, dropout_rate=0.0) - ).eval() - tokenizer = make_tokenizer() - 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) - proxy_projection = EdenLiveProxyProjection(in_channels=32, out_channels=8) - 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_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, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "transformer": transformer, - "scheduler": scheduler, - "proxy_encoder": proxy_encoder, - "proxy_projection": proxy_projection, - "image_encoder": image_encoder, - "image_processor": image_processor, - } - - def get_pipeline(self, **components): - causal = self.causal - pipe = self.pipeline_class( - **(components or self.get_dummy_components()), - profile=f"{self.profile}-{'causal' if causal else 'bidir'}", - proxy_image_size=(32, 32), - proxy_resize_mode="bicubic" if self.profile == "14b" else "bilinear", - proxy_unconditional="same", - ) - for component in pipe.components.values(): - if hasattr(component, "eval"): - component.eval() - pipe.set_progress_bar_config(disable=True) - return pipe - - def test_inference_and_standard_save_load(self, tmp_path): - pipe = self.get_pipeline().to(torch_device) - expected = pipe(**self.get_dummy_inputs()).frames - assert expected.shape == (1, *self.output_shape) - assert torch.isfinite(expected).all() - pipe.save_pretrained(tmp_path) - for loader in (self.pipeline_class,): - restored = loader.from_pretrained(tmp_path, local_files_only=True).to(torch_device) - restored.set_progress_bar_config(disable=True) - actual = restored(**self.get_dummy_inputs()).frames - assert type(restored) is self.pipeline_class - assert restored.config.profile == pipe.config.profile - torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) - - 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 TestEdenLivePipeline(EdenLivePipelineTesterConfig): - @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 - - def test_timeline_boundary(self): - from diffusers.pipelines.edenlive.pipeline_utils import encode_timeline - - pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs() - inputs["text_timeline"] = { - "num_frames": 9, - "fps": 4, - "segments": [ - {"start_sec": 0, "end_sec": 1, "text": "snow"}, - {"start_sec": 1, "end_sec": 2, "text": "fire"}, - ], - } - output = pipe(**inputs).frames - assert output.shape == (1, *self.output_shape) - _, _, layout = encode_timeline( - pipe, inputs["text_timeline"], inputs["negative_prompt"], 9, (1, 4, 3, 2, 2), 16, True - ) - assert layout["lengths"] == ((2, 1),) - - -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="14B profile requires"): - restored(**self.get_dummy_inputs()) - - -class TestEdenLive14BPipeline(EdenLivePipelineTesterConfig, EdenLive14BRequiredComponentsTesterMixin): - profile = "14b" - - def test_timeline_is_rejected(self): - pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs() - inputs["text_timeline"] = {} - with pytest.raises(ValueError, match="5B"): - pipe(**inputs) - - -class EdenLiveChunkCausalPipelineTesterConfig(EdenLivePipelineTesterConfig): - causal = True - - def get_dummy_components(self): - components = super().get_dummy_components() - components["scheduler"] = EdenLiveFlowMatchEulerDiscreteScheduler( - stochastic_sampling=True, default_timesteps=[1000, 500] - ) - return components - - -class TestEdenLiveChunkCausalPipeline(EdenLiveChunkCausalPipelineTesterConfig): - 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()) - - def test_chunk_prompt_changes_tail(self): - 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") - 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"] - with pytest.raises(ValueError, match="chunk_prompts"): - pipe(**inputs) - - 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) - - -class TestEdenLive14BChunkCausalPipeline( - EdenLiveChunkCausalPipelineTesterConfig, EdenLive14BRequiredComponentsTesterMixin -): - profile = "14b" - - def test_zero_proxy_changes_unconditional_prediction(self): - pipe = self.get_pipeline().to(torch_device) - pipe.register_to_config(proxy_unconditional="zero") - inputs = self.get_dummy_inputs() - inputs["output_type"] = "latent" - zero_proxy = pipe(**inputs).frames - pipe.register_to_config(proxy_unconditional="same") - inputs = self.get_dummy_inputs() - inputs["output_type"] = "latent" - same_proxy = pipe(**inputs).frames - assert (zero_proxy - same_proxy).abs().max() > 1e-6 - - -class TestEdenLiveUnifiedInterface: - @pytest.mark.parametrize( - "profile,steps,guidance", - [("5b-bidir", 48, 3), ("14b-bidir", 30, 3), ("5b-causal", 4, 1), ("14b-causal", 4, 1)], - ) - def test_profile_defaults(self, profile, steps, guidance): - fixture = ( - EdenLiveChunkCausalPipelineTesterConfig() if profile.endswith("causal") else EdenLivePipelineTesterConfig() - ) - fixture.profile = profile.split("-")[0] - components = fixture.get_dummy_components() - if fixture.causal: - components["scheduler"].register_to_config(default_timesteps=[1000, 750, 500, 250]) - pipe = fixture.get_pipeline(**components).to(torch_device) - inputs = fixture.get_dummy_inputs() - inputs.pop("num_inference_steps") - inputs.pop("guidance_scale") - inputs.update(num_frames=1, output_type="latent") - indices = [] - - def callback(pipeline, index, timestep, tensors): - indices.append(index) - return {} - - output = pipe(**inputs, callback_on_step_end=callback).frames - assert type(pipe) is EdenLivePipeline and pipe.config.profile == profile - assert pipe.guidance_scale == guidance - assert pipe.num_timesteps == steps - assert indices == list(range(steps)) - assert output.shape == (1, 4, 1, 2, 2) - - @pytest.mark.parametrize( - "profile,argument", - [ - ("5b-bidir", "chunk_prompts"), - ("14b-bidir", "chunk_prompts"), - ("14b-bidir", "text_timeline"), - ("5b-causal", "text_timeline"), - ("14b-causal", "text_timeline"), - ], - ) - def test_incompatible_temporal_prompts_fail_before_inference(self, profile, argument): - fixture = EdenLivePipelineTesterConfig() - fixture.profile, mode = profile.split("-") - fixture.causal = mode == "causal" - pipe = fixture.get_pipeline() - with pytest.raises(ValueError, match=argument): - pipe(**{argument: []}) - - @pytest.mark.parametrize("size", ["5b", "14b"]) - def test_causal_callback_exception_cleans_history(self, size): - fixture = EdenLiveChunkCausalPipelineTesterConfig() - fixture.profile = size - pipe = fixture.get_pipeline().to(torch_device) - inputs = fixture.get_dummy_inputs() - inputs.update(num_frames=13, output_type="latent") - inputs["video"] = np.concatenate([inputs["video"], inputs["video"][:4]]) - states = [] - steps = [] - - def record_state(module, args, kwargs): - state = kwargs["causal_state"] - if all(state is not item for item in states): - states.append(state) - - def fail_on_tail(pipeline, index, timestep, tensors): - steps.append((index, tensors["latents"].shape[2])) - if index == 2: - raise RuntimeError("callback failed on tail") - return {} - - handle = pipe.transformer.register_forward_pre_hook(record_state, with_kwargs=True) - try: - with pytest.raises(RuntimeError, match="callback failed on tail"): - pipe(**inputs, callback_on_step_end=fail_on_tail) - finally: - handle.remove() - assert steps == [(0, 3), (1, 3), (2, 1)] - assert len(states) == 2 and all(state.next_frame == 0 for state in states) - assert not any("cache" in key for key in pipe.transformer.state_dict()) - - @pytest.mark.parametrize("profile", ["5b-bidir", "14b-bidir", "5b-causal", "14b-causal"]) - def test_standard_profile_selector_conflict(self, tmp_path, profile): - fixture = EdenLivePipelineTesterConfig() - fixture.profile, mode = profile.split("-") - fixture.causal = mode == "causal" - pipe = fixture.get_pipeline() - pipe.save_pretrained(tmp_path) - restored = EdenLivePipeline.from_pretrained(tmp_path, profile=profile, local_files_only=True) - assert restored.config.profile == profile - wrong_profile = f"{fixture.profile}-{'bidir' if fixture.causal else 'causal'}" - with pytest.raises(ValueError, match="conflicts with saved profile"): - EdenLivePipeline.from_pretrained(tmp_path, profile=wrong_profile, local_files_only=True) - - -class RecordingDecoder(torch.nn.Module): - scale_factor_temporal = 4 - scale_factor_spatial = 8 - - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.zeros((), dtype=torch.float64)) - self.config = SimpleNamespace(latent_channels=4) - self.inputs = [] - - def decode(self, latents, return_dict=False): - self.inputs.append(latents.detach().clone()) - batch, _, frames, height, width = latents.shape - return (latents.new_zeros((batch, 3, (frames - 1) * 4 + 1, height * 8, width * 8)),) - - -class TestEdenLiveNativeNumerics: - @pytest.mark.parametrize("profile", ["5b", "14b"]) - @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float16, torch.bfloat16]) - def test_dino_input_uses_fp32_profile_preprocessing(self, monkeypatch, profile, input_dtype): - from diffusers.pipelines.edenlive.pipeline_utils import prepare_conditions - - fixture = EdenLivePipelineTesterConfig() - fixture.profile = profile - pipe = fixture.get_pipeline().to(torch_device) - # A bilinear override must not give 14B the 5B normalization round trip. - pipe.register_to_config(proxy_resize_mode="bilinear") - rgb = torch.arange(256).reshape(1, 1, 16, 16).expand(1, 3, -1, -1) - frame = (rgb.float() / 255).to(device=torch_device, dtype=input_dtype) - seen = [] - preprocess = pipe.proxy_video_processor.preprocess_proxy_frames - - def capture(value, *args): - seen.append(value.clone()) - return preprocess(value, *args) - - monkeypatch.setattr(pipe.proxy_video_processor, "preprocess_proxy_frames", capture) - prepare_conditions( - pipe, - fixture.get_dummy_inputs()["image"], - frame.unsqueeze(2), - 16, - 16, - 1, - 1, - torch.device(torch_device), - torch.float32, - ) - assert seen[0].dtype == torch.float32 - if profile == "5b": - # Fixed IEEE-754 values for the 1/255 input after the FP32 round trip. - expected = { - torch.float32: 0.003921568393707275, - torch.float16: 0.0039215087890625, - torch.bfloat16: 0.003936767578125, - }[input_dtype] - assert seen[0][0, 0, 0, 1].item() == expected - else: - torch.testing.assert_close(seen[0], frame.float(), rtol=0, atol=0) - - @pytest.mark.parametrize("guidance", [1.0, 3.0]) - def test_chunk_prompts_encode_unique_texts_individually(self, guidance): - fixture = EdenLiveChunkCausalPipelineTesterConfig() - 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("profile", ["5b", "14b"]) - @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) - def test_clean_commit_timestep_shape_and_cfg(self, profile, dtype): - fixture = EdenLiveChunkCausalPipelineTesterConfig() - fixture.profile = profile - pipe = fixture.get_pipeline().to(torch_device, dtype=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("decoder_kind,expected", [("tiny", 4.0625), ("custom", 4.03125), ("vae", 4.03125)]) - def test_bf16_decoder_denormalization(self, monkeypatch, decoder_kind, expected): - from diffusers.pipelines.edenlive.pipeline_utils import decode_latents - - fixture = EdenLivePipelineTesterConfig() - decoder = ( - EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) - if decoder_kind == "tiny" - else RecordingDecoder() - if decoder_kind == "custom" - else None - ) - pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) - pipe.vae.to(dtype=torch.bfloat16) - pipe.vae.register_to_config(latents_mean=[0.0] * 4, latents_std=[3.0] * 4) - seen = [] - - def decode(value, **kwargs): - seen.append(value.clone()) - return (value.new_zeros(1, 3, 1, 16, 16),) - - monkeypatch.setattr(decoder if decoder is not None else pipe.vae, "decode", decode) - z = torch.full((1, 4, 1, 2, 2), 1.3515625, dtype=torch.bfloat16) - assert decode_latents(pipe, z, "latent") is z and not seen - decode_latents(pipe, z, "pt") - assert len(seen) == 1 and (seen[0] == expected).all() - - def test_example_fsdp_casts_only_time_and_preserves_cache_identity(self, monkeypatch): - from torch.distributed import device_mesh, fsdp - - path = Path(__file__).resolve().parents[3] / "examples/edenlive/inference_edenlive.py" - spec = importlib.util.spec_from_file_location("edenlive_example_test", path) - example = importlib.util.module_from_spec(spec) - spec.loader.exec_module(example) - monkeypatch.setattr(device_mesh, "init_device_mesh", lambda *args, **kwargs: {"ulysses": object()}) - policies = [] - monkeypatch.setattr(fsdp, "fully_shard", lambda module, **kwargs: policies.append(kwargs["mp_policy"])) - monkeypatch.setattr(example, "ContextParallelConfig", lambda **kwargs: kwargs) - - class Recorder(torch.nn.Module): - def __init__(self): - super().__init__() - self.blocks = torch.nn.ModuleList([torch.nn.Identity()]) - - def enable_parallelism(self, **kwargs): - pass - - def forward(self, hidden_states, timestep, *, causal_state): - causal_state.next_frame += 3 - return hidden_states, timestep, causal_state - - model = Recorder() - example.configure_parallelism(model, torch.device("cpu"), 2, True) - assert len(policies) == 2 and all(not p.cast_forward_inputs for p in policies) - state = SimpleNamespace(next_frame=0, layers=[object()]) - hidden = torch.ones(1) - for positional, time in [(True, torch.tensor([999])), (False, torch.tensor([[0.0, 375.0]]))]: - args = (hidden, time) if positional else (hidden,) - kwargs = {} if positional else {"timestep": time} - actual_hidden, actual_time, actual_state = model(*args, **kwargs, causal_state=state) - assert actual_hidden is hidden and actual_state is state - assert actual_time.dtype == torch.bfloat16 - assert actual_time.tolist() == ([1000.0] if positional else [[0.0, 376.0]]) - assert state.next_frame == 6 - - -class TestEdenLiveDecoder: - @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()) - - @pytest.mark.parametrize("causal", [False, True]) - def test_custom_decoder_receives_raw_latents_once_and_latent_output_bypasses_it(self, causal): - from diffusers.pipelines.edenlive.pipeline_utils import decode_latents - - fixture = EdenLivePipelineTesterConfig() - fixture.causal = causal - decoder = RecordingDecoder() - pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) - pipe.vae.register_to_config(latents_mean=[0.25] * 4, latents_std=[2.0] * 4) - latents = torch.full((1, 4, 4, 2, 2), 0.5) - assert decode_latents(pipe, latents, "latent") is latents - assert decoder.inputs == [] - output = decode_latents(pipe, latents, "pt") - assert pipe.decoder is decoder and decoder.weight.dtype == torch.float64 - assert decoder.inputs[0].dtype == torch.float64 - torch.testing.assert_close(decoder.inputs[0], torch.full_like(decoder.inputs[0], 1.25)) - assert output.shape == (1, 13, 3, 16, 16) - - @pytest.mark.parametrize("causal", [False, True]) - @pytest.mark.parametrize("saved_decoder", ["tiny", "none", "legacy"]) - @pytest.mark.parametrize("size", ["5b", "14b"]) - @pytest.mark.parametrize("legacy_class", [False, True]) - def test_standard_snapshot_preserves_decoder_choice(self, tmp_path, causal, saved_decoder, size, legacy_class): - fixture = EdenLivePipelineTesterConfig() - fixture.profile = size - fixture.causal = causal - decoder = ( - EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) - if saved_decoder == "tiny" - else None - ) - pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) - pipe.save_pretrained(tmp_path) - index_path = tmp_path / "model_index.json" - index = json.loads(index_path.read_text()) - assert "_edenlive_original" not in index - if legacy_class: - index["_class_name"] = "EdenLiveChunkCausalPipeline" if causal else "EdenLivePipeline" - index["profile"] = size - if saved_decoder == "legacy": - index.pop("decoder", None) - index["_class_name"] = "EdenLiveChunkCausalPipeline" if causal else "EdenLivePipeline" - index["profile"] = size if legacy_class else pipe.config.profile - index_path.write_text(json.dumps(index)) - hashes = file_digests(tmp_path) - restored = EdenLivePipeline.from_pretrained(tmp_path, local_files_only=True) - assert type(restored) is EdenLivePipeline - assert restored.config.profile == f"{size}-{'causal' if causal else 'bidir'}" - assert (restored.config.chunk_size, restored.config.max_context_frames, restored.config.sink_frames) == ( - 3, - 18, - 6, - ) - assert file_digests(tmp_path) == hashes - if legacy_class: - restored.save_pretrained(tmp_path / "new-snapshot") - updated = json.loads((tmp_path / "new-snapshot/model_index.json").read_text()) - assert updated["_class_name"] == "EdenLivePipeline" - assert updated["profile"] == restored.config.profile - if saved_decoder == "tiny": - assert isinstance(restored.decoder, EdenLiveLightTAEDecoder) - assert restored.decoder.config.latent_channels == 4 - assert tuple(restored.decoder.config.block_out_channels) == (4, 4, 4, 4) - for key, value in decoder.state_dict().items(): - torch.testing.assert_close(restored.decoder.state_dict()[key], value, atol=0, rtol=0) - else: - assert restored.decoder is None - - @pytest.mark.parametrize("causal", [False, True]) - @pytest.mark.parametrize("saved_tiny", [False, True]) - def test_standard_snapshot_accepts_plain_decoder_override(self, tmp_path, causal, saved_tiny): - fixture = EdenLivePipelineTesterConfig() - fixture.causal = causal - decoder = EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) if saved_tiny else None - pipe = fixture.get_pipeline(**fixture.get_dummy_components(), decoder=decoder) - pipe.save_pretrained(tmp_path) - if saved_tiny: - for path in (tmp_path / "decoder").glob("*.safetensors"): - path.unlink() - override = RecordingDecoder() - restored = EdenLivePipeline.from_pretrained( - tmp_path, decoder=override, dtype=torch.bfloat16, local_files_only=True - ) - assert restored.decoder is override and override.weight.dtype == torch.float64 - - -class TestEdenLiveOriginalLoading: - @pytest.mark.parametrize("original_checkpoint", ["5b-bidir"], indirect=True) - def test_inactive_loading_options_are_accepted(self, original_checkpoint): - pipe = EdenLivePipeline.from_pretrained( - original_checkpoint.source, - profile=original_checkpoint.profile, - local_files_only=True, - device_map=None, - variant=None, - use_onnx=False, - offload_state_dict=False, - ) - assert pipe.config.profile == original_checkpoint.profile - - def test_original_profiles_components_and_default_dtype(self, original_checkpoint): - fixture = original_checkpoint - loader = EdenLivePipeline - pipe = loader.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) - assert type(pipe) is fixture.pipeline_class - assert pipe.config.profile == fixture.profile - assert isinstance(pipe.decoder, EdenLiveLightTAEDecoder) - expected_components = {**fixture.components, "decoder": fixture.decoder} - for name, expected in expected_components.items(): - if not isinstance(expected, torch.nn.Module): - continue - actual = getattr(pipe, name) - dtype = torch.float32 if name in ("proxy_encoder", "proxy_projection") else torch.bfloat16 - assert actual.dtype == dtype - assert set(actual.state_dict()) == set(expected.state_dict()) - for key, tensor in expected.state_dict().items(): - restored = actual.state_dict()[key] - torch.testing.assert_close(restored, tensor.to(restored.dtype), atol=0, rtol=0) - if fixture.profile == "5b-bidir": - assert isinstance(pipe.scheduler, UniPCMultistepScheduler) - assert pipe.scheduler.config.flow_shift == 5 - elif fixture.profile.endswith("causal"): - points = [1000, 937.5, 833.3333333333, 625] if fixture.profile.startswith("5b") else [1000, 900, 750, 500] - assert pipe.scheduler.config.stochastic_sampling - assert list(pipe.scheduler.config.default_timesteps) == points - assert (pipe.config.chunk_size, pipe.config.max_context_frames, pipe.config.sink_frames) == (3, 18, 6) - else: - assert pipe.scheduler.config.shift == 3 and not pipe.scheduler.config.stochastic_sampling - assert pipe.config.proxy_unconditional == "same" - expected_resize = "bilinear" if fixture.profile.startswith("5b") else "bicubic" - assert pipe.config.proxy_resize_mode == expected_resize - assert tuple(pipe.config.proxy_image_size) == (480, 832) - assert file_digests(fixture.source) == fixture.hashes - - def test_original_save_load_produces_self_contained_standard_snapshot(self, original_checkpoint, tmp_path): - fixture = original_checkpoint - pipe = EdenLivePipeline.from_pretrained( - fixture.source, profile=fixture.profile, dtype=torch.float32, local_files_only=True - ) - snapshot = tmp_path / "standard" - pipe.save_pretrained(snapshot) - index = json.loads((snapshot / "model_index.json").read_text()) - assert "_edenlive_original" not in index - assert index["decoder"] == ["diffusers", "EdenLiveLightTAEDecoder"] - restored = EdenLivePipeline.from_pretrained(snapshot, local_files_only=True) - assert type(restored) is fixture.pipeline_class - assert isinstance(restored.decoder, EdenLiveLightTAEDecoder) - for name in ("transformer", "vae", "decoder", "proxy_encoder", "proxy_projection"): - expected, actual = getattr(pipe, name).state_dict(), getattr(restored, name).state_dict() - for key, value in expected.items(): - torch.testing.assert_close(actual[key], value, atol=0, rtol=0) - assert file_digests(fixture.source) == fixture.hashes - - @pytest.mark.parametrize("dtype_kind", ["scalar", "mapping"]) - def test_original_dtype_override_and_component_identity(self, original_checkpoint, dtype_kind): - fixture = original_checkpoint - decoder = RecordingDecoder() - text_encoder = fixture.components["text_encoder"].to(dtype=torch.float64) - dtype = torch.float32 if dtype_kind == "scalar" else {"default": torch.float64, "proxy_encoder": torch.float32} - pipe = EdenLivePipeline.from_pretrained( - fixture.source, - profile=fixture.profile, - decoder=decoder, - text_encoder=text_encoder, - dtype=dtype, - local_files_only=True, - ) - assert pipe.decoder is decoder and decoder.weight.dtype == torch.float64 - assert pipe.text_encoder is text_encoder and text_encoder.dtype == torch.float64 - assert pipe.transformer.dtype == (torch.float32 if dtype_kind == "scalar" else torch.float64) - assert pipe.proxy_encoder.dtype == torch.float32 - assert pipe.proxy_projection.dtype == (torch.float32 if dtype_kind == "scalar" else torch.float64) - assert file_digests(fixture.source) == fixture.hashes - - def test_original_decoder_none_skips_missing_tiny_weights(self, original_checkpoint): - fixture = original_checkpoint - size = fixture.profile.split("-")[0] - for path in (fixture.source / f"decoder-{size}").iterdir(): - path.unlink() - with pytest.raises((OSError, ValueError), match="decoder|lighttae|No such file|does not appear"): - EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) - pipe = EdenLivePipeline.from_pretrained( - fixture.source, profile=fixture.profile, decoder=None, local_files_only=True - ) - assert pipe.decoder is None - assert isinstance(pipe.vae, AutoencoderKLWan) - - def test_original_custom_weight_names_and_shared_proxy(self, original_checkpoint): - fixture = original_checkpoint - transformer_path = fixture.transformer_dir / "diffusion_pytorch_model.safetensors" - custom_transformer = transformer_path.with_name("trained-transformer.safetensors") - transformer_path.rename(custom_transformer) - proxy_path = fixture.source / "vision-encoder/projection.safetensors" - custom_proxy = proxy_path.with_name("trained-backbone-and-projector.safetensors") - proxy_path.rename(custom_proxy) - size = fixture.profile.split("-")[0] - decoder_path = next((fixture.source / f"decoder-{size}").glob("*.pth")) - custom_decoder = decoder_path.with_name("trained-decoder.pth") - decoder_path.rename(custom_decoder) - index_path = fixture.source / "model_index.json" - index = json.loads(index_path.read_text()) - index["_edenlive_original"]["_path_overrides"] = { - "transformer": {fixture.profile: {"weights": str(custom_transformer.relative_to(fixture.source))}}, - "proxy_encoder": {"weights": str(custom_proxy.relative_to(fixture.source))}, - "proxy_projection": {"weights": str(custom_proxy.relative_to(fixture.source))}, - "decoder": {size: {"weights": str(custom_decoder.relative_to(fixture.source))}}, - } - index_path.write_text(json.dumps(index)) - hashes = file_digests(fixture.source) - pipe = EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) - torch.testing.assert_close( - pipe.proxy_projection.patch_projection.weight, - fixture.components["proxy_projection"].patch_projection.weight, - atol=0, - rtol=0, - ) - assert isinstance(pipe.decoder, EdenLiveLightTAEDecoder) - assert file_digests(fixture.source) == hashes - - def test_original_profile_validation(self, original_checkpoint): - fixture = original_checkpoint - with pytest.raises(ValueError, match="profile"): - EdenLivePipeline.from_pretrained(fixture.source, local_files_only=True) - - def test_original_download_selected_profile_files(self, original_hub, tmp_path): - hub = original_hub - pipe = EdenLivePipeline.from_pretrained( - hub.repo_id, - profile=hub.fixture.profile, - revision="release-test", - cache_dir=tmp_path / "cache", - token="fixture-token", - local_files_only=False, - force_download=True, - ) - assert type(pipe) is EdenLivePipeline and pipe.config.profile == hub.fixture.profile - assert len(hub.calls.metadata) == len(hub.calls.snapshots) == 1 - options = hub.calls.snapshots[0] - assert set(options["allow_patterns"]) == set(hub.fixture.hashes) - assert options["revision"] == hub.revision - assert options["cache_dir"] == tmp_path / "cache" - assert options["token"] == "fixture-token" and options["force_download"] - assert options["local_files_only"] is False - assert options["user_agent"]["pipeline_class"] == hub.fixture.pipeline_class.__name__ - - @pytest.mark.parametrize("original_checkpoint", ["5b-causal"], indirect=True) - @pytest.mark.parametrize("complete_cache", [False, True]) - def test_original_download_offline_cache(self, original_hub, monkeypatch, tmp_path, complete_cache): - hub = original_hub - - def unexpected_model_info(*args, **kwargs): - pytest.fail("Offline loading must not query the remote repository tree.") - - import huggingface_hub - - monkeypatch.setattr(huggingface_hub, "model_info", unexpected_model_info) - monkeypatch.setattr(huggingface_hub, "get_cached_repo_tree", unexpected_model_info) - if not complete_cache: - metadata_dir = tmp_path / "metadata-only-cache" - metadata_dir.mkdir() - index_path = metadata_dir / "model_index.json" - index_path.write_text((hub.fixture.source / "model_index.json").read_text()) - - def metadata_download(repository, filename, **kwargs): - assert repository == hub.repo_id and filename == "model_index.json" - assert kwargs["local_files_only"] is True - return str(index_path) - - monkeypatch.setattr(hub.module, "hf_hub_download", metadata_download) - pipe = EdenLivePipeline.from_pretrained( - hub.repo_id, profile=hub.fixture.profile, revision="release-test", local_files_only=True - ) - assert type(pipe) is EdenLivePipeline and pipe.config.profile == hub.fixture.profile - if complete_cache: - assert hub.calls.snapshots == [] - assert hub.calls.metadata[0]["local_files_only"] is True - else: - assert len(hub.calls.snapshots) == 1 - assert hub.calls.snapshots[0]["local_files_only"] is True - assert set(hub.calls.snapshots[0]["allow_patterns"]) == set(hub.fixture.hashes) - - @pytest.mark.parametrize("original_checkpoint", ["14b-causal"], indirect=True) - @pytest.mark.parametrize("options", [{"variant": "fp16"}, {"device_map": "balanced"}, {"use_safetensors": True}]) - def test_original_download_rejects_unsupported_before_weights(self, original_hub, options): - hub = original_hub - with pytest.raises(ValueError, match="does not support|includes PyTorch"): - EdenLivePipeline.from_pretrained( - hub.repo_id, profile=hub.fixture.profile, local_files_only=False, force_download=True, **options - ) - assert len(hub.calls.metadata) == 1 - assert hub.calls.snapshots == [] - - @pytest.mark.parametrize("original_checkpoint", ["5b-bidir"], indirect=True) - def test_original_rejects_unknown_proxy_checkpoint_key(self, original_checkpoint): - fixture = original_checkpoint - path = fixture.source / "vision-encoder/projection.safetensors" - state = {key: value.clone() for key, value in load_file(path).items()} - state["module.encoder.unrecognized.weight"] = torch.ones(1) - save_file(state, path) - with pytest.raises(ValueError, match="Unrecognized proxy checkpoint key"): - EdenLivePipeline.from_pretrained(fixture.source, profile=fixture.profile, local_files_only=True) - - -class EdenLiveTransformerTesterConfig: - 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, - "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), - "proxy_features": randn_tensor((batch_size, 8, 3, 4, 4), generator=generator, device=torch_device), - "encoder_hidden_states": randn_tensor((batch_size, 12, 16), generator=generator, device=torch_device), - "timestep": torch.tensor([500.0], device=torch_device).expand(batch_size), - } - - -class TestEdenLiveTransformerModel(EdenLiveTransformerTesterConfig): - @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() - 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 * state.tokens_per_frame for layer in state.layers) - if index == 1: - sink_keys = [layer.key[:, : 6 * state.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 * state.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()) - - @torch.no_grad() - def test_text_timeline_uses_complete_segment_bank(self): - model = self.model_class(**self.get_init_dict()).eval().to(torch_device) - attn = model.blocks[0].attn2 - hidden = randn_tensor((1, 12, 48), generator=self.generator, device=torch_device) - context = randn_tensor((1, 10, 48), generator=self.generator, device=torch_device) - layout = {"lengths": ((3, 9),), "text_length": 5} - expected = torch.cat((attn(hidden[:, :3], context[:, :5]), attn(hidden[:, 3:], context[:, 5:])), dim=1) - actual = attn(hidden, context, text_timeline=layout) - torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) - local = attn(hidden[:, 6:], context, text_timeline=layout, query_start=6) - torch.testing.assert_close(local, expected[:, 6:], atol=1e-6, rtol=1e-6) - - -class TestEdenLiveCausalCacheOrder: - def test_local_chunk_order_survives_eviction_and_short_tail(self, monkeypatch): - # Capture the real processor's dispatch boundary. Each rank keeps its - # own chunk slices; Ulysses performs the rank-major exchange downstream. - processor = transformer_edenlive.EdenLiveCausalAttnProcessor() - identity = torch.nn.Identity() - attn = SimpleNamespace( - fused_projections=False, - to_q=identity, - to_k=identity, - to_v=identity, - norm_q=identity, - norm_k=identity, - heads=1, - to_out=[identity, identity], - ) - cache = transformer_edenlive._EdenLiveLayerCache(torch.zeros(1, 6, 1, 2), torch.zeros(1, 6, 1, 2)) - seen = [] - - def dispatch(query, key, value, **kwargs): - seen.append((key.clone(), value.clone())) - return query - - monkeypatch.setattr(transformer_edenlive, "dispatch_attention_fn", dispatch) - - def call(ids, commit): - hidden = torch.tensor(ids, dtype=torch.float32).view(1, -1, 1).expand(-1, -1, 2) - rope = (torch.ones(1, 1, 1, 2), torch.zeros(1, 1, 1, 2)) - processor( - attn, - hidden, - rotary_emb=rope, - key_rotary_emb=rope, - layer_cache=cache, - cache_update=commit, - cache_window_tokens=6, - cache_sink_tokens=2, - ) - assert torch.equal(seen[-1][0], seen[-1][1]) - return seen[-1][0][0, :, 0, 0].tolist() - - assert call([0, 1], True) == [0, 1] - assert call([4, 5], True) == [0, 1, 4, 5] - assert call([8, 9], True) == [0, 1, 4, 5, 8, 9] - assert call([12, 13], True) == [0, 1, 8, 9, 12, 13] - stored = cache.key.clone() - assert call([16], False) == [0, 1, 9, 12, 13, 16] - assert cache.num_tokens == 6 and torch.equal(cache.key, stored) - assert call([16], True) == [0, 1, 9, 12, 13, 16] - - -def _parallel_outputs(model, inputs): - with torch.no_grad(): - output = model(**inputs, return_dict=False)[0] - state = model.create_causal_state( - batch_size=1, - height=8, - width=8, - chunk_size=3, - max_context_frames=18, - sink_frames=6, - device=inputs["hidden_states"].device, - dtype=inputs["hidden_states"].dtype, - ) - clean_inputs = {**inputs, "timestep": torch.zeros_like(inputs["timestep"])} - sink_keys = None - for index in range(8): - model(**clean_inputs, causal_state=state, cache_update=True) - assert state.next_frame == 3 * (index + 1) - assert all(layer.num_tokens <= 18 * state.tokens_per_frame // state.cp_size for layer in state.layers) - sink_tokens = 6 * state.tokens_per_frame // state.cp_size - if index == 1: - sink_keys = [layer.key[:, :sink_tokens].clone() for layer in state.layers] - elif index > 1: - for expected, layer in zip(sink_keys, state.layers): - torch.testing.assert_close(layer.key[:, :sink_tokens], expected, rtol=0, atol=0) - tail_inputs = { - **inputs, - "hidden_states": inputs["hidden_states"][:, :, :1], - "proxy_features": inputs["proxy_features"][:, :, :1], - "timestep": inputs["timestep"][:, :16] if inputs["timestep"].ndim == 2 else inputs["timestep"], - } - tail = model(**tail_inputs, causal_state=state, return_dict=False)[0] - assert state.next_frame == 24 - tail_clean = {**tail_inputs, "timestep": torch.zeros_like(tail_inputs["timestep"])} - model(**tail_clean, causal_state=state, cache_update=True) - assert state.next_frame == 25 - return [output.float().cpu(), tail.float().cpu()] - - -def _context_parallel_worker(rank, world_size, port, init_dict, state_dict, inputs, backend, use_fsdp, output_path): - # Spawned workers must establish the same settings as the single-device reference. - enable_full_determinism() - torch.set_float32_matmul_precision("highest") - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), RANK=str(rank), WORLD_SIZE=str(world_size)) - torch.cuda.set_device(rank) - device = torch.device("cuda", rank) - dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 - try: - dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=180)) - model = EdenLiveTransformer3DModel(**init_dict).eval() - model.load_state_dict(state_dict) - model.to(device=device, dtype=dtype) - model.set_attention_backend(backend) - inputs = { - name: tensor.to(device=device, dtype=torch.float32 if name == "timestep" else dtype) - for name, tensor in inputs.items() - } - - def check_timestep(_module, _args, kwargs): - assert kwargs["timestep"].dtype == torch.float32 - assert torch.isin(kwargs["timestep"], torch.tensor([0.0, 375.0], device=device)).all() - - model.register_forward_pre_hook(check_timestep, with_kwargs=True) - from torch.distributed.device_mesh import init_device_mesh - - mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) - model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) - if use_fsdp: - from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard - - policy = MixedPrecisionPolicy( - param_dtype=dtype, reduce_dtype=torch.float32, output_dtype=dtype, cast_forward_inputs=False - ) - for block in model.blocks: - fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - fully_shard(model, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - outputs = _parallel_outputs(model, inputs) - if rank == 0: - torch.save(outputs, output_path) - finally: - if dist.is_initialized(): - dist.destroy_process_group() - - -class TestEdenLiveContextParallel(EdenLiveTransformerTesterConfig): - @pytest.mark.context_parallel - @pytest.mark.attention - @pytest.mark.parametrize("world_size", [2, 4]) - @pytest.mark.parametrize( - "backend,use_fsdp", [("native", False), ("_native_flash", False), ("_native_flash", True)] - ) - @pytest.mark.parametrize("num_layers,token_time", [(2, False), (1, True)]) - def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsdp, num_layers, token_time): - from diffusers.models.attention_dispatch import _AttentionBackendRegistry - - if not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < world_size: - pytest.skip(f"This case requires {world_size} CUDA devices with NCCL.") - previous_backend, _ = _AttentionBackendRegistry.get_active_backend() - deterministic = torch.are_deterministic_algorithms_enabled() - warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - try: - enable_full_determinism() - config = {**self.get_init_dict(), "attention_head_dim": 16, "num_layers": num_layers} - model = EdenLiveTransformer3DModel(**config).eval() - state = {name: tensor.cpu().clone() for name, tensor in model.state_dict().items()} - inputs = {name: tensor.cpu() for name, tensor in self.get_dummy_inputs().items()} - # This value would be rounded to 376 by an indiscriminate BF16 input cast. - inputs["timestep"] = torch.tensor([375.0], dtype=torch.float32) - if token_time: - inputs["timestep"] = torch.full((1, 48), 375.0) - inputs["timestep"][:, :16] = 0 - dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 - model.to(device="cuda:0", dtype=dtype).set_attention_backend(backend) - reference_inputs = { - name: tensor.to(device="cuda:0", dtype=torch.float32 if name == "timestep" else dtype) - for name, tensor in inputs.items() - } - expected = _parallel_outputs(model, reference_inputs) - del model, reference_inputs - torch.cuda.empty_cache() - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - port = sock.getsockname()[1] - output_path = str(tmp_path / "context_parallel.pt") - mp.spawn( - _context_parallel_worker, - args=(world_size, port, config, state, inputs, backend, use_fsdp, output_path), - nprocs=world_size, - join=True, - ) - actual = torch.load(output_path, map_location="cpu", weights_only=True) - for reference, output in zip(expected, actual): - assert torch.isfinite(output).all() - # Keep the existing small-model CP backend tolerance. - torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2) - finally: - _AttentionBackendRegistry.set_active_backend(previous_backend) - torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) - - -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) - - -class TestEdenLiveStreaming(EdenLiveChunkCausalPipelineTesterConfig): - def setup_stream(self, profile="5b", output_type="latent", guidance=1.0): - self.profile = profile - components = self.get_dummy_components() - components["decoder"] = EdenLiveLightTAEDecoder(latent_channels=4, block_out_channels=(4, 4, 4, 4)) - pipe = self.get_pipeline(**components).to(torch_device) - inputs = self.get_dummy_inputs() - options = { - "image": inputs["image"], - "prompt": "snow", - "height": 16, - "width": 16, - "num_inference_steps": 2, - "max_sequence_length": 16, - "output_type": output_type, - "guidance_scale": guidance, - } - state = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) - return pipe, state, inputs["video"][:3], options - - @pytest.mark.parametrize("profile", ["5b", "14b"]) - @pytest.mark.parametrize("guidance", [1.0, 2.0]) - def test_eight_chunks_prompt_and_bounded_history(self, profile, guidance): - from unittest.mock import patch - - pipe, state, proxy, _ = self.setup_stream(profile, guidance=guidance) - cache = state.condition_encoder_cache - with patch.object(pipe, "encode_prompt", wraps=pipe.encode_prompt) as encode: - capacities = [] - for index in range(8): - result = pipe(stream_state=state, proxy_frames=proxy, prompt="snow" if index < 4 else "fire") - assert result.frames.shape == (1, 4, 3, 2, 2) - assert torch.isfinite(result.frames).all() - assert all(s.next_frame == (index + 1) * 3 for s in state.causal_states) - capacities.append( - sum(layer.key.numel() + layer.value.numel() for s in state.causal_states for layer in s.layers) - ) - assert encode.call_count == 1 - assert capacities[-1] == capacities[-2] - pipe(stream_state=state, proxy_frames=proxy, prompt_embeds=state.prompt_embeds) - assert state.encoded_prompt is None - pipe(stream_state=state, proxy_frames=proxy, prompt="fire") - assert encode.call_count == 2 - assert state.condition_encoder_cache is cache - assert state.image is None - pipe.close_stream(state) - pipe.close_stream(state) - assert state.closed and not state.valid and not state.causal_states - with pytest.raises(ValueError): - pipe(stream_state=state, proxy_frames=proxy) - - @pytest.mark.parametrize("profile", ["5b", "14b"]) - def test_interleaved_streams_and_dynamic_prompt(self, profile): - pipe, first, proxy, options = self.setup_stream(profile) - second = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) - for _ in range(2): - a = pipe(stream_state=first, proxy_frames=proxy).frames - b = pipe(stream_state=second, proxy_frames=proxy).frames - torch.testing.assert_close(a, b, rtol=1e-4, atol=1e-5) - a = pipe(stream_state=first, proxy_frames=proxy, prompt="fire").frames - b = pipe(stream_state=second, proxy_frames=proxy, prompt="ocean").frames - assert (a - b).abs().max() > 1e-6 - pipe.close_stream(first) - assert pipe(stream_state=second, proxy_frames=proxy).frames.shape[2] == 3 - pipe.close_stream(second) - - @pytest.mark.parametrize("profile", ["5b", "14b"]) - def test_rgb_frames_and_invalid_input_then_execution_failure(self, profile): - from unittest.mock import patch - - pipe, state, proxy, _ = self.setup_stream(profile, output_type="pt") - rng = state.generator.get_state().clone() - for kwargs in ( - {"proxy_frames": proxy[:2]}, - {"proxy_frames": proxy, "latents": torch.zeros(1)}, - {"proxy_frames": proxy, "prompt_embeds": torch.zeros(2, 4, 32)}, - ): - with pytest.raises(ValueError): - pipe(stream_state=state, **kwargs) - assert state.valid and state.causal_states[0].next_frame == 0 - assert torch.equal(rng, state.generator.get_state()) - assert pipe(stream_state=state, proxy_frames=proxy).frames.shape == (1, 9, 3, 16, 16) - assert pipe(stream_state=state, proxy_frames=proxy).frames.shape == (1, 12, 3, 16, 16) - with patch.object(pipe.proxy_encoder, "forward", side_effect=RuntimeError("injected failure")): - with pytest.raises(RuntimeError, match="injected failure"): - pipe(stream_state=state, proxy_frames=proxy) - assert not state.valid - with pytest.raises(ValueError, match="invalid"): - pipe(stream_state=state, proxy_frames=proxy) - pipe.close_stream(state) - - @pytest.mark.parametrize("profile", ["5b", "14b"]) - def test_offline_stream_numerics_with_fixed_noise(self, profile): - from unittest.mock import patch - - from diffusers.pipelines.edenlive import pipeline_utils - - pipe, state, proxy, options = self.setup_stream(profile, guidance=2.0) - full_proxy = np.concatenate( - [proxy[:1].repeat(4, axis=0), proxy[1:2].repeat(4, axis=0), proxy[2:3].repeat(13, axis=0)] - ) - - def noise(pipe, shape, generator, device, dtype): - return torch.full(shape, 0.25, device=device, dtype=dtype) - - with patch.object(pipeline_utils, "sample_noise", side_effect=noise): - expected = pipe(**options, video=full_proxy, num_frames=21).frames - chunks = [ - pipe(stream_state=state, proxy_frames=full_proxy[:9:4]).frames, - pipe(stream_state=state, proxy_frames=full_proxy[12:21:4]).frames, - ] - torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=1e-4, atol=1e-5) - pipe.close_stream(state) - - @pytest.mark.parametrize("patch_size", [1, 2]) - @torch.no_grad() - def test_lighttae_continuity_reset_and_isolation(self, patch_size): - from diffusers import EdenLiveLightTAEDecodeCache - - decoder = EdenLiveLightTAEDecoder( - latent_channels=4, patch_size=patch_size, block_out_channels=(4, 4, 4, 4) - ).to(torch_device) - latents = torch.randn(1, 4, 9, 2, 2, device=torch_device) - expected = decoder.decode(latents).sample - cache = EdenLiveLightTAEDecodeCache() - outputs = [] - for chunk in latents.split(3, dim=2): - outputs.append(decoder.decode(chunk, cache=cache).sample) - decoder.decode(chunk + 1, cache=EdenLiveLightTAEDecodeCache()) - assert [x.shape[2] for x in outputs] == [9, 12, 12] - torch.testing.assert_close(torch.cat(outputs, dim=2), expected, rtol=1e-4, atol=1e-5) - cache.reset() - torch.testing.assert_close(decoder.decode(latents, cache=cache).sample, expected, rtol=1e-4, atol=1e-5) - - def test_proxy_chunk_is_not_subsampled(self): - from diffusers.pipelines.edenlive.video_processor import EdenLiveVideoProcessor - - processor = EdenLiveVideoProcessor() - frames = np.stack([np.full((16, 16, 3), x, dtype=np.float32) for x in (0.0, 0.5, 1.0)]) - result = processor.preprocess_proxy_chunk(frames) - assert result.shape == (1, 3, 3, 16, 16) - torch.testing.assert_close(result[0, 0, :, 0, 0], torch.tensor([0.0, 0.5, 1.0])) - - def test_initialization_validation_and_owner_lifetime(self): - import gc - import weakref - from unittest.mock import patch - - pipe, state, proxy, options = self.setup_stream() - for image in (np.full((16, 16, 3), np.nan, dtype=np.float32), np.ones((16, 16, 3), dtype=np.float32) * 2): - with patch.object(pipe.vae, "encode", side_effect=AssertionError("must validate before encoding")): - with pytest.raises(ValueError, match="reference image"): - pipe.create_stream_state(**{**options, "image": image}) - with patch.object(pipe.decoder, "decode", side_effect=AssertionError("latent output must bypass decode")): - pipe(stream_state=state, proxy_frames=proxy) - pipe.close_stream(state) - owner = weakref.ref(pipe) - del pipe - gc.collect() - assert owner() is None - - def test_rejects_foreign_state_and_component_replacement(self): - pipe, state, proxy, _ = self.setup_stream() - other = self.get_pipeline().to(torch_device) - with pytest.raises(ValueError, match="belong"): - other(stream_state=state, proxy_frames=proxy) - original = pipe.scheduler - pipe.scheduler = EdenLiveFlowMatchEulerDiscreteScheduler(default_timesteps=[1000, 500]) - with pytest.raises(ValueError, match="components"): - pipe(stream_state=state, proxy_frames=proxy) - assert state.valid and state.causal_states[0].next_frame == 0 - pipe.scheduler = original - pipe.close_stream(state) - - @torch.no_grad() - def test_condition_encoding_continuity_isolation_and_root_hook(self): - from unittest.mock import Mock, patch - - from diffusers.pipelines.edenlive.pipeline_edenlive import ( - _EdenLiveConditionEncoderCache, - _encode_condition_chunk, - ) - - self.profile = "14b" - vae = self.get_dummy_components()["vae"].to(torch_device).eval() - video = randn_tensor((1, 3, 33, 16, 16), generator=torch.Generator().manual_seed(12), device=torch_device) - expected = vae.encode(video).latent_dist.mode() - internal_map, internal_index = vae._enc_feat_map, vae._enc_conv_idx - map_values, index_values = list(internal_map), list(internal_index) - cache, other = _EdenLiveConditionEncoderCache(), _EdenLiveConditionEncoderCache() - hook = Mock() - with patch.object(vae, "_hf_hook", hook, create=True): - first = _encode_condition_chunk(vae, video[:, :, :9], cache) - cache_map = cache.feat_map - _encode_condition_chunk(vae, video[:, :, :9] + 1, other) - middle = _encode_condition_chunk(vae, video[:, :, 9:21], cache) - tail = _encode_condition_chunk(vae, video[:, :, 21:], cache) - assert cache.feat_map is cache_map - assert hook.pre_forward.call_count == 4 - hook.pre_forward.assert_called_with(vae) - torch.testing.assert_close(torch.cat([first, middle, tail], dim=2), expected, rtol=1e-4, atol=1e-5) - assert vae._enc_feat_map is internal_map and vae._enc_conv_idx is internal_index - assert vae._enc_feat_map == map_values and vae._enc_conv_idx == index_values - cache.reset() - seed = _encode_condition_chunk(vae, video[:, :, :1], cache) - rest = _encode_condition_chunk(vae, video[:, :, 1:], cache) - torch.testing.assert_close(torch.cat([seed, rest], dim=2), expected, rtol=1e-4, atol=1e-5) - - @torch.no_grad() - def test_condition_encoding_rejects_invalid_chunks_and_memory_modes(self): - from diffusers.pipelines.edenlive.pipeline_edenlive import ( - _EdenLiveConditionEncoderCache, - _encode_condition_chunk, - ) - - self.profile = "14b" - vae = self.get_dummy_components()["vae"].to(torch_device).eval() - video = torch.zeros(1, 3, 9, 16, 16, device=torch_device) - cache = _EdenLiveConditionEncoderCache() - with pytest.raises(ValueError): - _encode_condition_chunk(vae, video[:, :, :4], cache) - assert cache.feat_map is None - _encode_condition_chunk(vae, video[:, :, :1], cache) - with pytest.raises(ValueError): - _encode_condition_chunk(vae, video[:, :, :1], cache) - for mode in ("slicing", "tiling"): - getattr(vae, f"enable_{mode}")() - with pytest.raises(ValueError): - _encode_condition_chunk(vae, video, _EdenLiveConditionEncoderCache()) - getattr(vae, f"disable_{mode}")() - - def test_stream_dispatch_conflicts_are_preflight_and_tuple_output(self): - from unittest.mock import patch - - pipe, state, proxy, options = self.setup_stream("14b", output_type="pt") - rng = state.generator.get_state().clone() - fixed = { - "image": options["image"], - "video": proxy, - "height": 16, - "width": 16, - "num_frames": 9, - "num_inference_steps": 2, - "guidance_scale": 1.0, - "negative_prompt": "", - "negative_prompt_embeds": state.prompt_embeds, - "generator": state.generator, - "text_timeline": {}, - "chunk_prompts": [], - "max_sequence_length": 512, - "output_type": "np", - "attention_kwargs": {}, - "callback_on_step_end": lambda *args: {}, - "callback_on_step_end_tensor_inputs": ["latents"], - } - with patch.object(pipe.proxy_encoder, "forward", side_effect=AssertionError("must validate before encoding")): - for name, value in fixed.items(): - with pytest.raises(ValueError, match=name): - pipe(stream_state=state, proxy_frames=proxy, **{name: value}) - with pytest.raises(ValueError, match="stream_state"): - pipe(proxy_frames=proxy) - with pytest.raises(ValueError): - pipe(stream_state=state) - assert state.valid and state.causal_states[0].next_frame == 0 - assert state.condition_encoder_cache.feat_map is None and not state.decoder_cache.dec_state - assert torch.equal(rng, state.generator.get_state()) - result = pipe(stream_state=state, proxy_frames=proxy, return_dict=False) - assert isinstance(result, tuple) and len(result) == 1 - assert result[0].shape == (1, 9, 3, 16, 16) - pipe.close_stream(state) - - def test_offline_effective_defaults(self): - from unittest.mock import patch - - pipe, state, _, options = self.setup_stream() - pipe.close_stream(state) - inputs = self.get_dummy_inputs() - inputs.pop("output_type") - inputs.pop("max_sequence_length") - with patch.object(pipe, "encode_prompt", wraps=pipe.encode_prompt) as encode: - actual = pipe(**inputs).frames - assert encode.call_args.args[-1] == 512 - inputs = self.get_dummy_inputs() - inputs.update(max_sequence_length=512, output_type="np", callback_on_step_end_tensor_inputs=["latents"]) - expected = pipe(**inputs).frames - assert isinstance(actual, np.ndarray) - np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5) - - @pytest.mark.parametrize("profile", ["5b", "14b"]) - def test_model_cpu_offload_matches_device_placement(self, profile): - if torch.device(torch_device).type == "cpu": - pytest.skip("Model CPU offload requires an accelerator.") - pipe, state, proxy, options = self.setup_stream(profile, output_type="pt", guidance=2.0) - expected = [pipe(stream_state=state, proxy_frames=proxy).frames.cpu() for _ in range(2)] - pipe.close_stream(state) - pipe.enable_model_cpu_offload(device=torch_device) - state = pipe.create_stream_state(**options, generator=torch.Generator().manual_seed(12)) - cache = state.condition_encoder_cache - try: - for chunk in expected: - actual = pipe(stream_state=state, proxy_frames=proxy).frames.cpu() - torch.testing.assert_close(actual, chunk, rtol=1e-4, atol=1e-5) - assert state.condition_encoder_cache is cache - finally: - pipe.close_stream(state) - assert pipe.vae.device.type == "cpu" and pipe.transformer.device.type == "cpu" diff --git a/tests/pipelines/edenlive/test_pipeline_edenlive.py b/tests/pipelines/edenlive/test_pipeline_edenlive.py new file mode 100644 index 000000000000..86111aad8fee --- /dev/null +++ b/tests/pipelines/edenlive/test_pipeline_edenlive.py @@ -0,0 +1,243 @@ +# 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, + EdenLiveProxyProjection, + 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) + proxy_projection = EdenLiveProxyProjection(in_channels=32, out_channels=8) + 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_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, + "proxy_projection": proxy_projection, + "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..6f8fe905f7d7 --- /dev/null +++ b/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py @@ -0,0 +1,246 @@ +# 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, + EdenLivePipeline, + EdenLiveTransformer3DModel, + UniPCMultistepScheduler, +) + +from ...testing_utils import torch_device +from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin +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_components_from_bidirectional_root(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + expected = pipe(**self.get_dummy_inputs()).frames + pipe.transformer.save_pretrained(tmp_path / "transformer_causal") + pipe.scheduler.save_pretrained(tmp_path / "scheduler_causal") + components = self.get_dummy_components() + components["scheduler"] = UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5 + ) + EdenLivePipeline(**components).save_pretrained(tmp_path) + restored = self.pipeline_class.from_pretrained( + tmp_path, + transformer=EdenLiveTransformer3DModel.from_pretrained(tmp_path / "transformer_causal"), + scheduler=EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained(tmp_path / "scheduler_causal"), + ).to(torch_device) + assert restored.decoder is None + 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 = fixture.get_pipeline().to(torch_device, dtype=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_failure_clears_history_and_next_call_is_independent(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 + assert all(state.next_frame == 0 and all(layer.num_tokens == 0 for layer in state.layers) for state in states) + 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) From d4fdc866b6e89555ca698a53cf288bcff0f31396 Mon Sep 17 00:00:00 2001 From: skygoo2000 Date: Wed, 21 Oct 2026 20:05:00 +0800 Subject: [PATCH 6/6] feat: remove custom cp; split causal ckpt; simplify KV cache state; bugfix: device_map, offload, and BF16 handling. --- .../en/api/models/edenlive_transformer_3d.md | 21 +- docs/source/en/api/pipelines/edenlive.md | 36 +-- docs/source/en/api/schedulers/edenlive.md | 6 +- examples/edenlive/README.md | 31 +- examples/edenlive/inference_edenlive.py | 201 +++++-------- examples/edenlive/requirements.txt | 2 +- scripts/convert_edenlive_to_diffusers.py | 65 +++-- src/diffusers/__init__.py | 2 - src/diffusers/models/__init__.py | 2 - src/diffusers/models/transformers/__init__.py | 2 +- .../transformers/transformer_edenlive.py | 225 +++++--------- .../pipelines/edenlive/pipeline_edenlive.py | 122 ++++---- .../edenlive/pipeline_edenlive_causal.py | 276 ++++++++---------- .../pipelines/edenlive/pipeline_utils.py | 17 -- src/diffusers/utils/dummy_pt_objects.py | 15 - .../test_models_edenlive_proxy_projection.py | 91 ------ .../test_models_transformer_edenlive.py | 31 +- ...test_models_transformer_edenlive_causal.py | 235 ++++++--------- .../edenlive/test_pipeline_edenlive.py | 5 +- .../edenlive/test_pipeline_edenlive_causal.py | 26 +- 20 files changed, 523 insertions(+), 888 deletions(-) delete mode 100644 tests/models/transformers/test_models_edenlive_proxy_projection.py diff --git a/docs/source/en/api/models/edenlive_transformer_3d.md b/docs/source/en/api/models/edenlive_transformer_3d.md index aa345cf42a3c..d0bac2de858a 100644 --- a/docs/source/en/api/models/edenlive_transformer_3d.md +++ b/docs/source/en/api/models/edenlive_transformer_3d.md @@ -12,9 +12,9 @@ specific language governing permissions and limitations under the License. # EdenLiveTransformer3DModel -[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) uses a video diffusion transformer conditioned on -text, a reference image, and projected DINO proxy features. The same model class supports bidirectional and causal -attention. See the [pipeline guide](../pipelines/edenlive) for video generation. +[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 @@ -25,8 +25,11 @@ transformer = EdenLiveTransformer3DModel.from_pretrained( ) ``` -Use `subfolder="transformer_causal"` for the causal weights. Keep model timesteps in FP32 when calling the transformer -directly, including when its weights are BF16. +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 @@ -34,14 +37,6 @@ directly, including when its weights are BF16. - forward - create_causal_state -## EdenLiveProxyProjection - -Projects DINO features into the transformer's conditioning space. Use the trained projection and DINO encoder from -the same checkpoint; the projection returns FP32 features. - -[[autodoc]] EdenLiveProxyProjection - - forward - ## EdenLiveCausalState The causal pipeline manages this state automatically. For custom loops, create a separate state for each CFG branch diff --git a/docs/source/en/api/pipelines/edenlive.md b/docs/source/en/api/pipelines/edenlive.md index 830d9a0525ff..f4791bb2ef04 100644 --- a/docs/source/en/api/pipelines/edenlive.md +++ b/docs/source/en/api/pipelines/edenlive.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. # EdenLive -[EdenLive](https://github.com/VAST-AI-Research/EdenLive-Official) generates video from a proxy video, a reference +[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. @@ -28,7 +28,7 @@ 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, "proxy_projection": torch.float32}, + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, ).to("cuda") image = load_image("reference.png") @@ -43,25 +43,19 @@ frames = pipe( export_to_video(frames, "edenlive.mp4", fps=24) ``` -Keep the proxy encoder and projection in FP32 when loading the other models in BF16, as shown above. +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 transformer and scheduler from the same repository. This works with either model size and uses -four fixed denoising steps per chunk. +Load the causal pipeline from its dedicated repository. Both model sizes use four fixed denoising steps per chunk. ```python -from diffusers import EdenLiveCausalPipeline, EdenLiveTransformer3DModel, EdenLiveFlowMatchEulerDiscreteScheduler +from diffusers import EdenLiveCausalPipeline causal_pipe = EdenLiveCausalPipeline.from_pretrained( - "VAST-AI/EdenLive-5B-Diffusers", - transformer=EdenLiveTransformer3DModel.from_pretrained( - "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer_causal", torch_dtype=torch.bfloat16 - ), - scheduler=EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( - "VAST-AI/EdenLive-5B-Diffusers", subfolder="scheduler_causal" - ), - torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, + "VAST-AI/EdenLive-5B-Causal-Diffusers", + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, ).to("cuda") ``` @@ -99,16 +93,18 @@ pipe = EdenLivePipeline.from_pretrained( decoder=EdenLiveLightTAEDecoder.from_pretrained( "VAST-AI/EdenLive-5B-Diffusers", subfolder="decoder", torch_dtype=torch.bfloat16 ), - torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, + torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, ).to("cuda") ``` ## Notes -| Checkpoint | Default height × width | Example fps | Bidirectional steps | -|---|---|---|---| -| [5B](https://huggingface.co/VAST-AI/EdenLive-5B-Diffusers) | 704 × 1280 | 24 | 48 | -| [14B](https://huggingface.co/VAST-AI/EdenLive-14B-Diffusers) | 480 × 832 | 16 | 30 | +| 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. @@ -116,7 +112,7 @@ pipe = EdenLivePipeline.from_pretrained( 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, multi-GPU inference, and checkpoint conversion. + video resampling, CPU offload, and checkpoint conversion. ## EdenLivePipeline diff --git a/docs/source/en/api/schedulers/edenlive.md b/docs/source/en/api/schedulers/edenlive.md index f4a2c86165c8..94bbe1124b11 100644 --- a/docs/source/en/api/schedulers/edenlive.md +++ b/docs/source/en/api/schedulers/edenlive.md @@ -16,9 +16,9 @@ specific language governing permissions and limitations under the License. bidirectional generation and stochastic sampling for the distilled causal checkpoints. The 5B bidirectional model uses [`UniPCMultistepScheduler`]. -Use the scheduler saved with the checkpoint: `scheduler` for bidirectional generation and `scheduler_causal` for -causal generation. The released causal schedules have four fixed steps per chunk. See the -[pipeline guide](../pipelines/edenlive#causal-generation) for loading the causal components. +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 diff --git a/examples/edenlive/README.md b/examples/edenlive/README.md index 9e360474e60e..f69c109be028 100644 --- a/examples/edenlive/README.md +++ b/examples/edenlive/README.md @@ -18,33 +18,29 @@ The [requirements file](requirements.txt) also lists the example dependencies. ```bash python examples/edenlive/inference_edenlive.py \ - --model VAST-AI/EdenLive-5B-Diffusers --causal \ + --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` for the causal weights or omit it for bidirectional generation. `--model` also accepts a local +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 or `--vae-tiling` to reduce VAE memory use. Run +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. -## Multiple GPUs +## Model repositories -Use Ulysses context parallelism and FSDP2 on two or four GPUs: +| 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) | -```bash -torchrun --standalone --nproc_per_node=2 examples/edenlive/inference_edenlive.py \ - --model VAST-AI/EdenLive-14B-Diffusers --causal \ - --video proxy.mp4 --image reference.png \ - --prompt "A character explores a snowy landscape." \ - --num-frames 81 -``` - -Set `--nproc_per_node=4` for four GPUs. Add `--no-fsdp` to keep transformer weights replicated across GPUs. +Each repository contains one complete pipeline. Select the matching repository and pipeline variant. ## Changing prompts @@ -68,8 +64,9 @@ To convert a local original release instead of loading the published Diffusers w ```bash python scripts/convert_edenlive_to_diffusers.py \ - --source /path/to/EdenLive --size 5b --output /path/to/EdenLive-5B-Diffusers + --source /path/to/EdenLive --size 5b --variant bidirectional \ + --output /path/to/EdenLive-5B-Diffusers ``` -Use `--size 14b` and a separate output directory for 14B. Each package includes both sampling variants and the -optional LightTAE decoder. +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 index e368eb0e4af3..015b7ee6a6c9 100644 --- a/examples/edenlive/inference_edenlive.py +++ b/examples/edenlive/inference_edenlive.py @@ -12,27 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline EdenLive inference, optionally using Ulysses and FSDP2 through torchrun.""" +"""Offline EdenLive inference on one GPU, with optional CPU offload.""" import argparse import json import math -import os import time from pathlib import Path import numpy as np import torch -import torch.distributed as dist from PIL import Image from diffusers import ( - ContextParallelConfig, EdenLiveCausalPipeline, - EdenLiveFlowMatchEulerDiscreteScheduler, EdenLiveLightTAEDecoder, EdenLivePipeline, - EdenLiveTransformer3DModel, ) from diffusers.utils import export_to_video @@ -76,45 +71,10 @@ def load_video(path, fps, height, width, num_frames=None): return frames, {"source_frame_indices": indices.tolist(), "fps": fps, "num_frames": count} -def configure_parallelism(transformer, device, world_size, use_fsdp): - if world_size == 1: - transformer.to(device) - return - if world_size not in (2, 4): - raise ValueError("This example supports 1, 2, or 4 processes.") - from torch.distributed.device_mesh import init_device_mesh - from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard - - mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) - transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) - if use_fsdp: - policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - output_dtype=torch.bfloat16, - cast_forward_inputs=False, - ) - for block in transformer.blocks: - fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - fully_shard(transformer, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - - def native_timestep(module, args, kwargs): - # Match native FSDP time quantization without copying mutable cache dataclasses. - if "timestep" in kwargs: - kwargs = {**kwargs, "timestep": kwargs["timestep"].float().to(torch.bfloat16)} - elif len(args) > 1: - args = (args[0], args[1].float().to(torch.bfloat16), *args[2:]) - return args, kwargs - - transformer.register_forward_pre_hook(native_timestep, with_kwargs=True, prepend=True) - else: - transformer.to(device) - - 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="Use the causal transformer and scheduler subfolders.") + 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." ) @@ -132,7 +92,7 @@ def parse_args(): 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("--no-fsdp", action="store_true", help="Use replicated DiT weights with context parallelism.") + 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() @@ -145,100 +105,69 @@ def parse_args(): def main(): args = parse_args() - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size not in (1, 2, 4): - raise ValueError("Launch this example with 1, 2, or 4 processes.") - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) - if world_size > 1: - dist.init_process_group("nccl", device_id=device) - rank = dist.get_rank() if dist.is_initialized() else 0 - try: - loading_kwargs = {"local_files_only": args.local_files_only} - components = {"decoder": None} - pipeline_class = EdenLiveCausalPipeline if args.causal else EdenLivePipeline - if args.causal: - components["transformer"] = EdenLiveTransformer3DModel.from_pretrained( - args.model, subfolder="transformer_causal", torch_dtype=torch.bfloat16, **loading_kwargs - ) - components["scheduler"] = EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( - args.model, subfolder="scheduler_causal", **loading_kwargs - ) - 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, "proxy_projection": torch.float32}, - **components, - **loading_kwargs, + 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 ) - if args.attention_backend is not None: - pipeline.transformer.set_attention_backend(args.attention_backend) - configure_parallelism(pipeline.transformer, device, world_size, not args.no_fsdp) - for name, component in pipeline.components.items(): - if name != "transformer" and isinstance(component, torch.nn.Module): - component.to(device=device) - if args.vae_tiling: - pipeline.vae.enable_tiling() - pipeline.set_progress_bar_config(disable=rank != 0) - 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 = { - "rank": rank, - "seconds": time.perf_counter() - started, - "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), - "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), - } - all_metrics = [None] * world_size - if world_size > 1: - dist.all_gather_object(all_metrics, metrics) - else: - all_metrics[0] = metrics - if rank == 0: - 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, - "world_size": world_size, - "fsdp2": world_size > 1 and not args.no_fsdp, - "metrics": all_metrics, - } - args.output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n") - finally: - if dist.is_initialized(): - dist.destroy_process_group() + 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__": diff --git a/examples/edenlive/requirements.txt b/examples/edenlive/requirements.txt index 9935e2adae5f..f68559fc7f2b 100644 --- a/examples/edenlive/requirements.txt +++ b/examples/edenlive/requirements.txt @@ -1,6 +1,6 @@ torch torchvision -transformers +transformers>=4.56.0 accelerate av imageio diff --git a/scripts/convert_edenlive_to_diffusers.py b/scripts/convert_edenlive_to_diffusers.py index 6af78c18a7f5..090420b21790 100644 --- a/scripts/convert_edenlive_to_diffusers.py +++ b/scripts/convert_edenlive_to_diffusers.py @@ -35,7 +35,6 @@ AutoencoderKLWan, EdenLiveFlowMatchEulerDiscreteScheduler, EdenLiveLightTAEDecoder, - EdenLiveProxyProjection, EdenLiveTransformer3DModel, UniPCMultistepScheduler, ) @@ -129,23 +128,34 @@ def _remap_wan22_vae_key(key): return key -def convert(source, output, size): +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 = {} - for folder, destination in ((f"dit-{size}-bidir", "transformer"), (f"dit-{size}", "transformer_causal")): - 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_channels"] = state["patch_embedding.weight"].shape[1] - config["in_channels"] - config["condition_rmsnorm"] = "patch_embedding.condition_norm.weight" in state - 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 / destination, manifest) - del model, state - gc.collect() + 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" @@ -182,13 +192,7 @@ def convert(source, output, size): key = f"model.{key}" backbone[key] = value save_component(model, backbone, output / "proxy_encoder", manifest) - del model, backbone - projection = {key: value for key, value in state.items() if key.startswith("patch_projection.")} - weight = projection["patch_projection.weight"] - with torch.device("meta"): - model = EdenLiveProxyProjection(in_channels=weight.shape[1], out_channels=weight.shape[0]) - save_component(model, projection, output / "proxy_projection", manifest) - del model, state, projection, weight + del model, state, backbone gc.collect() for name in ("text_encoder", "tokenizer"): @@ -227,31 +231,31 @@ def convert(source, output, size): 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") - EdenLiveFlowMatchEulerDiscreteScheduler(stochastic_sampling=True, default_timesteps=timesteps).save_pretrained( - output / "scheduler_causal" - ) index = { - "_class_name": "EdenLivePipeline", + "_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"], - "proxy_projection": ["diffusers", "EdenLiveProxyProjection"], "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", "proxy_projection")), + ("vision-encoder", ("proxy_encoder", "transformer")), ("clip", ("image_encoder", "image_processor") if size == "14b" else ()), ): for destination in destinations: @@ -267,5 +271,6 @@ def convert(source, output, size): "--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) + convert(args.source, args.output, args.size, args.variant) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 6b81504cceff..f859454abb06 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -300,7 +300,6 @@ "EasyAnimateTransformer3DModel", "EdenLiveCausalState", "EdenLiveLightTAEDecoder", - "EdenLiveProxyProjection", "EdenLiveTransformer3DModel", "ErnieImageTransformer2DModel", "Flux2Transformer2DModel", @@ -1223,7 +1222,6 @@ EasyAnimateTransformer3DModel, EdenLiveCausalState, EdenLiveLightTAEDecoder, - EdenLiveProxyProjection, EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index adaca45dcd5b..61b238152846 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -123,7 +123,6 @@ _import_structure["transformers.transformer_easyanimate"] = ["EasyAnimateTransformer3DModel"] _import_structure["transformers.transformer_edenlive"] = [ "EdenLiveCausalState", - "EdenLiveProxyProjection", "EdenLiveTransformer3DModel", ] _import_structure["transformers.transformer_ernie_image"] = ["ErnieImageTransformer2DModel"] @@ -277,7 +276,6 @@ DualTransformer2DModel, EasyAnimateTransformer3DModel, EdenLiveCausalState, - EdenLiveProxyProjection, EdenLiveTransformer3DModel, ErnieImageTransformer2DModel, Flux2Transformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index c55fc7d13901..876353d7c842 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -31,7 +31,7 @@ from .transformer_cosmos import CosmosTransformer3DModel from .transformer_cosmos3 import Cosmos3OmniTransformer from .transformer_easyanimate import EasyAnimateTransformer3DModel - from .transformer_edenlive import EdenLiveCausalState, EdenLiveProxyProjection, EdenLiveTransformer3DModel + 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 index 3acc4a381e8c..10b091d763f1 100644 --- a/src/diffusers/models/transformers/transformer_edenlive.py +++ b/src/diffusers/models/transformers/transformer_edenlive.py @@ -24,8 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ...utils import apply_lora_scale -from ...utils.torch_utils import maybe_adjust_dtype_for_device, maybe_allow_in_graph -from .._modeling_parallel import ContextParallelInput, ContextParallelOutput +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 @@ -70,18 +69,12 @@ class _EdenLiveLayerCache: @dataclass class EdenLiveCausalState: - """Per-rollout, per-guidance-branch KV storage. Positions count latent frames.""" + """KV storage for one fixed-resolution rollout and guidance branch. Positions count latent frames.""" layers: list[_EdenLiveLayerCache] - batch_size: int - height: int - width: int chunk_size: int max_context_frames: int sink_frames: int - tokens_per_frame: int - cp_size: int = 1 - cp_rank: int = 0 next_frame: int = 0 valid: bool = True @@ -95,7 +88,6 @@ def reset(self) -> None: class EdenLiveAttnProcessor: _attention_backend = None - _parallel_config = None def __call__( self, @@ -103,7 +95,6 @@ def __call__( hidden_states, encoder_hidden_states=None, attention_mask=None, - rotary_emb=None, *, text_context_length=512, ): @@ -121,10 +112,6 @@ def __call__( 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) - key = _apply_rotary_emb(key, rotary_emb) - image_output = None if image_context is not None: image_key, image_value = _get_added_kv_projections(attn, image_context) @@ -139,7 +126,6 @@ def __call__( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=None, ) .flatten(2, 3) .type_as(query) @@ -153,7 +139,6 @@ def __call__( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config if encoder_hidden_states is None else None, ) output = output.flatten(2, 3).type_as(query) if image_output is not None: @@ -161,7 +146,9 @@ def __call__( return attn.to_out[1](attn.to_out[0](output)) -class EdenLiveCausalAttnProcessor(EdenLiveAttnProcessor): +class EdenLiveCausalAttnProcessor: + _attention_backend = None + def __call__( self, attn, @@ -176,16 +163,29 @@ def __call__( cache_window_tokens=None, cache_sink_tokens=0, ): - if layer_cache is None: - return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, rotary_emb) if encoder_hidden_states is not None or attention_mask is not None: - raise ValueError("Causal KV storage is only used by unmasked self-attention.") + 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)) - query = _apply_rotary_emb(query, rotary_emb) + 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 @@ -225,7 +225,6 @@ def __call__( ) key = _apply_rotary_emb(key, key_rotary_emb) - # Ulysses retains the native rank-major KV order, including across chunks. output = dispatch_attention_fn( query, key, @@ -234,51 +233,11 @@ def __call__( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, ) output = output.flatten(2, 3).type_as(query) return attn.to_out[1](attn.to_out[0](output)) -class EdenLiveProxyProjection(ModelMixin, ConfigMixin): - r""" - Learned DINO patch projection, separate from the Transformers vision backbone. - - Args: - in_channels (`int`, defaults to `1024`): - Number of input DINO feature channels. - out_channels (`int`, defaults to `1024`): - Number of projected proxy feature channels. - """ - - _no_split_modules = [] - - @register_to_config - def __init__(self, in_channels: int = 1024, out_channels: int = 1024): - super().__init__() - self.patch_projection = nn.Linear(in_channels, out_channels) - - def forward( - self, hidden_states: torch.Tensor, return_dict: bool = True - ) -> Transformer2DModelOutput | tuple[torch.Tensor]: - r""" - Args: - hidden_states (`torch.Tensor`): - DINO patch features with `in_channels` in the final dimension. - return_dict (`bool`, defaults to `True`): - Whether to return a structured output instead of a tuple. - - Returns: - [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: - Unit-normalized projected features in float32, with `out_channels` in the final dimension. - """ - hidden_states = self.patch_projection(hidden_states) - hidden_states = F.normalize(hidden_states, p=2, dim=-1).float() - if not return_dict: - return (hidden_states,) - return Transformer2DModelOutput(sample=hidden_states) - - 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) @@ -390,16 +349,14 @@ def unfuse_projections(self): self.fused_projections = False - # Copied from diffusers.models.transformers.transformer_wan.WanAttention.forward def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, - rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: - return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, **kwargs) class EdenLiveImageEmbedding(torch.nn.Module): @@ -639,6 +596,8 @@ class EdenLiveTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Atte 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. """ @@ -647,18 +606,8 @@ class EdenLiveTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Atte _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"] - _cp_plan = { - "rope": { - 0: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), - 1: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), - }, - "condition_embedder": { - 1: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), - }, - "blocks.0": {"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3)}, - "head_input": ContextParallelOutput(gather_dim=1, expected_dims=3), - } @register_to_config def __init__( @@ -678,6 +627,7 @@ def __init__( 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__() @@ -685,13 +635,14 @@ def __init__( 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 proxy_channels < 1 or num_layers < 1: - raise ValueError("proxy_channels and num_layers must be positive.") + 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 ) @@ -704,26 +655,11 @@ def __init__( for _ in range(num_layers) ] ) - self.head_input = nn.Identity() 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 _context_parallel_info(self): - config = self._parallel_config - if config is None or config.context_parallel_config is None: - return 1, 0 - config = config.context_parallel_config - if config.ring_degree != 1 or config.ulysses_anything or config.ring_anything: - raise ValueError( - "EdenLive supports equal-sized Ulysses context parallelism without Ring or Anything mode." - ) - size = config.ulysses_degree - if self.config.num_attention_heads % size: - raise ValueError("num_attention_heads must be divisible by the Ulysses degree.") - return size, config._ulysses_local_rank - def create_causal_state( self, batch_size: int, @@ -736,7 +672,7 @@ def create_causal_state( dtype=None, ) -> EdenLiveCausalState: r""" - Allocate independent KV storage for one guidance branch and rollout. + 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. @@ -745,11 +681,13 @@ def create_causal_state( 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; defaults to the model device. + 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 the current context-parallel layout. + `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: @@ -760,34 +698,31 @@ def create_causal_state( 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.") - cp_size, cp_rank = self._context_parallel_info() tokens_per_frame = (height // patch_height) * (width // patch_width) - if tokens_per_frame % cp_size: - raise ValueError("Latent spatial token count must be divisible by the Ulysses degree.") shape = ( batch_size, - max_context_frames * tokens_per_frame // cp_size, + max_context_frames * tokens_per_frame, self.config.num_attention_heads, self.config.attention_head_dim, ) - layers = [ - _EdenLiveLayerCache( - torch.zeros(shape, device=device or self.device, dtype=dtype or self.dtype), - torch.zeros(shape, device=device or self.device, dtype=dtype or self.dtype), + 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), + ) ) - for _ in self.blocks - ] return EdenLiveCausalState( layers, - batch_size, - height, - width, chunk_size, max_context_frames, sink_frames, - tokens_per_frame, - cp_size, - cp_rank, ) @apply_lora_scale("attention_kwargs") @@ -818,8 +753,9 @@ def forward( encoder_hidden_states_image (`torch.Tensor`, *optional*): CLIP image embeddings required by the 14B model. proxy_features (`torch.Tensor`, *optional*): - Required proxy features of shape `(batch_size, proxy_channels, num_frames, proxy_height, proxy_width)`. - Spatial dimensions are interpolated to match the latent input. + 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`): @@ -844,7 +780,7 @@ def forward( ) 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_channels): + 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.") @@ -855,11 +791,8 @@ def forward( 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.") - cp_size, cp_rank = self._context_parallel_info() tokens_per_frame = (height // patch_height) * (width // patch_width) num_tokens = num_frames * tokens_per_frame - if num_tokens % cp_size: - raise ValueError("The input token count must be divisible by the Ulysses degree.") text_context_length = encoder_hidden_states.shape[1] key_rotary_emb = None cache_window_tokens = None @@ -871,39 +804,35 @@ def forward( state = causal_state if not state.valid: raise ValueError("A previous cache commit failed. Reset or recreate the causal state.") - if (batch_size, height, width, cp_size, cp_rank) != ( - state.batch_size, - state.height, - state.width, - state.cp_size, - state.cp_rank, - ) or len(state.layers) != len(self.blocks): - raise ValueError("The causal state does not match this model's batch, latent dimensions or CP layout.") + 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 // cp_size - cache_sink_tokens = state.sink_frames * tokens_per_frame // cp_size + 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) - if cp_size > 1: - # Each rank caches its contiguous slice of every chunk. Keep the same - # chunk-local order for RoPE; Ulysses exchanges the sharded KV tensors. - chunk_tokens = state.chunk_size * tokens_per_frame - key_rotary_emb = tuple( - torch.cat( - [ - chunk.narrow(1, cp_rank * (chunk.shape[1] // cp_size), chunk.shape[1] // cp_size) - for chunk in frequency.split(chunk_tokens, dim=1) - ], - dim=1, - ) - for frequency in key_rotary_emb - ) 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( @@ -948,16 +877,16 @@ def forward( else: hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, **kwargs) - hidden_states = self.head_input(hidden_states) 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 = ( - self.norm_out(hidden_states.float()) * (1 + scale.to(hidden_states.device)) - + shift.to(hidden_states.device) - ).type_as(hidden_states) + 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, diff --git a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py index 3c34cfab06d9..7cb5e727b8cf 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_edenlive.py +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive.py @@ -21,13 +21,14 @@ import torch.nn.functional as F from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel -from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel +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, sample_noise +from .pipeline_utils import prepare_model_input from .video_processor import EdenLiveVideoProcessor @@ -44,7 +45,7 @@ >>> 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, "proxy_projection": torch.float32}, + ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32}, ... ).to("cuda") >>> frames = pipe( ... image=load_image("reference.png"), video=load_video("proxy.mp4"), @@ -71,8 +72,6 @@ class EdenLivePipeline(DiffusionPipeline): Scheduler corresponding to the transformer checkpoint. proxy_encoder (`DINOv3ViTModel`): Trained RGB proxy encoder. - proxy_projection ([`EdenLiveProxyProjection`]): - Learned projection of DINO patch features. image_encoder (`CLIPVisionModel`, *optional*): Reference image encoder required by the 14B model. image_processor (`CLIPImageProcessor`, *optional*): @@ -85,7 +84,7 @@ class EdenLivePipeline(DiffusionPipeline): Width of the proxy images passed to DINO. """ - model_cpu_offload_seq = "text_encoder->image_encoder->proxy_encoder->proxy_projection->transformer->vae->decoder" + 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"] @@ -97,7 +96,6 @@ def __init__( transformer: EdenLiveTransformer3DModel, scheduler: EdenLiveFlowMatchEulerDiscreteScheduler | UniPCMultistepScheduler, proxy_encoder: DINOv3ViTModel, - proxy_projection: EdenLiveProxyProjection, image_encoder: CLIPVisionModel | None = None, image_processor: CLIPImageProcessor | None = None, decoder: EdenLiveLightTAEDecoder | None = None, @@ -112,7 +110,6 @@ def __init__( transformer=transformer, scheduler=scheduler, proxy_encoder=proxy_encoder, - proxy_projection=proxy_projection, image_encoder=image_encoder, image_processor=image_processor, decoder=decoder, @@ -231,7 +228,7 @@ def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, 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 sample_noise(self, shape, generator, device, 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.""" @@ -296,9 +293,7 @@ def prepare_conditions(self, image, proxy_video, height, width, num_frames, batc hidden = hidden[:, start : start + grid_height * grid_width].reshape( batch_size, grid_height, grid_width, -1 ) - with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): - projected = self.proxy_projection(hidden, return_dict=False)[0] - proxy_features.append(projected.float().permute(0, 3, 1, 2)) + 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 @@ -380,8 +375,7 @@ def __call__( 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. Under context parallelism, the - first rank generates and broadcasts each complete noise tensor. + 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. @@ -445,64 +439,62 @@ def __call__( ) if batch_size < 1: raise ValueError("Provide a positive prompt or prompt embeddings.") - try: - 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( + 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, - prompt_embeds, + negative_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) - finally: - self.maybe_free_model_hooks() + 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 index 07ee8b9ca288..45fb9b86b76a 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py +++ b/src/diffusers/pipelines/edenlive/pipeline_edenlive_causal.py @@ -21,13 +21,14 @@ import torch.nn.functional as F from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, DINOv3ViTModel, UMT5EncoderModel -from ...models import AutoencoderKLWan, EdenLiveLightTAEDecoder, EdenLiveProxyProjection, EdenLiveTransformer3DModel +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, sample_noise +from .pipeline_utils import prepare_model_input from .video_processor import EdenLiveVideoProcessor @@ -40,17 +41,11 @@ Examples: ```python >>> import torch - >>> from diffusers import EdenLiveCausalPipeline, EdenLiveTransformer3DModel, EdenLiveFlowMatchEulerDiscreteScheduler + >>> from diffusers import EdenLiveCausalPipeline >>> from diffusers.utils import export_to_video, load_image, load_video - >>> transformer = EdenLiveTransformer3DModel.from_pretrained( - ... "VAST-AI/EdenLive-5B-Diffusers", subfolder="transformer_causal", torch_dtype=torch.bfloat16 - ... ) - >>> scheduler = EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained( - ... "VAST-AI/EdenLive-5B-Diffusers", subfolder="scheduler_causal" - ... ) >>> pipe = EdenLiveCausalPipeline.from_pretrained( - ... "VAST-AI/EdenLive-5B-Diffusers", transformer=transformer, scheduler=scheduler, - ... torch_dtype={"default": torch.bfloat16, "proxy_encoder": torch.float32, "proxy_projection": torch.float32}, + ... "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"), @@ -77,8 +72,6 @@ class EdenLiveCausalPipeline(DiffusionPipeline): Scheduler corresponding to the transformer checkpoint. proxy_encoder (`DINOv3ViTModel`): Trained RGB proxy encoder. - proxy_projection ([`EdenLiveProxyProjection`]): - Learned projection of DINO patch features. image_encoder (`CLIPVisionModel`, *optional*): Reference image encoder required by the 14B model. image_processor (`CLIPImageProcessor`, *optional*): @@ -97,7 +90,7 @@ class EdenLiveCausalPipeline(DiffusionPipeline): Initial latent frames retained when the attention window fills. """ - model_cpu_offload_seq = "text_encoder->image_encoder->proxy_encoder->proxy_projection->transformer->vae->decoder" + 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"] @@ -109,7 +102,6 @@ def __init__( transformer: EdenLiveTransformer3DModel, scheduler: EdenLiveFlowMatchEulerDiscreteScheduler, proxy_encoder: DINOv3ViTModel, - proxy_projection: EdenLiveProxyProjection, image_encoder: CLIPVisionModel | None = None, image_processor: CLIPImageProcessor | None = None, decoder: EdenLiveLightTAEDecoder | None = None, @@ -127,7 +119,6 @@ def __init__( transformer=transformer, scheduler=scheduler, proxy_encoder=proxy_encoder, - proxy_projection=proxy_projection, image_encoder=image_encoder, image_processor=image_processor, decoder=decoder, @@ -253,7 +244,7 @@ def prepare_latents(self, batch_size, height, width, num_frames, dtype, device, 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 sample_noise(self, shape, generator, device, 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): @@ -319,9 +310,7 @@ def prepare_conditions(self, image, proxy_video, height, width, num_frames, batc hidden = hidden[:, start : start + grid_height * grid_width].reshape( batch_size, grid_height, grid_width, -1 ) - with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): - projected = self.proxy_projection(hidden, return_dict=False)[0] - proxy_features.append(projected.float().permute(0, 3, 1, 2)) + 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 @@ -405,8 +394,7 @@ def __call__( 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. Under context parallelism, the - first rank generates and broadcasts each complete noise tensor. + 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. @@ -490,154 +478,144 @@ def __call__( ) if batch_size < 1: raise ValueError("Provide a positive prompt or prompt embeddings or chunk_prompts.") - states = [] - try: - 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 + 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, ) - 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 = sample_noise( - self, current.shape, generator, current.device, 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 - ) + 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, - current.new_zeros(()), - start, - causal=True, + self, current, condition, first_frame, timestep, 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( + 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=True, + cache_update=False, attention_kwargs=attention_kwargs, return_dict=False, - ) + )[0] if self.do_classifier_free_guidance: - self.transformer( + 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=True, + 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 ) - generated.append(current) - progress_bar.update() - latents = torch.cat(generated, dim=2) - frames = self.decode_latents(latents, output_type) - finally: - for state in states: - state.reset() - self.maybe_free_model_hooks() + 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_utils.py b/src/diffusers/pipelines/edenlive/pipeline_utils.py index 3e70bd9547c0..8141f2856087 100644 --- a/src/diffusers/pipelines/edenlive/pipeline_utils.py +++ b/src/diffusers/pipelines/edenlive/pipeline_utils.py @@ -14,23 +14,6 @@ import torch -from ...utils.torch_utils import randn_tensor - - -def sample_noise(pipe, shape, generator, device, dtype): - parallel_config = getattr(pipe.transformer, "_parallel_config", None) - context_parallel_config = getattr(parallel_config, "context_parallel_config", None) - if context_parallel_config is None: - return randn_tensor(shape, generator=generator, device=device, dtype=dtype) - group = context_parallel_config._flattened_mesh.get_group() - source_rank = torch.distributed.get_global_rank(group, 0) - if torch.distributed.get_rank() == source_rank: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - noise = torch.empty(shape, device=device, dtype=dtype) - torch.distributed.broadcast(noise, src=source_rank, group=group) - return noise - def prepare_model_input(pipe, latents, condition, first_frame, timestep, start=0, causal=False): model_input = latents diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index f37686781430..ad94d5ea2d1f 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1354,21 +1354,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class EdenLiveProxyProjection(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"] diff --git a/tests/models/transformers/test_models_edenlive_proxy_projection.py b/tests/models/transformers/test_models_edenlive_proxy_projection.py deleted file mode 100644 index 9cf8d6bbf743..000000000000 --- a/tests/models/transformers/test_models_edenlive_proxy_projection.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 EdenLiveProxyProjection -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 EdenLiveProxyProjectionTesterConfig(BaseModelTesterConfig): - model_class = EdenLiveProxyProjection - main_input_name = "hidden_states" - input_shape = (4, 4, 16) - output_shape = (4, 4, 8) - - def get_init_dict(self): - return {"in_channels": 16, "out_channels": 8} - - def get_dummy_inputs(self, height=4, width=4): - return { - "hidden_states": randn_tensor( - (1, height, width, 16), - generator=torch.Generator().manual_seed(0), - device=torch_device, - dtype=self.torch_dtype, - ) - } - - -class TestEdenLiveProxyProjectionModel(EdenLiveProxyProjectionTesterConfig, ModelTesterMixin): - pass - - -class TestEdenLiveProxyProjectionMemory(EdenLiveProxyProjectionTesterConfig, 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] - - @torch.no_grad() - def test_cpu_offload(self, base_model_output): - from accelerate import cpu_offload - - torch.manual_seed(0) - model = self.model_class(**self.get_init_dict()).eval() - cpu_offload(model, execution_device=torch.device(torch_device)) - actual = model(**self.get_dummy_inputs()).sample - torch.testing.assert_close(actual, base_model_output, atol=1e-5, rtol=0) - - @pytest.mark.skip(reason="One linear layer cannot be partitioned between an execution device and disk.") - def test_disk_offload_without_safetensors(self): - pass - - @pytest.mark.skip(reason="One linear layer cannot be partitioned between an execution device and disk.") - def test_disk_offload_with_safetensors(self): - pass - - -class TestEdenLiveProxyProjectionTorchCompile(EdenLiveProxyProjectionTesterConfig, 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 index 26d0eb2c6537..d28defc9ebc1 100644 --- a/tests/models/transformers/test_models_transformer_edenlive.py +++ b/tests/models/transformers/test_models_transformer_edenlive.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import EdenLiveTransformer3DModel @@ -21,7 +22,6 @@ from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, - ContextParallelTesterMixin, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, @@ -54,6 +54,7 @@ def get_init_dict(self): "ffn_dim": 32, "num_layers": 2, "proxy_channels": 8, + "proxy_input_channels": 16, "rope_max_seq_len": 64, } @@ -64,7 +65,7 @@ def get_dummy_inputs(self, batch_size=1, height=8, width=8): (batch_size, 4, 3, height, width), generator=generator, device=torch_device, dtype=self.torch_dtype ), "proxy_features": randn_tensor( - (batch_size, 8, 3, 4, 4), generator=generator, device=torch_device, dtype=self.torch_dtype + (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 @@ -74,10 +75,26 @@ def get_dummy_inputs(self, batch_size=1, height=8, width=8): 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, @@ -97,12 +114,12 @@ def test_causal_state_commits_only_clean_chunks_and_preserves_sink(self): inputs["timestep"].zero_() model(**inputs, causal_state=state, cache_update=True) assert state.next_frame == 3 * (index + 1) - assert all(layer.num_tokens <= 18 * state.tokens_per_frame for layer in state.layers) + assert all(layer.num_tokens <= 18 * tokens_per_frame for layer in state.layers) if index == 1: - sink_keys = [layer.key[:, : 6 * state.tokens_per_frame].clone() for layer in state.layers] + 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 * state.tokens_per_frame], expected, atol=0, rtol=0) + 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()) @@ -128,10 +145,6 @@ def test_torch_compile_without_proxy_resize(self): torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) -class TestEdenLiveTransformerContextParallel(EdenLiveTransformerTesterConfig, ContextParallelTesterMixin): - pass - - class TestEdenLiveTransformerTraining(EdenLiveTransformerTesterConfig, TrainingTesterMixin): pass diff --git a/tests/models/transformers/test_models_transformer_edenlive_causal.py b/tests/models/transformers/test_models_transformer_edenlive_causal.py index 77b8dbe09f68..9197fc187854 100644 --- a/tests/models/transformers/test_models_transformer_edenlive_causal.py +++ b/tests/models/transformers/test_models_transformer_edenlive_causal.py @@ -12,23 +12,45 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import socket -from datetime import timedelta - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from diffusers import ContextParallelConfig, EdenLiveTransformer3DModel from diffusers.models.transformers import transformer_edenlive -from ...testing_utils import enable_full_determinism +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) @@ -82,143 +104,64 @@ def call(ids, expected_ids, commit): call([16], [0, 1, 9, 12, 13, 16], True) -def _parallel_outputs(model, inputs): - with torch.no_grad(): - output = model(**inputs, return_dict=False)[0] - state = model.create_causal_state( - batch_size=1, - height=8, - width=8, - chunk_size=3, - max_context_frames=18, - sink_frames=6, - device=inputs["hidden_states"].device, - dtype=inputs["hidden_states"].dtype, +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, + } ) - clean_inputs = {**inputs, "timestep": torch.zeros_like(inputs["timestep"])} - sink_keys = None - for index in range(8): - model(**clean_inputs, causal_state=state, cache_update=True) - assert state.next_frame == 3 * (index + 1) - assert all(layer.num_tokens <= 18 * state.tokens_per_frame // state.cp_size for layer in state.layers) - sink_tokens = 6 * state.tokens_per_frame // state.cp_size - if index == 1: - sink_keys = [layer.key[:, :sink_tokens].clone() for layer in state.layers] - elif index > 1: - for expected, layer in zip(sink_keys, state.layers): - torch.testing.assert_close(layer.key[:, :sink_tokens], expected, rtol=0, atol=0) - tail_inputs = { - **inputs, - "hidden_states": inputs["hidden_states"][:, :, :1], - "proxy_features": inputs["proxy_features"][:, :, :1], - "timestep": inputs["timestep"][:, :16] if inputs["timestep"].ndim == 2 else inputs["timestep"], + 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, } - tail = model(**tail_inputs, causal_state=state, return_dict=False)[0] - assert state.next_frame == 24 - tail_clean = {**tail_inputs, "timestep": torch.zeros_like(tail_inputs["timestep"])} - model(**tail_clean, causal_state=state, cache_update=True) - assert state.next_frame == 25 - return [output.float().cpu(), tail.float().cpu()] - - -def _context_parallel_worker(rank, world_size, port, init_dict, state_dict, inputs, backend, use_fsdp, output_path): - # Spawned workers must establish the same settings as the single-device reference. - enable_full_determinism() - torch.set_float32_matmul_precision("highest") - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), RANK=str(rank), WORLD_SIZE=str(world_size)) - torch.cuda.set_device(rank) - device = torch.device("cuda", rank) - dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 - try: - dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=180)) - model = EdenLiveTransformer3DModel(**init_dict).eval() - model.load_state_dict(state_dict) - model.to(device=device, dtype=dtype) - model.set_attention_backend(backend) - inputs = { - name: tensor.to(device=device, dtype=torch.float32 if name == "timestep" else dtype) - for name, tensor in inputs.items() - } - - def check_timestep(_module, _args, kwargs): - assert kwargs["timestep"].dtype == torch.float32 - assert torch.isin(kwargs["timestep"], torch.tensor([0.0, 375.0], device=device)).all() - - model.register_forward_pre_hook(check_timestep, with_kwargs=True) - from torch.distributed.device_mesh import init_device_mesh - - mesh = init_device_mesh("cuda", (1, world_size), mesh_dim_names=("ring", "ulysses")) - model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world_size, mesh=mesh)) - if use_fsdp: - from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard - - policy = MixedPrecisionPolicy( - param_dtype=dtype, reduce_dtype=torch.float32, output_dtype=dtype, cast_forward_inputs=False - ) - for block in model.blocks: - fully_shard(block, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - fully_shard(model, mesh=mesh["ulysses"], mp_policy=policy, reshard_after_forward=True) - outputs = _parallel_outputs(model, inputs) - if rank == 0: - torch.save(outputs, output_path) - finally: - if dist.is_initialized(): - dist.destroy_process_group() - - -class TestEdenLiveContextParallel(EdenLiveTransformerTesterConfig): - @pytest.mark.context_parallel - @pytest.mark.attention - @pytest.mark.parametrize("world_size", [2, 4]) - @pytest.mark.parametrize( - "backend,use_fsdp", [("native", False), ("_native_flash", False), ("_native_flash", True)] - ) - @pytest.mark.parametrize("num_layers,token_time", [(2, False), (1, True)]) - def test_context_parallel_inference(self, tmp_path, world_size, backend, use_fsdp, num_layers, token_time): - from diffusers.models.attention_dispatch import _AttentionBackendRegistry - - if not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < world_size: - pytest.skip(f"This case requires {world_size} CUDA devices with NCCL.") - previous_backend, _ = _AttentionBackendRegistry.get_active_backend() - deterministic = torch.are_deterministic_algorithms_enabled() - warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - try: - enable_full_determinism() - config = {**self.get_init_dict(), "attention_head_dim": 16, "num_layers": num_layers} - model = EdenLiveTransformer3DModel(**config).eval() - state = {name: tensor.cpu().clone() for name, tensor in model.state_dict().items()} - inputs = {name: tensor.cpu() for name, tensor in self.get_dummy_inputs().items()} - # This value would be rounded to 376 by an indiscriminate BF16 input cast. - inputs["timestep"] = torch.tensor([375.0], dtype=torch.float32) - if token_time: - inputs["timestep"] = torch.full((1, 48), 375.0) - inputs["timestep"][:, :16] = 0 - dtype = torch.bfloat16 if backend == "_native_flash" else torch.float32 - model.to(device="cuda:0", dtype=dtype).set_attention_backend(backend) - reference_inputs = { - name: tensor.to(device="cuda:0", dtype=torch.float32 if name == "timestep" else dtype) - for name, tensor in inputs.items() - } - expected = _parallel_outputs(model, reference_inputs) - del model, reference_inputs - torch.cuda.empty_cache() - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - port = sock.getsockname()[1] - output_path = str(tmp_path / "context_parallel.pt") - mp.spawn( - _context_parallel_worker, - args=(world_size, port, config, state, inputs, backend, use_fsdp, output_path), - nprocs=world_size, - join=True, - ) - actual = torch.load(output_path, map_location="cpu", weights_only=True) - for reference, output in zip(expected, actual): - assert torch.isfinite(output).all() - # Keep the existing small-model CP backend tolerance. - torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2) - finally: - _AttentionBackendRegistry.set_active_backend(previous_backend) - torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + 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/test_pipeline_edenlive.py b/tests/pipelines/edenlive/test_pipeline_edenlive.py index 86111aad8fee..fb89a2b6b33f 100644 --- a/tests/pipelines/edenlive/test_pipeline_edenlive.py +++ b/tests/pipelines/edenlive/test_pipeline_edenlive.py @@ -32,7 +32,6 @@ EdenLiveFlowMatchEulerDiscreteScheduler, EdenLiveLightTAEDecoder, EdenLivePipeline, - EdenLiveProxyProjection, EdenLiveTransformer3DModel, UniPCMultistepScheduler, ) @@ -83,14 +82,13 @@ def get_dummy_components(self): ) ).eval() torch.manual_seed(0) - proxy_projection = EdenLiveProxyProjection(in_channels=32, out_channels=8) - 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, @@ -129,7 +127,6 @@ def get_dummy_components(self): "transformer": transformer, "scheduler": scheduler, "proxy_encoder": proxy_encoder, - "proxy_projection": proxy_projection, "image_encoder": image_encoder, "image_processor": image_processor, } diff --git a/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py b/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py index 6f8fe905f7d7..7ca96012dc7e 100644 --- a/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py +++ b/tests/pipelines/edenlive/test_pipeline_edenlive_causal.py @@ -19,13 +19,11 @@ from diffusers import ( EdenLiveCausalPipeline, EdenLiveFlowMatchEulerDiscreteScheduler, - EdenLivePipeline, - EdenLiveTransformer3DModel, - UniPCMultistepScheduler, ) 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, @@ -47,22 +45,13 @@ def get_dummy_components(self): class TestEdenLiveCausalPipeline( EdenLiveOptionalDecoderTesterMixin, EdenLiveCausalPipelineTesterConfig, PipelineTesterMixin ): - def test_load_causal_components_from_bidirectional_root(self, tmp_path): + 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.transformer.save_pretrained(tmp_path / "transformer_causal") - pipe.scheduler.save_pretrained(tmp_path / "scheduler_causal") - components = self.get_dummy_components() - components["scheduler"] = UniPCMultistepScheduler( - prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5 - ) - EdenLivePipeline(**components).save_pretrained(tmp_path) - restored = self.pipeline_class.from_pretrained( - tmp_path, - transformer=EdenLiveTransformer3DModel.from_pretrained(tmp_path / "transformer_causal"), - scheduler=EdenLiveFlowMatchEulerDiscreteScheduler.from_pretrained(tmp_path / "scheduler_causal"), - ).to(torch_device) + 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): @@ -157,7 +146,7 @@ def record_dit(module, args, kwargs): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_clean_commit_timestep_shape_and_cfg(self, dtype): fixture = self - pipe = fixture.get_pipeline().to(torch_device, dtype=dtype) + pipe = cast_pipeline_to_dtype(fixture.get_pipeline().to(torch_device), dtype) calls = [] def capture(module, args, kwargs): @@ -198,7 +187,7 @@ def inputs(): actual = pipe(**inputs(), chunk_prompts=["snow", "snow"]).frames torch.testing.assert_close(actual, expected, rtol=0, atol=0) - def test_failure_clears_history_and_next_call_is_independent(self): + def test_next_call_is_independent_after_failure(self): pipe = self.get_pipeline().to(torch_device) def inputs(): @@ -229,7 +218,6 @@ def fail_in_second_chunk(module, args, kwargs): finally: handle.remove() assert len(states) == 2 - assert all(state.next_frame == 0 and all(layer.num_tokens == 0 for layer in state.layers) for state in states) actual = pipe(**inputs()).frames torch.testing.assert_close(actual, expected, atol=0, rtol=0)