modify log

aegis.py

@@ -15,7 +15,7 @@
 
 import argparse
 
-from . import log
+from .log import log
 import torch
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer

ar_model.py

@@ -19,7 +19,7 @@ import time
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set
 
-from . import log
+from .log import log
 import torch
 from safetensors.torch import load_file
 from torch.nn.modules.module import _IncompatibleKeys

ar_networks.py

@@ -20,7 +20,7 @@ from torch import nn
 
 from .ar_tokenizer_modules import CausalConv3d, DecoderFactorized, EncoderFactorized
 from .ar_tokenizer_quantizers import FSQuantizer
-from . import log
+from .log import log
 
 NetworkEval = namedtuple("NetworkEval", ["reconstructions", "quant_loss", "quant_info"])
 

ar_tokenizer_image_text_tokenizer.py

@@ -22,7 +22,7 @@ from transformers import AutoImageProcessor
 from transformers.image_utils import ImageInput, is_valid_image, load_image
 
 from .ar_tokenizer_text_tokenizer import TextTokenizer
-from . import log
+from .log import log
 
 # Configuration for different vision-language models
 IMAGE_CONFIGS = {

ar_tokenizer_modules.py

@@ -41,7 +41,7 @@ from .ar_tokenizer_utils import (
     space2batch,
     time2batch,
 )
-from . import log
+from .log import log
 
 
 class CausalConv3d(nn.Module):

ar_tokenizer_text_tokenizer.py

@@ -19,7 +19,7 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer
 
-from . import log
+from .log import log
 
 
 def get_tokenizer_path(model_family: str, is_instruct_model: bool = False):

ar_transformer.py

@@ -29,7 +29,7 @@ from .ar_modules_mlp import MLP
 from .ar_modules_normalization import create_norm
 from .checkpoint import process_state_dict, substrings_to_ignore
 from .ar_utils_misc import maybe_convert_to_namespace
-from . import log
+from .log import log
 
 
 class TransformerBlock(nn.Module):

blocklist.py

@@ -19,7 +19,7 @@ import re
 import string
 from difflib import SequenceMatcher
 
-from . import log
+from .log import log
 import nltk
 from better_profanity import profanity
 

blocks.py

@@ -23,7 +23,7 @@ from einops.layers.torch import Rearrange
 from torch import nn
 
 from .attention import Attention, GPT2FeedForward
-from . import log
+from .log import log
 
 
 def modulate(x, shift, scale):

conditioner.py

@@ -24,7 +24,7 @@ import torch
 import torch.nn as nn
 
 from .batch_ops import batch_mul
-from . import log
+from .log import log
 from .lazy_config_init import instantiate
 
 

config_helper.py

@@ -27,7 +27,7 @@ from hydra import compose, initialize
 from hydra.core.config_store import ConfigStore
 from omegaconf import DictConfig, OmegaConf
 
-from . import log
+from .log import log
 from .config import Config
 
 

cosmos1/models/autoregressive/diffusion_decoder/inference.py

@@ -22,7 +22,7 @@ import torch
 from inference_config import DiffusionDecoderSamplingConfig
 from cosmos1.models.autoregressive.diffusion_decoder.model import LatentDiffusionDecoderModel
 from cosmos1.models.autoregressive.diffusion_decoder.utils import linear_blend_video_list, split_with_overlap
-from . import log
+from .log import log
 
 
 def diffusion_decoder_process_tokens(

cosmos1/models/autoregressive/inference/base.py

@@ -21,7 +21,7 @@ import torch
 
 from cosmos1.models.autoregressive.inference.world_generation_pipeline import ARBaseGenerationPipeline
 from cosmos1.models.autoregressive.utils.inference import add_common_arguments, load_vision_input, validate_args
-from . import log
+from .log import log
 
 
 def parse_args():

cosmos1/models/autoregressive/inference/video2world.py

@@ -21,7 +21,7 @@ import torch
 
 from cosmos1.models.autoregressive.inference.world_generation_pipeline import ARVideo2WorldGenerationPipeline
 from cosmos1.models.autoregressive.utils.inference import add_common_arguments, load_vision_input, validate_args
-from . import log
+from .log import log
 from io import read_prompts_from_file
 
 

cosmos1/models/autoregressive/inference/world_generation_pipeline.py

@@ -17,7 +17,7 @@ import gc
 import os
 from typing import List, Optional, Tuple
 
-from . import log
+from .log import log
 import numpy as np
 import torch
 from einops import rearrange

cosmos1/models/autoregressive/nemo/cosmos.py

@@ -29,7 +29,7 @@ from nemo.lightning import OptimizerModule, io
 from nemo.lightning.base import teardown
 from torch import Tensor, nn
 
-from . import log
+from .log import log
 
 
 class RotaryEmbedding3D(RotaryEmbedding):

cosmos1/models/autoregressive/nemo/inference/general.py

@@ -37,7 +37,7 @@ from cosmos1.models.autoregressive.nemo.utils import run_diffusion_decoder_model
 from discrete_video import DiscreteVideoFSQJITTokenizer
 from cosmos1.models.autoregressive.utils.inference import load_vision_input
 from . import presets as guardrail_presets
-from . import log
+from .log import log
 
 torch._C._jit_set_texpr_fuser_enabled(False)
 

cosmos1/models/autoregressive/nemo/post_training/prepare_dataset.py

@@ -24,7 +24,7 @@ from nemo.collections.nlp.data.language_modeling.megatron import indexed_dataset
 
 from cosmos1.models.autoregressive.nemo.utils import read_input_videos
 from discrete_video import DiscreteVideoFSQJITTokenizer
-from . import log
+from .log import log
 
 TOKENIZER_COMPRESSION_FACTOR = [8, 16, 16]
 DATA_RESOLUTION_SUPPORTED = [640, 1024]

cosmos1/models/autoregressive/nemo/utils.py

@@ -31,7 +31,7 @@ from inference_utils import (
     load_tokenizer_model,
     skip_init_linear,
 )
-from . import log
+from .log import log
 from config_helper import get_config_module, override
 
 TOKENIZER_COMPRESSION_FACTOR = [8, 16, 16]

cosmos1/models/autoregressive/utils/inference.py

@@ -26,7 +26,7 @@ import torchvision
 from PIL import Image
 
 from inference_config import SamplingConfig
-from . import log
+from .log import log
 
 _IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", "webp"]
 _VIDEO_EXTENSIONS = [".mp4"]

cosmos1/models/diffusion/nemo/inference/general.py

@@ -37,7 +37,7 @@ from nemo.collections.diffusion.sampler.cosmos.cosmos_diffusion_pipeline import
 from transformers import T5EncoderModel, T5TokenizerFast
 
 from cosmos1.models.diffusion.nemo.inference.inference_utils import process_prompt, save_video
-from . import log
+from .log import log
 
 EXAMPLE_PROMPT = (
     "The teal robot is cooking food in a kitchen. Steam rises from a simmering pot "

cosmos1/models/diffusion/nemo/inference/inference_utils.py

@@ -30,7 +30,7 @@ from presets import (
     run_text_guardrail,
     run_video_guardrail,
 )
-from . import log
+from .log import log
 
 
 def get_upsampled_prompt(

cosmos1/models/diffusion/nemo/post_training/general.py

@@ -57,7 +57,7 @@ def cosmos_diffusion_7b_text2world_finetune() -> run.Partial:
     recipe.resume.resume_if_exists = False
 
     # Directory to save checkpoints / logs
-    recipe.log.log_dir = "nemo_experiments/cosmos_diffusion_7b_text2world_finetune"
+    recipe.log_log_dir = "nemo_experiments/cosmos_diffusion_7b_text2world_finetune"
 
     return recipe
 
@@ -102,7 +102,7 @@ def cosmos_diffusion_14b_text2world_finetune() -> run.Partial:
     recipe.resume.resume_if_exists = False
 
     # Directory to save checkpoints / logs
-    recipe.log.log_dir = "nemo_experiments/cosmos_diffusion_14b_text2world_finetune"
+    recipe.log_log_dir = "nemo_experiments/cosmos_diffusion_14b_text2world_finetune"
 
     return recipe
 

cosmos1/models/diffusion/nemo/post_training/prepare_dataset.py

@@ -27,7 +27,7 @@ from nemo.collections.diffusion.models.model import DiT7BConfig
 from tqdm import tqdm
 from transformers import T5EncoderModel, T5TokenizerFast
 
-from . import log
+from .log import log
 
 
 def get_parser():

cosmos1/models/diffusion/networks/general_dit.py

@@ -34,7 +34,7 @@ from blocks import (
     Timesteps,
 )
 from position_embedding import LearnablePosEmbAxis, VideoRopePosition3DEmb
-from . import log
+from .log import log
 
 
 class GeneralDIT(nn.Module):
@@ -390,16 +390,16 @@ class GeneralDIT(nn.Module):
             latent_condition_sigma=latent_condition_sigma,
         )
         # logging affline scale information
-        affline_scale_log_info = {}
+        affline_scale_log.info = {}
 
         timesteps_B_D, adaln_lora_B_3D = self.t_embedder(timesteps.flatten())
         affline_emb_B_D = timesteps_B_D
-        affline_scale_log_info["timesteps_B_D"] = timesteps_B_D.detach()
+        affline_scale_log.info["timesteps_B_D"] = timesteps_B_D.detach()
 
         if scalar_feature is not None:
             raise NotImplementedError("Scalar feature is not implemented yet.")
 
-        affline_scale_log_info["affline_emb_B_D"] = affline_emb_B_D.detach()
+        affline_scale_log.info["affline_emb_B_D"] = affline_emb_B_D.detach()
         affline_emb_B_D = self.affline_norm(affline_emb_B_D)
 
         if self.use_cross_attn_mask:

cosmos1/models/diffusion/networks/general_dit_video_conditioned.py

@@ -22,7 +22,7 @@ from torch import nn
 from conditioner import DataType
 from blocks import TimestepEmbedding, Timesteps
 from cosmos1.models.diffusion.networks.general_dit import GeneralDIT
-from . import log
+from .log import log
 
 
 class VideoExtendGeneralDIT(GeneralDIT):
@@ -155,11 +155,11 @@ class VideoExtendGeneralDIT(GeneralDIT):
             latent_condition_sigma=latent_condition_sigma,
         )
         # logging affline scale information
-        affline_scale_log_info = {}
+        affline_scale_log.info = {}
 
         timesteps_B_D, adaln_lora_B_3D = self.t_embedder(timesteps.flatten())
         affline_emb_B_D = timesteps_B_D
-        affline_scale_log_info["timesteps_B_D"] = timesteps_B_D.detach()
+        affline_scale_log.info["timesteps_B_D"] = timesteps_B_D.detach()
 
         if scalar_feature is not None:
             raise NotImplementedError("Scalar feature is not implemented yet.")
@@ -173,7 +173,7 @@ class VideoExtendGeneralDIT(GeneralDIT):
 
             affline_augment_sigma_emb_B_D, _ = self.augment_sigma_embedder(condition_video_augment_sigma.flatten())
             affline_emb_B_D = affline_emb_B_D + affline_augment_sigma_emb_B_D
-        affline_scale_log_info["affline_emb_B_D"] = affline_emb_B_D.detach()
+        affline_scale_log.info["affline_emb_B_D"] = affline_emb_B_D.detach()
         affline_emb_B_D = self.affline_norm(affline_emb_B_D)
 
         if self.use_cross_attn_mask:

distributed.py

@@ -27,7 +27,7 @@ import pynvml
 import torch
 import torch.distributed as dist
 
-from . import log
+from .log import log
 from .device import Device
 
 

face_blur_filter.py

@@ -16,7 +16,7 @@
 import argparse
 import os
 
-from . import log
+from .log import log
 import numpy as np
 import torch
 from pytorch_retinaface.data import cfg_re50

guardrail_blocklist_utils.py

@@ -16,7 +16,7 @@
 import os
 import re
 
-from . import log
+from .log import log
 
 
 def read_keyword_list_from_dir(folder_path: str) -> list[str]:

guardrail_core.py

@@ -17,7 +17,7 @@ from typing import Any, Tuple
 
 import numpy as np
 
-from . import log
+from .log import log
 
 
 class ContentSafetyGuardrail:

guardrail_io_utils.py

@@ -19,7 +19,7 @@ from dataclasses import dataclass
 import imageio
 import numpy as np
 
-from . import log
+from .log import log
 
 
 @dataclass

inference_utils.py

@@ -24,7 +24,7 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as transforms_F
 
-from. model_t2w import DiffusionT2WModel
+from .model_t2w import DiffusionT2WModel
 from .model_v2w import DiffusionV2WModel
 from .config_helper import get_config_module, override
 from .utils_io import load_from_fileobj

log.py

@@ -90,36 +90,39 @@ def _rank0_only_filter(record: Any) -> bool:
     return not is_rank0
 
 
-def trace(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).trace(message)
+class log():
 
+    @staticmethod
+    def trace(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).trace(message)
 
-def debug(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).debug(message)
+    @staticmethod
+    def debug(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).debug(message)
 
+    @staticmethod
+    def info(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).info(message)
 
-def info(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).info(message)
+    @staticmethod
+    def success(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).success(message)
 
+    @staticmethod
+    def warning(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).warning(message)
 
-def success(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).success(message)
+    @staticmethod
+    def error(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).error(message)
 
+    @staticmethod
+    def critical(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).critical(message)
 
-def warning(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).warning(message)
-
-
-def error(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).error(message)
-
-
-def critical(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).critical(message)
-
-
-def exception(message: str, rank0_only: bool = True) -> None:
-    logger.opt(depth=1).bind(rank0_only=rank0_only).exception(message)
+    @staticmethod
+    def exception(message: str, rank0_only: bool = True) -> None:
+        logger.opt(depth=1).bind(rank0_only=rank0_only).exception(message)
 
 
 def _get_rank(group: Optional[dist.ProcessGroup] = None) -> int:

misc.py

@@ -24,7 +24,7 @@ import time
 from contextlib import ContextDecorator
 from typing import Any, Callable, TypeVar
 
-from . import log
+from .log import log
 import numpy as np
 import termcolor
 import torch

model_config.py

@@ -25,7 +25,7 @@ from .ar_config_tokenizer import (
 )
 from .ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
 from .ar_tokenizer_text_tokenizer import TextTokenizer
-from . import log
+from .log import log
 from .lazy_config_init import LazyCall as L
 
 # Common architecture specifications

model_t2w.py

@@ -27,7 +27,7 @@ from .blocks import FourierFeatures
 from .pretrained_vae import BaseVAE
 from . import misc
 from . import instantiate as lazy_instantiate
-from . import log
+from .log import log
 
 
 class EDMSDE:

model_v2w.py

@@ -16,7 +16,7 @@
 from dataclasses import dataclass
 from typing import Callable, Dict, Optional, Tuple, Union
 
-from . import log
+from .log import log
 import torch
 from torch import Tensor
 

presets.py

@@ -22,7 +22,7 @@ from .blocklist import Blocklist
 from .guardrail_core import GuardrailRunner
 from .face_blur_filter import RetinaFaceFilter
 from .video_content_safety_filter import VideoContentSafetyFilter
-from . import log
+from .log import log
 
 
 def create_text_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:

retinaface_utils.py

@@ -17,7 +17,7 @@ import numpy as np
 import torch
 from pytorch_retinaface.utils.nms.py_cpu_nms import py_cpu_nms
 
-from . import log
+from .log import log
 
 
 # Adapted from https://github.com/biubug6/Pytorch_Retinaface/blob/master/detect.py

t5_text_encoder.py

@@ -19,7 +19,7 @@ import torch
 import transformers
 from transformers import T5EncoderModel, T5TokenizerFast
 
-from . import log
+from .log import log
 
 transformers.logging.set_verbosity_error()
 

text2world.py

@@ -16,7 +16,7 @@
 import argparse
 import os
 
-from . import log
+from .log import log
 import torch
 
 from .inference_utils import add_common_arguments, validate_args

text2world_hf.py

@@ -5,7 +5,7 @@ from transformers import PreTrainedModel, PretrainedConfig
 
 from .inference_utils import add_common_arguments, validate_args
 from .world_generation_pipeline import DiffusionText2WorldGenerationPipeline
-from . import log
+from .log import log
 from . import misc
 from .utils_io import read_prompts_from_file, save_video
 

text2world_prompt_upsampler_inference.py

@@ -27,7 +27,7 @@ from .model_config import create_text_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
 from . import presets as guardrail_presets
-from . import log
+from .log import log
 
 
 def create_prompt_upsampler(checkpoint_dir: str) -> AutoRegressiveModel:

video2world.py

@@ -16,7 +16,7 @@
 import argparse
 import os
 
-from . import log
+from .log import log
 import torch
 
 from .inference_utils import add_common_arguments, check_input_frames, validate_args

video2world_prompt_upsampler_inference.py

@@ -30,7 +30,7 @@ from .model_config import create_vision_language_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
 from . import presets as guardrail_presets
-from . import log
+from .log import log
 from .utils_io import load_from_fileobj
 
 

video_content_safety_filter.py

@@ -18,7 +18,7 @@ import json
 import os
 from typing import Iterable, Tuple, Union
 
-from . import log
+from .log import log
 import torch
 from PIL import Image
 

vit.py

@@ -28,7 +28,7 @@ import torch.nn as nn
 
 from .ar_modules_normalization import create_norm
 from .ar_transformer import TransformerBlock
-from . import log
+from .log import log
 
 
 def get_vit_config(model_name: str) -> Mapping[str, Any]:

world_generation_pipeline.py

@@ -43,7 +43,7 @@ from .video2world_prompt_upsampler_inference import (
 from .video2world_prompt_upsampler_inference import (
     run_chat_completion as run_chat_completion_vlm,
 )
-from . import log
+from .log import log
 
 MODEL_NAME_DICT = {
     "Cosmos-1.0-Diffusion-7B-Text2World": "Cosmos_1_0_Diffusion_Text2World_7B",