Initial release: Memento LoRA + KeyframeQuery weights

README.md

@@ -1,3 +1,161 @@
 ---
 license: apache-2.0
+tags:
+  - video-generation
+  - wan
+  - lora
+  - memory-to-video
+  - character-consistency
+  - long-video
+base_model:
+  - Wan-AI/Wan2.2-T2V-A14B
+  - Wan-AI/Wan2.2-I2V-A14B
 ---
+
+# Memento - LoRA Weights for Consistent Long Video Generation
+
+Memento enables character-consistent multi-shot long video generation by extending Wan2.2-A14B with a learnable memory mechanism and identity reconstruction.
+
+## Model Details
+
+| Property | Value |
+|----------|-------|
+| Base model | Wan2.2-A14B (T2V + I2V) |
+| Architecture | Dual DiT (low noise + high noise) with flow matching, boundary at t=0.9 |
+| LoRA rank | 128 |
+| LoRA targets | Self-attention (Q/K/V/O) + Cross-attention (Q/K/V/O) + FFN layers |
+| Additional modules | KeyframeQuery (learnable memory selection, 12 queries) |
+| Parameters | ~3.2 GB per noise model (LoRA + KeyframeQuery) |
+| Precision | bfloat16 |
+| Resolution | 832x480, 81 frames per shot (~5s at 16fps) |
+
+## Key Features
+
+- **Memory-to-Video (M2V)**: Maintains a latent memory pool across shots, enabling character consistency across an entire story
+- **Learnable KeyframeQuery**: Selects the most relevant memory frames per shot using split local+global queries
+- **Identity Reconstruction**: Uses character appearance descriptions to condition identity preservation
+- **Split Identity Attention**: Separates identity and memory attention paths for better consistency
+- **Dual Query Mode**: Low noise and high noise models use independent keyframe queries optimized for their respective denoising stages
+
+## Files
+
+```
+backbone_low_noise.pth   # Low noise model: LoRA weights + KeyframeQuery (t < 0.9)
+backbone_high_noise.pth  # High noise model: LoRA weights + KeyframeQuery (t >= 0.9)
+```
+
+Each file contains:
+- `trainable_state_dict`: 800 LoRA parameter tensors + 24 KeyframeQuery parameter tensors
+- `config`: training configuration metadata
+
+## Usage
+
+### Requirements
+
+- 8x NVIDIA GPUs (A100 80GB recommended)
+- PyTorch 2.0+
+- Base models: [Wan2.2-T2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B) and [Wan2.2-I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B)
+- Dependencies: `transformers`, `peft`, `json5`, `av`, `decord`, `diffusers`
+
+### Directory Setup
+
+```
+memento_src/
+├── models/
+│   ├── Wan2.2-T2V-A14B/           # Base T2V model
+│   ├── Wan2.2-I2V-A14B/           # Base I2V model
+│   └── memento_lora/              # This repo's weights
+│       ├── backbone_low_noise.pth
+│       └── backbone_high_noise.pth
+└── ...
+```
+
+### Single Story Inference
+
+```bash
+cd memento_src
+
+torchrun --nproc_per_node=8 --master_port=8200 \
+  pipeline_learnable_acce.py \
+  --story_script_path ./story_rewritten_aligned/showcase/astronaut.json \
+  --t2v_model_path ./models/Wan2.2-T2V-A14B \
+  --i2v_model_path ./models/Wan2.2-I2V-A14B \
+  --lora_weight_path ./models/memento_lora \
+  --size "832*480" \
+  --max_memory_size 8 \
+  --max_memory_frames 8 \
+  --output_dir ./results/astronaut \
+  --ulysses_size 8 \
+  --offload_model \
+  --lora_rank 128 \
+  --mi2v \
+  --t2v_first_shot \
+  --t5_fsdp \
+  --dit_fsdp \
+  --split_identity_attn \
+  --split_learnable_query \
+  --global_query_num 2 \
+  --use_subject_recon \
+  --use_both_query \
+  --compile_dit
+```
+
+### Batch Inference
+
+```bash
+bash run_inference.sh ./story_rewritten_aligned/showcase ./results ./models/memento_lora
+```
+
+### Key Inference Flags
+
+| Flag | Description |
+|------|-------------|
+| `--t2v_first_shot` | Generate first shot with T2V model (no memory conditioning) |
+| `--mi2v` | Use last frame of previous shot as I2V condition for continuous shots |
+| `--split_identity_attn` | Enable split identity+memory attention mechanism |
+| `--split_learnable_query` | Use separate local and global queries for memory selection |
+| `--global_query_num 2` | Number of global (appearance-focused) memory queries |
+| `--use_subject_recon` | Enable identity reconstruction conditioning |
+| `--use_both_query` | Each noise model uses its own KeyframeQuery for memory selection |
+| `--compile_dit` | Apply torch.compile for ~15% speedup (slower first step) |
+| `--max_memory_size 8` | Maximum memory frames fed to DiT |
+| `--max_memory_frames 8` | Frames sampled from previous shot for memory update |
+
+### Story Script Format
+
+```json
+{
+  "reconstruct_target": "[Person A] physical appearance description\n[Person B] physical appearance description",
+  "scenes": [
+    {
+      "scene_num": 1,
+      "cut": [true, false, false],
+      "video_prompts": [
+        "global caption: [Person A]: appearance; shot caption: [Person A] action in present tense...",
+        "global caption: [Person A]: appearance; shot caption: [Person A] continues action...",
+        "global caption: [Person A]: appearance; shot caption: [Person A] final action..."
+      ]
+    }
+  ]
+}
+```
+
+- `reconstruct_target`: Character identity descriptions (appearance only, no actions). One per line.
+- `cut[i]`: `true` = new camera angle (no I2V conditioning), `false` = continuous from previous shot
+- `global caption`: Lists appearances of characters **visible in this shot only**
+- `shot caption`: 1-3 sentences describing the visual action in present tense
+
+## How It Works
+
+1. **Shot 1**: Generated purely by T2V model using the first video_prompt
+2. **Subsequent shots**: Memory pool accumulates keyframes from all previous shots. The KeyframeQuery module selects the most relevant 8 frames based on the current shot's prompt. These memory frames + identity description condition the generation to maintain character consistency.
+3. **Continuous shots** (`cut=false`): Additionally use the last frame of the previous shot as an I2V starting frame for smooth visual transitions.
+
+## Citation
+
+```bibtex
+@article{memento2025,
+  title={Memento: Reconstruct to Remember for Consistent Long Video Generation},
+  year={2025}
+}
+```

backbone_high_noise.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abeea090c9e01c7b5292b7ec5a83e91655e3776e985b6cc49a7c9bed59759168
+size 1647301008

backbone_low_noise.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d863230da54c34758b77a251982c047f355bc1c8e7eb6bc468209635219a2d72
+size 1647300184

config.json

@@ -0,0 +1,72 @@
+{
+  "model_type": "memento",
+  "model_name": "Memento",
+  "version": "1.0",
+  "license": "apache-2.0",
+  "description": "Memento LoRA + KeyframeQuery weights for character-consistent multi-shot long video generation, built on Wan2.2-A14B.",
+  "base_model": {
+    "t2v": "Wan-AI/Wan2.2-T2V-A14B",
+    "i2v": "Wan-AI/Wan2.2-I2V-A14B",
+    "architecture": "Dual DiT (low-noise + high-noise) with flow matching",
+    "boundary": 0.9
+  },
+  "lora": {
+    "rank": 128,
+    "alpha": 128,
+    "dropout": 0.0,
+    "hidden_dim": 5120,
+    "target_modules": [
+      "self_attn.q",
+      "self_attn.k",
+      "self_attn.v",
+      "self_attn.o",
+      "cross_attn.q",
+      "cross_attn.k",
+      "cross_attn.v",
+      "cross_attn.o",
+      "ffn.0",
+      "ffn.2"
+    ]
+  },
+  "keyframe_query": {
+    "num_local_queries": 10,
+    "num_global_queries": 2,
+    "num_keyframes": 10,
+    "query_dim": 5120,
+    "memory_patch_in_channels": 16,
+    "memory_patch_size": [1, 2, 2]
+  },
+  "precision": "bfloat16",
+  "video": {
+    "resolution": "832x480",
+    "frames_per_shot": 81,
+    "fps": 16
+  },
+  "files": {
+    "low_noise_weights": "backbone_low_noise.safetensors",
+    "high_noise_weights": "backbone_high_noise.safetensors"
+  },
+  "inference": {
+    "num_gpus": 8,
+    "ulysses_size": 8,
+    "lora_rank": 128,
+    "size": "832*480",
+    "max_memory_size": 8,
+    "max_memory_frames": 8,
+    "global_query_num": 2,
+    "sample_guide_scale": 3.5,
+    "sample_solver": "unipc",
+    "flags": {
+      "mi2v": true,
+      "t2v_first_shot": true,
+      "t5_fsdp": true,
+      "dit_fsdp": true,
+      "split_identity_attn": true,
+      "split_learnable_query": true,
+      "use_subject_recon": true,
+      "use_both_query": true,
+      "offload_model": true,
+      "compile_dit": true
+    }
+  }
+}