Add Echo-Memory community pipeline for Wan 2.1 1.3B

examples/community/README.md

@@ -10,6 +10,7 @@ Please also check out our [Community Scripts](https://github.com/huggingface/dif
 
 | Example                                                                                                                               | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              | Code Example                                                                              | Colab                                                                                                                                                                                                              |                                                        Author |
 |:--------------------------------------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------:|
+|Echo-Memory Wan 2.1|[Echo-Memory](https://arxiv.org/abs/2606.09803) overlays the released `context_k1` DiT fine-tune onto official Wan 2.1 1.3B after remapping original DiffSynth/Wan keys to Diffusers transformer names.|[Echo-Memory Wan 2.1](#echo-memory-wan-21)|-|[Weiyang Jin](https://github.com/WayneJin0918)|
 |Spatiotemporal Skip Guidance (STG)|[Spatiotemporal Skip Guidance for Enhanced Video Diffusion Sampling](https://huggingface.co/papers/2411.18664) (CVPR 2025) enhances video diffusion models by generating a weaker model through layer skipping and using it as guidance, improving fidelity in models like HunyuanVideo, LTXVideo, and Mochi.|[Spatiotemporal Skip Guidance](#spatiotemporal-skip-guidance)|-|[Junha Hyung](https://junhahyung.github.io/), [Kinam Kim](https://kinam0252.github.io/), and [Ednaordinary](https://github.com/Ednaordinary)|
 |Adaptive Mask Inpainting|Adaptive Mask Inpainting algorithm from [Beyond the Contact: Discovering Comprehensive Affordance for 3D Objects from Pre-trained 2D Diffusion Models](https://github.com/snuvclab/coma) (ECCV '24, Oral) provides a way to insert human inside the scene image without altering the background, by inpainting with adapting mask.|[Adaptive Mask Inpainting](#adaptive-mask-inpainting)|-|[Hyeonwoo Kim](https://sshowbiz.xyz),[Sookwan Han](https://jellyheadandrew.github.io)|
 |Flux with CFG|[Flux with CFG](https://github.com/ToTheBeginning/PuLID/blob/main/docs/pulid_for_flux.md) provides an implementation of using CFG in [Flux](https://blackforestlabs.ai/announcing-black-forest-labs/).|[Flux with CFG](#flux-with-cfg)|[Notebook](https://github.com/huggingface/notebooks/blob/main/diffusers/flux_with_cfg.ipynb)|[Linoy Tsaban](https://github.com/linoytsaban), [Apolinário](https://github.com/apolinario), and [Sayak Paul](https://github.com/sayakpaul)|
@@ -98,6 +99,44 @@ pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion
 
 ## Example usages
 
+### Echo-Memory Wan 2.1
+
+**Weiyang Jin and Echo Team**
+
+[Echo-Memory](https://arxiv.org/abs/2606.09803) is a controlled study of memory in action world models on Wan 2.1 1.3B. This community pipeline loads official [`Wan-AI/Wan2.1-T2V-1.3B-Diffusers`](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) and overlays the released [`context_k1`](https://huggingface.co/Echo-Team/Echo-Memory/tree/main/context_k1) row after remapping original DiffSynth / Wan keys to Diffusers transformer names (825 / 825 official keys). Extra action-MLP / SSM slots stay in the [Echo-Memory](https://github.com/Echo-Team-Joy-Future-Academy-JD/Echo-Memory) research stack.
+
+Already-converted transformer weights: [`Wayne-King/echo-memory-diffusers`](https://huggingface.co/Wayne-King/echo-memory-diffusers) (`context_k1-diffusers/diffusion_pytorch_model.safetensors`).
+
+#### Usage example
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.utils import export_to_video
+
+pipe = DiffusionPipeline.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    custom_pipeline="pipeline_echo_memory",
+    torch_dtype=torch.bfloat16,
+)
+pipe.load_echo_memory_weights()
+# or: pipe.load_converted_echo_memory_weights()
+pipe.to("cuda")
+
+frames = pipe(
+    prompt="A golden retriever running across a sunny green field, cinematic camera follow.",
+    negative_prompt="blurry, static, low quality, deformed",
+    height=480,
+    width=832,
+    num_frames=33,
+    num_inference_steps=30,
+    guidance_scale=5.0,
+).frames[0]
+export_to_video(frames, "echo_memory_context_k1.mp4", fps=16)
+```
+
+The same pipeline class is also hosted at [`Wayne-King/echo-memory-diffusers`](https://huggingface.co/Wayne-King/echo-memory-diffusers) for `custom_pipeline="Wayne-King/echo-memory-diffusers"` with `trust_remote_code=True`.
+
 ### Spatiotemporal Skip Guidance
 
 **Junha Hyung\*, Kinam Kim\*, Susung Hong, Min-Jung Kim, Jaegul Choo**

examples/community/pipeline_echo_memory.py

@@ -0,0 +1,159 @@
+# Copyright 2026 Echo Team 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.
+
+"""Echo-Memory community pipeline for official Wan 2.1 Diffusers weights.
+
+Loads `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, then overlays the released
+`context_k1` row from `Echo-Team/Echo-Memory` after remapping original
+DiffSynth / Wan keys onto the Diffusers transformer.
+
+Paper: https://arxiv.org/abs/2606.09803
+Code: https://github.com/Echo-Team-Joy-Future-Academy-JD/Echo-Memory
+"""
+
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from diffusers import WanPipeline
+
+
+DEFAULT_BASE_MODEL = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DEFAULT_REPO_ID = "Echo-Team/Echo-Memory"
+DEFAULT_FILENAME = "context_k1/epoch-0.safetensors"
+DEFAULT_CONVERTED_REPO_ID = "Wayne-King/echo-memory-diffusers"
+DEFAULT_CONVERTED_FILENAME = "context_k1-diffusers/diffusion_pytorch_model.safetensors"
+
+SKIP_SUBSTRINGS = (
+    "action_mlp",
+    "self_attn_with_action",
+    "block_wise_ssm",
+    "videossm_hybrid",
+    "spatial_memory_module",
+)
+
+# Same mapping as `scripts/convert_wan_to_diffusers.py` for Wan 2.1 T2V.
+TRANSFORMER_KEYS_RENAME_DICT = {
+    "time_embedding.0": "condition_embedder.time_embedder.linear_1",
+    "time_embedding.2": "condition_embedder.time_embedder.linear_2",
+    "text_embedding.0": "condition_embedder.text_embedder.linear_1",
+    "text_embedding.2": "condition_embedder.text_embedder.linear_2",
+    "time_projection.1": "condition_embedder.time_proj",
+    "head.modulation": "scale_shift_table",
+    "head.head": "proj_out",
+    "modulation": "scale_shift_table",
+    "ffn.0": "ffn.net.0.proj",
+    "ffn.2": "ffn.net.2",
+    # The original model names norms as norm1, norm3, norm2.
+    # Diffusers uses norm1, norm2, norm3.
+    "norm2": "norm__placeholder",
+    "norm3": "norm2",
+    "norm__placeholder": "norm3",
+    "self_attn.q": "attn1.to_q",
+    "self_attn.k": "attn1.to_k",
+    "self_attn.v": "attn1.to_v",
+    "self_attn.o": "attn1.to_out.0",
+    "self_attn.norm_q": "attn1.norm_q",
+    "self_attn.norm_k": "attn1.norm_k",
+    "cross_attn.q": "attn2.to_q",
+    "cross_attn.k": "attn2.to_k",
+    "cross_attn.v": "attn2.to_v",
+    "cross_attn.o": "attn2.to_out.0",
+    "cross_attn.norm_q": "attn2.norm_q",
+    "cross_attn.norm_k": "attn2.norm_k",
+}
+
+
+def is_diffusers_transformer_state_dict(keys: Iterable[str]) -> bool:
+    keys = list(keys)
+    return any(key.startswith("condition_embedder.") or ".attn1." in key for key in keys)
+
+
+def convert_echo_memory_transformer_state_dict(
+    state_dict: Dict[str, torch.Tensor],
+    skip_substrings: Iterable[str] = SKIP_SUBSTRINGS,
+) -> Tuple[Dict[str, torch.Tensor], List[str]]:
+    """Convert original Echo-Memory / DiffSynth Wan keys to Diffusers names."""
+    skip_substrings = tuple(skip_substrings)
+    if is_diffusers_transformer_state_dict(state_dict):
+        converted = {
+            key: value
+            for key, value in state_dict.items()
+            if not any(token in key for token in skip_substrings)
+        }
+        skipped = [key for key in state_dict if key not in converted]
+        return converted, skipped
+
+    converted = {}
+    skipped = []
+    for key, value in state_dict.items():
+        if any(token in key for token in skip_substrings):
+            skipped.append(key)
+            continue
+        new_key = key
+        for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
+            new_key = new_key.replace(replace_key, rename_key)
+        converted[new_key] = value
+    return converted, skipped
+
+
+class EchoMemoryPipeline(WanPipeline):
+    """Wan 2.1 T2V pipeline with an Echo-Memory `context_k1` overlay."""
+
+    def load_echo_memory_weights(
+        self,
+        repo_id: str = DEFAULT_REPO_ID,
+        filename: str = DEFAULT_FILENAME,
+        local_path: Optional[str] = None,
+        strict: bool = False,
+    ):
+        """Download one Echo-Memory row and overlay it on `self.transformer`."""
+        ckpt_path = local_path or hf_hub_download(repo_id=repo_id, filename=filename)
+        raw = load_file(ckpt_path)
+        converted, skipped = convert_echo_memory_transformer_state_dict(raw)
+        missing, unexpected = self.transformer.load_state_dict(converted, strict=strict)
+        print(
+            f"[Echo-Memory] overlaid {len(converted)}/{len(raw)} transformer keys from {ckpt_path} "
+            f"(skipped={len(skipped)}, missing={len(missing)}, unexpected={len(unexpected)})"
+        )
+        return missing, unexpected, skipped
+
+    def load_converted_echo_memory_weights(
+        self,
+        repo_id: str = DEFAULT_CONVERTED_REPO_ID,
+        filename: str = DEFAULT_CONVERTED_FILENAME,
+        local_path: Optional[str] = None,
+        strict: bool = False,
+    ):
+        """Overlay the already-remapped `context_k1` transformer weights."""
+        return self.load_echo_memory_weights(
+            repo_id=repo_id,
+            filename=filename,
+            local_path=local_path,
+            strict=strict,
+        )
+
+    @classmethod
+    def from_echo_memory(
+        cls,
+        pretrained_model_name_or_path: str = DEFAULT_BASE_MODEL,
+        echo_memory_repo: str = DEFAULT_REPO_ID,
+        echo_memory_filename: str = DEFAULT_FILENAME,
+        **kwargs,
+    ):
+        pipe = cls.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        pipe.load_echo_memory_weights(repo_id=echo_memory_repo, filename=echo_memory_filename)
+        return pipe