UniRelight (Blackwell)

Docker image + web UI for Lilgeneric/UniRelight on NVIDIA Blackwell (SM_120).

Base	nvidia/cuda:12.9.0-devel-ubuntu24.04
Python	3.10
CUDA / Torch	12.9 / 2.12.1+cu129
UI	http://localhost:8080

This stack does not recreate the upstream CUDA 12.4 / Torch 2.6 conda env. See DEPENDENCY_DEVIATIONS.md.

---

Start from scratch

For a new machine with an NVIDIA GPU (Blackwell / SM_120 recommended). Commands below are PowerShell; on Linux/macOS use the same docker / docker compose lines (and ./build.sh instead of .\build.ps1).

Prerequisites

1. NVIDIA driver recent enough for CUDA 12.9
2. Docker Desktop (Windows/macOS) or Docker Engine + Compose plugin (Linux)
3. NVIDIA Container Toolkit so containers see the GPU (docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu24.04 nvidia-smi should work)
4. Disk space: ~40+ GB for the image, plus ~30 GB for model weights
5. Hugging Face access for checkpoint download (login if the repo requires it)

1. Get this project

cd E:\Dockerbuilds\uniRelight   # or wherever you cloned / copied this repo

2. Create host folders + HF cache volume

mkdir -Force data\input, data\hdri, data\outputs, data\checkpoints, data\models, data\jobs, data\frames
docker volume create unirelight-hf-cache

3. Build the runtime image (once, slow)

This builds unirelight-blackwell:cu129 (CUDA, Python 3.10, PyTorch cu129, TE, etc.). Expect a long first build (often an hour+).

.\build.ps1 -MaxJobs 8

Linux:

./build.sh

4. Build the web UI overlay (fast)

Compose layers the UI/entrypoint on top of that runtime image → unirelight-blackwell:cu129-web.

docker compose build

5. Download model weights (once, large)

Weights are not inside the image; they land under ./data/checkpoints on the host (~27 GB+).

docker compose run --rm unirelight bash -lc "python scripts/download_unirelight_checkpoints.py --checkpoint_dir /data/checkpoints"

Confirm:

data/checkpoints/UniRelight/model.pt
data/checkpoints/Cosmos-Tokenize1-CV8x8x8-720p/mean_std.pt

6. Start the UI

docker compose up -d

Open http://localhost:8080

• Status should become healthy; /healthz returns "ok": true and "ckpt_ready": true
• Empty data/input / data/hdri may be seeded with example assets on first start
• Or upload your own video / frame folder and HDRI in the UI

7. Run a relight

In the UI:

1. Select a video or frame folder
2. Select an HDRI
3. Pick resolution (default 848 × 480)
4. Leave outputs as Relit + Albedo (or enable PNG sequence / rotate light)
5. Press Execute

While the job runs, the right panel shows live logs:

When it finishes you get separate albedo and relit previews (and matching files under data/outputs/<job-id>/):

8. Stop (data stays on disk)

docker compose down

./data and the HF volume are kept. Next time: docker compose up -d.

Day-to-day after the first setup

Change	Command
Start UI again	docker compose up -d
UI / entrypoint / patches only	docker compose build && docker compose up -d
Full Torch/CUDA stack change	.\build.ps1 then docker compose build && docker compose up -d

---

Persistent data layout

Host folders under ./data are bind-mounted into the container. Deleting the container does not delete models or outputs.

data/
  input/          # videos (.mp4) or frame folders (…/00000.rgb.png)
  hdri/           # environment maps (.hdr / .exr)
  outputs/        # relight results (.mp4 and optional PNG sequences)
  checkpoints/    # UniRelight + Cosmos tokenizer weights
  models/         # optional extra models
  jobs/           # job logs (created by the UI)
  frames/         # extracted frames cache (created by the UI)

docker-compose.yml mounts each subfolder explicitly so host files are not shadowed by anonymous Docker volumes.

Hugging Face cache: named volume unirelight-hf-cache (see step 2 above).

Ports

Port	Service
8080	Web UI + API (/, /api/*, /healthz)

Change mapping in docker-compose.yml ("8080:8080") or set UNIRELIGHT_PORT.

CLI (optional)

docker compose run --rm unirelight bash

python -m cosmos_predict1.diffusion.inference.single_gpu_relight `
  --config_file cosmos_predict1/diffusion/training/config/config_relight.py `
  --config unirelight_cosmos_f57_480p `
  --ckpt_path /data/checkpoints/UniRelight/model.pt `
  --dataset_name /data/frames/<dataset> `
  --output_path /data/outputs/manual `
  --sample_n_frames 57 `
  --env_light_path /data/hdri/warm_bar_2k.hdr `
  --resolution 480 848 `
  --resize_resolution 486 864 `
  --save_condition false `
  --output_panels relit albedo

--output_panels accepts any of: relit, albedo, input, comparison, condition.
Default (patched image) is separate relit + albedo files — not a side-by-side.

GPU smoke test:

docker compose run --rm unirelight python /usr/local/bin/validate_runtime.py

What was fixed for Blackwell (baked into the image)

1. PyTorch cu129 instead of Torch 2.6 / CUDA 12.4
2. Transformer Engine 2.x shim — re-exports apply_rotary_pos_emb for Cosmos imports
3. nvdiffrast installed with --no-build-isolation (required for HDR projection)
4. ffmpegcv for Cosmos video helpers
5. Real example assets fetched from GitHub LFS media URLs (shallow clones only have pointer files)
6. --output_panels patch — export separate relit/albedo (upstream always wrote a side-by-side)

API sketch

• GET /api/assets — list videos / HDRIs
• POST /api/upload/video · POST /api/upload/hdri · POST /api/upload/frames
• POST /api/jobs — start relight (video_id, hdri_id, resolution, output_panels[], save_images, rotate_light, …)
• GET /api/jobs/{id} · GET /api/jobs/{id}/log
• GET /api/outputs/{id}/…