import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("omnimem/self_forcing_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]self_forcing_from_ode_init
OmniMem ordinary self-forcing checkpoint initialized from ODE-trained generator weights.
This release contains the generator EMA weights at self-forcing training step 800. It is the shared initialization checkpoint used for the plain and prompt-switching rolling-training branches.
Checkpoint lineage
ODE step-800 generator weights → ordinary self-forcing step-800 EMA → rolling-training branches.
The downstream releases are:
Contents
config.json- Two
safetensorsweight shards and their index SHA256SUMSfor payload integrity checking
The weights retain their original FP32 storage. The model 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/self_forcing_from_ode_init"
)
The original checkpoint configuration is preserved unchanged. Its saved class name
records the training wrapper, and _name_or_path records the upstream ODE initialization source.
The released payload is the ordinary self-forcing step-800 EMA checkpoint.
- Downloads last month
- 6