Inference requirements

[Nole21]

Thanks for sharing such wonderful work. I wonder what is the gpu memory requirement for running the inference pipeline of unirelight. I failed on 48G L40s GPU.

[luckyday2022]

+1

[raulmosa]

@Nole21 @luckyday2022 I ran into the same issue.

I noticed that the inference script loads the EMA tracker alongside the model. The EMA tracker is typically used during training to maintain averaged weights, but I think it gets initialized during inference as well.

I tried applying these two small changes and it seems to fix the issue:

cosmos_predict1/diffusion/inference/inference_utils.py — add an inference_mode flag to load_model_by_config:

def load_model_by_config(
    config_job_name,
    config_file="projects/cosmos_video/config/config.py",
    model_class=DiffusionT2WModel,
    inference_mode=False,           # ← add this
):
    config = override(config, ["--", f"experiment={config_job_name}"])

    # Disable EMA during inference to avoid duplicating model weights on GPU
    if inference_mode and hasattr(config.model, "ema"):
        config.model.ema.enabled = False

    config.validate()
    config.freeze()

cosmos_predict1/diffusion/inference/single_gpu_relight.py — pass inference_mode=True when loading the model:

model, config = load_model_by_config(
    config_job_name=args.config,
    config_file=args.config_file,
    model_class=MultiviewExtendDiffusionModelRelight,
    inference_mode=True,            # ← add this
)

With these changes, inference ran fine on a L40S.

Anyway the authors can confirm whether this is the right approach.