Cosmos-Predict2.5 (Video Generation)¶

Model Compilation for Multiple NPUs¶

Start by compiling the models in the Cosmos-Predict2.5 pipeline with the RBLNCosmos2_5_PredictBasePipeline class provided by optimum-rbln. To compile the models, you must set the export argument to True. The class's from_pretrained() method then downloads the original models from the HuggingFace Hub and compiles them using the RBLN Compiler. Since the diffusers-format weights of nvidia/Cosmos-Predict2.5-2B are provided on a separate revision of the HuggingFace Hub repository, pass revision="diffusers/base/post-trained". After compilation, save the model artifacts to disk using the save_pretrained() method. This creates a directory (e.g., rbln-cosmos-predict2.5) containing the compiled model.

Before compiling the models, specify the shape of the output video. Cosmos-Predict2.5 generates a 93-frame clip (about 5 seconds at 16 FPS) with a resolution of 704x1280 pixels. Conditioning media (an image or a video, when given) enters through the VAE encoder at runtime, so one compiled pipeline serves Text2World, Image2World, and Video2World modes.

In addition, several modules in the pipeline need multiple NPUs for inference, so specify the number of NPUs at compile time. The example code for compilation is as follows.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

import torch
from optimum.rbln import RBLNCosmos2_5_PredictBasePipeline

# Define the HuggingFace model ID
model_id = "nvidia/Cosmos-Predict2.5-2B"

# Compile Cosmos core models (transformer, text_encoder, vae, safety_checker)
pipe = RBLNCosmos2_5_PredictBasePipeline.from_pretrained(
    model_id,
    revision="diffusers/base/post-trained",
    torch_dtype=torch.bfloat16,
    export=True,
    rbln_config={
        "height": 704,
        "width": 1280,
        "num_frames": 93,
        "max_seq_len": 512,
        "create_runtimes": False,
        "text_encoder": {
            "num_devices": 8,
        },
        "transformer": {
            "num_devices": 8,
        },
    },
)

pipe.save_pretrained("rbln-cosmos-predict2.5")

Load the Compiled RBLN Model¶

To generate videos, load the compiled models saved in the previous step. You can load them with the from_pretrained() method of the RBLNCosmos2_5_PredictBasePipeline class. In this step, set the export argument to False to indicate that you are loading precompiled models from a local path.

The Cosmos-Predict2.5 pipeline consists of many modules, some of which occupy a large amount of device memory. Loading all of them on the same device causes a device out-of-memory error. Even with the same number of NPUs, loading can succeed or fail depending on which modules are placed on which devices, so set the device IDs in rbln_config carefully. In the example code below, every submodule is placed within the 8 NPUs that the transformer spans. text_encoder uses all eight, and the VAE encoder/decoder and the safety_checker are placed on the first three, so the whole pipeline fits in one pool of 8 NPUs.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

from optimum.rbln import RBLNCosmos2_5_PredictBasePipeline

pipe = RBLNCosmos2_5_PredictBasePipeline.from_pretrained(
    model_id="rbln-cosmos-predict2.5",
    export=False,
    rbln_config={
        "text_encoder": {
            "device": [0, 1, 2, 3, 4, 5, 6, 7],
        },
        "transformer": {
            "device": [0, 1, 2, 3, 4, 5, 6, 7],
        },
        "vae": {
            "device_map": {"encoder": 0, "decoder": 1},
        },
        "safety_checker": {
            "qwen3guard": {"device": 2},
            "face_blur_filter": {"device": 2},
        },
    },
)

Video Generation¶

After the compiled pipeline is loaded, generate videos as follows. Call the pipeline with your image and prompt to generate a video whose first frame follows the input image (Image2World). The negative prompt defaults to the value defined inside the pipeline, and you can pass your own prompt through the negative_prompt argument.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

import torch
from diffusers.utils import export_to_video, load_image

# Conditioning image: the first frame the generated world video continues from.
image = load_image(
    "https://media.githubusercontent.com/media/nvidia-cosmos/cosmos-predict2.5/refs/heads/main/assets/base/bus_terminal.jpg"
)
prompt = "A nighttime city bus terminal gradually shifts from stillness to subtle movement. At first, multiple double-decker buses are parked under the glow of overhead lights, with a central bus labeled '87D' facing forward and stationary. As the video progresses, the bus in the middle moves ahead slowly, its headlights brightening the surrounding area and casting reflections onto adjacent vehicles. The motion creates space in the lineup, signaling activity within the otherwise quiet station. It then comes to a smooth stop, resuming its position in line. Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."

# Image2World: condition on a single image and generate a 93-frame world video.
video = pipe(
    image=image,
    prompt=prompt,
    num_frames=93,
    generator=torch.Generator().manual_seed(1),
).frames[0]
export_to_video(video, "generated_video.mp4", fps=16)

Example Output:

The full example code is available in the RBLN Model Zoo GitHub repository.