import torch
from tqdm import tqdm
from diffusers.utils import export_to_video
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState

repo_id = "krea/krea-realtime-video"
pipe = ModularPipeline.from_pretrained(repo_id, trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
)
for block in pipe.transformer.blocks:
    block.self_attn.fuse_projections()

num_blocks = 9

frames = []
state = PipelineState()
prompt = ["a cat sitting on a boat"]

generator = torch.Generator(device=pipe.device).manual_seed(42)
for block_idx in tqdm(range(num_blocks)):
    state = pipe(
        state,
        prompt=prompt,
        num_inference_steps=6,
        num_blocks=num_blocks,
        block_idx=block_idx,
        generator=generator,
    )
    frames.extend(state.values["videos"][0])

export_to_video(frames, "output.mp4", fps=24)

import torch
from tqdm import tqdm
from diffusers.utils import load_video, export_to_video
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState

repo_id = "krea/krea-realtime-video"
pipe = ModularPipeline.from_pretrained(repo_id, trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
)
for block in pipe.transformer.blocks:
    block.self_attn.fuse_projections()

num_blocks = 9
video = load_video("https://app-uploads.krea.ai/public/a8218957-1a80-43dc-81b2-da970b5f2221-video.mp4")

frames = []
prompt = ["A car racing down a snowy mountain"]

state = PipelineState()
generator = torch.Generator("cuda").manual_seed(42)
for block_idx in tqdm(range(num_blocks)):
    state = pipe(
        state,
        video=video,
        prompt=prompt,
        num_inference_steps=6,
        strength=0.3,
        block_idx=block_idx,
        generator=generator,
    )
    frames.extend(state.values["videos"][0])

export_to_video(frames, "output-v2v.mp4", fps=24)

Using the video_stream input will process video frames in as they arrive, while maintaining temporal consistency across chunks.

import torch
from collections import deque
from tqdm import tqdm
from diffusers.utils import load_video, export_to_video
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState

repo_id = "krea/krea-realtime-video"
pipe = ModularPipeline.from_pretrained(repo_id, trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
)
for block in pipe.transformer.blocks:
    block.self_attn.fuse_projections()

n_samples = 9
frame_sample_len = 12
video = load_video(
    "https://app-uploads.krea.ai/public/a8218957-1a80-43dc-81b2-da970b5f2221-video.mp4"
)

# Simulate streaming video input
frame_samples = [
    video[sample_start : sample_start + frame_sample_len]
    for sample_start in range(0, n_samples * frame_sample_len, frame_sample_len)
]

frames = []
state = PipelineState()
prompt = ["A car racing down a snowny mountain road"]

block_idx = 0
generator = torch.Generator("cpu").manual_seed(42)
for frame_sample in tqdm(frame_samples):
    state = pipe(
        state,
        video_stream=frame_sample,
        prompt=prompt,
        num_inference_steps=6,
        strength=0.3,
        block_idx=block_idx,
        generator=generator,
    )
    frames.extend(state.values["videos"][0])

    block_idx += 1

export_to_video(frames, "output-v2v-streaming.mp4", fps=24)

import torch
from collections import deque
from tqdm import tqdm
from diffusers.utils import export_to_video
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState

repo_id = "krea/krea-realtime-video"
pipe = ModularPipeline.from_pretrained(repo_id, trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
)
pipe.transformer.load_lora_adapter(
    "shauray/Origami_WanLora",
    prefix="diffusion_model",
    weight_name="origami_000000500.safetensors",
    adapter_name="origami",
)
for block in pipe.transformer.blocks:
    block.self_attn.fuse_projections()

num_blocks = 9

frames = []
state = PipelineState()
prompt = ["[origami] a cat sitting on a boat"]

generator = torch.Generator("cuda").manual_seed(42)
for block_idx in tqdm(range(num_blocks)):
    state = pipe(
        state,
        prompt=prompt,
        num_inference_steps=6,
        num_blocks=num_blocks,
        block_idx=block_idx,
        generator=generator,
    )
    frames.extend(state.values["videos"][0])

export_to_video(frames, "output.mp4", fps=24)

To optimize inference speed and memory usage on Hopper level GPUs (H100s), we recommend using torch.compile, Sageattention and FP8 quantization with torchao.

First let's set up our depedencies by enabling Sageattention via Hub kernels and installing the torchao and kernels packages.

export DIFFUSERS_ENABLE_HUB_KERNELS=true
pip install -U kernels torchao

Alternatively, you can use Flash Attention 3 via kernels by disabling Sageattention:

export DISABLE_SAGEATTENTION=1

Then we will iterate over the blocks of the transformer and apply quantization and torch.compile.

import torch
from collections import deque
from tqdm import tqdm
from diffusers.utils import export_to_video
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, quantize_

repo_id = "krea/krea-realtime-video"
pipe = ModularPipeline.from_pretrained(repo_id, trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True,
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
)

for block in pipe.transformer.blocks:
    block.self_attn.fuse_projections()

# Quantize just the transformer blocks
for block in pipe.transformer.blocks:
    quantize_(block, Float8DynamicActivationFloat8WeightConfig())

# Compile just the attention modules
for submod in pipe.transformer.modules():
    if submod.__class__.__name__ in ["CausalWanAttentionBlock"]:
        submod.compile(fullgraph=False)

num_blocks = 9

state = PipelineState()
prompt = ["a cat sitting on a boat"]

# Compile warmup
for block_idx in range(num_blocks):
    state = pipe(
        state,
        prompt=prompt,
        num_inference_steps=2,
        num_blocks=num_blocks,
        num_frames_per_block=num_frames_per_block,
        block_idx=block_idx,
        generator=torch.Generator("cuda").manual_seed(42),
    )

# Reset state
state = PipelineState()
generator = torch.Generator("cuda").manual_seed(42)
for block_idx in tqdm(range(num_blocks)):
    state = pipe(
        state,
        prompt=prompt,
        num_inference_steps=6,
        num_blocks=num_blocks,
        block_idx=block_idx,
        generator=generator,
    )
    frames.extend(state.values["videos"][0])

export_to_video(frames, "output.mp4", fps=24)