Spaces:
Running on L40S
Running on L40S
File size: 3,729 Bytes
9f818c5 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
from __future__ import annotations
import html
import json
import tempfile
from typing import Literal
import torch
import wandb
from einops import rearrange
from cosmos_framework.utils.config import JobConfig
from cosmos_framework.utils import callback, distributed
from cosmos_framework.utils.easy_io import easy_io
class VisualizationLoggingCallback(callback.WandBCallback):
def __init__(
self,
every_n: int = 1,
input_normalization: str | None = None,
job: JobConfig | None = None,
config: callback.Config | None = None,
trainer: callback.ImaginaireTrainer | None = None,
):
super().__init__(config, trainer)
self.every_n = every_n
self.input_normalization = input_normalization
self.job = job
def on_training_step_end(
self,
model,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int = 0,
) -> None:
if iteration % self.every_n == 0 and distributed.is_rank0() and wandb.run is not None:
# log images to wandb for rank0
self.log_videos(log_type="train", data=data_batch, output=output_batch, iteration=iteration)
def on_validation_step_end(
self,
model,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int = 0,
) -> None: # Collect the validation batch and aggregate the overall loss.
super().on_validation_step_end(model, data_batch, output_batch, loss, iteration)
if iteration % self.every_n == 0 and distributed.is_rank0() and wandb.run is not None:
# log images to wandb for rank0
self.log_videos(log_type="val", data=data_batch, output=output_batch, iteration=iteration)
@torch.no_grad()
def log_videos(
self,
log_type: Literal["train", "val"],
data: dict[str, torch.Tensor],
output: dict[str, torch.Tensor],
iteration: int = 0,
):
if "raw_image" in data:
video = data["raw_image"][0].cpu() # [3,T,H,W], range [0, 255], uint8
elif "raw_video" in data:
video = data["raw_video"][0].cpu() # [3,T,H,W], range [0, 255], uint8
video = video.permute(1, 0, 2, 3) # [T,3,H,W]
video = video.numpy()
wandb_video = rearrange(video, "t c h w -> t h w c") # [T,H,W,3]
# create temp file
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
temp_file_path = temp_file.name
easy_io.dump(wandb_video, temp_file_path, file_format="mp4", format="mp4", fps=10, quality=5)
wandb.log({f"{log_type}/video": wandb.Video(temp_file_path)}, step=iteration)
# Convert dialog string to JSON and format it for HTML display
try:
dialog_json = json.loads(data["dialog_str"][0], indent=4)
formatted_html = f"""
<pre style='font-size: 0.5em;'>{dialog_json}</pre>
"""
wandb.log({f"{log_type}/prompt": wandb.Html(formatted_html)}, step=iteration)
except Exception as e:
# Fallback to original format if JSON conversion fails
# Escape HTML tags in the dialog string to display them properly
escaped_dialog = html.escape(data["dialog_str"][0])
wandb.log(
{f"{log_type}/prompt": wandb.Html(f"<pre style='font-size: 0.5em;'>{escaped_dialog}</pre>")},
step=iteration,
)
|