import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("omnimem/long_rolling_switch_prompt_from_ode_init", dtype=torch.bfloat16, device_map="cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]long_rolling_switch_prompt_from_ode_init
OmniMem long rolling video generation with prompt switching, initialized from an ODE-trained lineage.
This release contains the non-EMA generator checkpoint at training step 1008.
Contents
config.json- Two
safetensorsweight shards and their index SHA256SUMSfor payload integrity checking
The weights retain their original FP32 storage. Inference was evaluated in BF16. The configuration uses NSA top-12 block selection.
Usage
Use compatible OmniMem code with CausalWanNSAModel and CausalWanT2VPipeline.
The Wan2.1 VAE, UMT5 text encoder, and tokenizer are separate required components.
from omnimem.models.transformers.causal_wan_nsa_model import CausalWanNSAModel
model = CausalWanNSAModel.from_pretrained(
"omnimem/long_rolling_switch_prompt_from_ode_init"
)
The original checkpoint configuration is preserved unchanged. Its saved class name
records the training wrapper, and _name_or_path records the initialization source.
This release is the step-1008 non-EMA checkpoint.
Evaluated inference settings
| Setting | Value |
|---|---|
| Resolution | 832 × 480 |
| Video duration | 60 seconds |
| Frame rate | 16 FPS |
| Exported frames | 960 (969 generated, tail cropped) |
| Sampling steps | 4 |
| Guidance scale | 1 |
| Timestep shift | 5 |
| Seed | 42 |
| Frame block size | 3 |
| Prompt switch latent indices | 42, 81, 120, 162, 201 |
| KV history window | 3 |
| KV cache LRU maximum | 7 |
| Prompt recache mode | full |
| VAE decode chunk size | 3 |
The checkpoint was evaluated on five prompt-switching sequences.
- Downloads last month
- 5