update all relative import

ar_config_base_model.py

@@ -17,7 +17,7 @@ from typing import Optional
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import TokenizerConfig
+from .ar_config_base_tokenizer import TokenizerConfig
 
 
 @attrs.define

ar_config_base_model_config.py

@@ -16,17 +16,17 @@
 import copy
 from typing import Callable, List, Optional
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_model import ModelConfig
-from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import (
+from .ar_config_base_model import ModelConfig
+from .ar_config_base_tokenizer import (
     TextTokenizerConfig,
     TokenizerConfig,
     VideoTokenizerConfig,
     create_discrete_video_fsq_tokenizer_state_dict_config,
 )
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_text_tokenizer import TextTokenizer
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
+from .ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
+from .ar_tokenizer_text_tokenizer import TextTokenizer
+from .log import log
+from .lazy_config_init import LazyCall as L
 
 # Common architecture specifications
 BASE_CONFIG = {"n_kv_heads": 8, "norm_type": "rmsnorm", "norm_eps": 1e-5, "ffn_hidden_size": 14336}

ar_config_base_tokenizer.py

@@ -17,10 +17,10 @@ from typing import Optional
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQStateDictTokenizer
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_networks import CausalDiscreteVideoTokenizer
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQStateDictTokenizer
+from .ar_tokenizer_networks import CausalDiscreteVideoTokenizer
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 
 def create_discrete_video_fsq_tokenizer_state_dict_config(

ar_config_inference_inference_config.py

@@ -17,7 +17,7 @@ from typing import Any, List, Union
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_model import ModelConfig, TokenizerConfig
+from .ar_config_base_model import ModelConfig, TokenizerConfig
 
 
 @attrs.define(slots=False)

ar_diffusion_decoder_config_base_conditioner.py

@@ -18,8 +18,8 @@ from typing import Dict, Optional
 
 import torch
 
-from AutoregressiveVideo2WorldGeneration.df_conditioner import BaseVideoCondition, GeneralConditioner
-from AutoregressiveVideo2WorldGeneration.df_config_base_conditioner import (
+from .df_conditioner import BaseVideoCondition, GeneralConditioner
+from .df_config_base_conditioner import (
     FPSConfig,
     ImageSizeConfig,
     LatentConditionConfig,
@@ -28,8 +28,8 @@ from AutoregressiveVideo2WorldGeneration.df_config_base_conditioner import (
     PaddingMaskConfig,
     TextConfig,
 )
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 
 @dataclass

ar_diffusion_decoder_config_config_latent_diffusion_decoder.py

@@ -17,11 +17,11 @@ from typing import Any, List
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_config_registry import register_configs as register_dd_configs
-from AutoregressiveVideo2WorldGeneration.df_config_base_model import LatentDiffusionDecoderModelConfig
-from AutoregressiveVideo2WorldGeneration.df_config_registry import register_configs
-from AutoregressiveVideo2WorldGeneration import config
-from AutoregressiveVideo2WorldGeneration.config_helper import import_all_modules_from_package
+from .ar_diffusion_decoder_config_registry import register_configs as register_dd_configs
+from .df_config_base_model import LatentDiffusionDecoderModelConfig
+from .df_config_registry import register_configs
+from .config import config
+from .config_helper import import_all_modules_from_package
 
 
 @attrs.define(slots=False)

ar_diffusion_decoder_config_inference_cosmos_diffusiondecoder_7b.py

@@ -15,9 +15,9 @@
 
 from hydra.core.config_store import ConfigStore
 
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_network import DiffusionDecoderGeneralDIT
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .ar_diffusion_decoder_network import DiffusionDecoderGeneralDIT
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 num_frames = 57
 Cosmos_DiffusionDecoder_7B_INFERENCE_ONLY: LazyDict = LazyDict(

ar_diffusion_decoder_config_registry.py

@@ -15,12 +15,12 @@
 
 from hydra.core.config_store import ConfigStore
 
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_config_base_conditioner import (
+from .ar_diffusion_decoder_config_base_conditioner import (
     VideoLatentDiffusionDecoderConditionerConfig,
 )
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
-from AutoregressiveVideo2WorldGeneration.df_module_pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from .df_module_pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
+from .lazy_config_init import LazyCall as L
 
 
 def get_cosmos_video_discrete_tokenizer_comp8x16x16(

ar_diffusion_decoder_inference.py

@@ -19,10 +19,10 @@ from typing import List
 
 import torch
 
-from AutoregressiveVideo2WorldGeneration.ar_config_inference_inference_config import DiffusionDecoderSamplingConfig
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_model import LatentDiffusionDecoderModel
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_utils import linear_blend_video_list, split_with_overlap
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_config_inference_inference_config import DiffusionDecoderSamplingConfig
+from .ar_diffusion_decoder_model import LatentDiffusionDecoderModel
+from .ar_diffusion_decoder_utils import linear_blend_video_list, split_with_overlap
+from .log import log
 
 
 def diffusion_decoder_process_tokens(

ar_diffusion_decoder_model.py

@@ -19,11 +19,11 @@ from typing import Callable, Dict, Optional, Tuple
 import torch
 from torch import Tensor
 
-from AutoregressiveVideo2WorldGeneration.df_conditioner import BaseVideoCondition
-from AutoregressiveVideo2WorldGeneration.df_df_functional_batch_ops import batch_mul
-from AutoregressiveVideo2WorldGeneration.df_df_module_res_sampler import COMMON_SOLVER_OPTIONS
-from AutoregressiveVideo2WorldGeneration.df_model_model_t2w import DiffusionT2WModel as VideoDiffusionModel
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import instantiate as lazy_instantiate
+from .df_conditioner import BaseVideoCondition
+from .df_df_functional_batch_ops import batch_mul
+from .df_df_module_res_sampler import COMMON_SOLVER_OPTIONS
+from .df_model_model_t2w import DiffusionT2WModel as VideoDiffusionModel
+from .lazy_config_init import instantiate as lazy_instantiate
 
 
 @dataclass

ar_diffusion_decoder_network.py

@@ -20,8 +20,8 @@ from einops import rearrange
 from torch import nn
 from torchvision import transforms
 
-from AutoregressiveVideo2WorldGeneration.df_module_blocks import PatchEmbed
-from AutoregressiveVideo2WorldGeneration.df_network_general_dit import GeneralDIT
+from .df_module_blocks import PatchEmbed
+from .df_network_general_dit import GeneralDIT
 
 
 class DiffusionDecoderGeneralDIT(GeneralDIT):

ar_model.py

@@ -19,24 +19,24 @@ import time
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set
 
-from AutoregressiveVideo2WorldGeneration import misc
+from .misc import misc
 import torch
 from safetensors.torch import load_file
 from torch.nn.modules.module import _IncompatibleKeys
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_model import ModelConfig
-from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import TokenizerConfig
-from AutoregressiveVideo2WorldGeneration.ar_module_mm_projector import MultimodalProjector
-from AutoregressiveVideo2WorldGeneration.ar_network_transformer import Transformer
-from AutoregressiveVideo2WorldGeneration.ar_network_vit import VisionTransformer, get_vit_config
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_tokenizer import DiscreteMultimodalTokenizer, update_vocab_size
-from AutoregressiveVideo2WorldGeneration.ar_utils_checkpoint import (
+from .ar_config_base_model import ModelConfig
+from .ar_config_base_tokenizer import TokenizerConfig
+from .ar_module_mm_projector import MultimodalProjector
+from .ar_network_transformer import Transformer
+from .ar_network_vit import VisionTransformer, get_vit_config
+from .ar_tokenizer_tokenizer import DiscreteMultimodalTokenizer, update_vocab_size
+from .ar_utils_checkpoint import (
     get_partial_state_dict,
     process_state_dict,
     substrings_to_ignore,
 )
-from AutoregressiveVideo2WorldGeneration.ar_utils_sampling import decode_n_tokens, decode_one_token, prefill
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_utils_sampling import decode_n_tokens, decode_one_token, prefill
+from .log import log
 
 
 class AutoRegressiveModel(torch.nn.Module):

ar_module_attention.py

@@ -19,8 +19,8 @@ from typing import Optional, Union
 import torch
 from torch import nn
 
-from AutoregressiveVideo2WorldGeneration.ar_module_embedding import RotaryPositionEmbedding
-from AutoregressiveVideo2WorldGeneration.ar_module_normalization import create_norm
+from .ar_module_embedding import RotaryPositionEmbedding
+from .ar_module_normalization import create_norm
 
 
 class Attention(nn.Module):

ar_network_transformer.py

@@ -19,17 +19,17 @@ import torch
 import torch.nn as nn
 from torch.nn.modules.module import _IncompatibleKeys
 
-from AutoregressiveVideo2WorldGeneration.ar_module_attention import Attention
-from AutoregressiveVideo2WorldGeneration.ar_module_embedding import (
+from .ar_module_attention import Attention
+from .ar_module_embedding import (
     RotaryPositionEmbeddingPytorchV1,
     RotaryPositionEmbeddingPytorchV2,
     SinCosPosEmbAxisTE,
 )
-from AutoregressiveVideo2WorldGeneration.ar_module_mlp import MLP
-from AutoregressiveVideo2WorldGeneration.ar_module_normalization import create_norm
-from AutoregressiveVideo2WorldGeneration.ar_utils_checkpoint import process_state_dict, substrings_to_ignore
-from AutoregressiveVideo2WorldGeneration.ar_utils_misc import maybe_convert_to_namespace
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_module_mlp import MLP
+from .ar_module_normalization import create_norm
+from .ar_utils_checkpoint import process_state_dict, substrings_to_ignore
+from .ar_utils_misc import maybe_convert_to_namespace
+from .log import log
 
 
 class TransformerBlock(nn.Module):

ar_network_vit.py

@@ -26,9 +26,9 @@ from typing import Any, Callable, Mapping, Optional, Tuple
 import torch
 import torch.nn as nn
 
-from AutoregressiveVideo2WorldGeneration.ar_module_normalization import create_norm
-from AutoregressiveVideo2WorldGeneration.ar_network_transformer import TransformerBlock
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_module_normalization import create_norm
+from .ar_network_transformer import TransformerBlock
+from .log import log
 
 
 def get_vit_config(model_name: str) -> Mapping[str, Any]:

ar_tokenizer_discrete_video.py

@@ -18,7 +18,7 @@ from typing import Optional
 import torch
 from einops import rearrange
 
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_quantizers import FSQuantizer
+from .ar_tokenizer_quantizers import FSQuantizer
 
 # Make sure jit model output consistenly during consecutive calls
 # Check here: https://github.com/pytorch/pytorch/issues/74534

ar_tokenizer_image_text_tokenizer.py

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

ar_tokenizer_modules.py

@@ -29,8 +29,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_patching import Patcher3D, UnPatcher3D
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_utils import (
+from .ar_tokenizer_patching import Patcher3D, UnPatcher3D
+from .ar_tokenizer_utils import (
     CausalNormalize,
     batch2space,
     batch2time,
@@ -41,7 +41,7 @@ from AutoregressiveVideo2WorldGeneration.ar_tokenizer_utils import (
     space2batch,
     time2batch,
 )
-from AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 
 class CausalConv3d(nn.Module):

ar_tokenizer_networks.py

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

ar_tokenizer_quantizers.py

@@ -21,7 +21,7 @@ import torch
 import torch.nn as nn
 from einops import rearrange
 
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_utils import default, pack_one, round_ste, unpack_one
+from .ar_tokenizer_utils import default, pack_one, round_ste, unpack_one
 
 
 class FSQuantizer(nn.Module):

ar_tokenizer_text_tokenizer.py

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

ar_tokenizer_tokenizer.py

@@ -19,8 +19,8 @@ from typing import Optional
 import torch
 from einops import rearrange
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import TokenizerConfig
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import instantiate as lazy_instantiate
+from .ar_config_base_tokenizer import TokenizerConfig
+from .lazy_config_init import instantiate as lazy_instantiate
 
 
 def update_vocab_size(

ar_utils_inference.py

@@ -25,8 +25,8 @@ import torch
 import torchvision
 from PIL import Image
 
-from AutoregressiveVideo2WorldGeneration.ar_config_inference_inference_config import SamplingConfig
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_config_inference_inference_config import SamplingConfig
+from .log import log
 
 _IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", "webp"]
 _VIDEO_EXTENSIONS = [".mp4"]

ar_utils_sampling.py

@@ -17,7 +17,7 @@ from typing import Optional, Tuple
 
 import torch
 
-from AutoregressiveVideo2WorldGeneration.ar_network_transformer import Transformer
+from .ar_network_transformer import Transformer
 
 
 def sample_top_p(logits, temperature, top_p, return_probs: bool = False):

base.py

@@ -19,9 +19,9 @@ import os
 import imageio
 import torch
 
-from AutoregressiveVideo2WorldGeneration.world_generation_pipeline import ARBaseGenerationPipeline
-from AutoregressiveVideo2WorldGeneration.ar_utils_inference import add_common_arguments, load_vision_input, validate_args
-from AutoregressiveVideo2WorldGeneration import log
+from .world_generation_pipeline import ARBaseGenerationPipeline
+from .ar_utils_inference import add_common_arguments, load_vision_input, validate_args
+from .log import log
 
 
 def parse_args():

base_world_generation_pipeline.py

@@ -21,8 +21,8 @@ from typing import Any
 import numpy as np
 import torch
 
-from AutoregressiveVideo2WorldGeneration.t5_text_encoder import CosmosT5TextEncoder
-from AutoregressiveVideo2WorldGeneration import guardrail_common_presets as guardrail_presets
+from .t5_text_encoder import CosmosT5TextEncoder
+from .guardrail_common_presets import guardrail_common_presets as guardrail_presets
 
 
 class BaseWorldGenerationPipeline(ABC):

config.py

@@ -19,8 +19,8 @@ from typing import Any, TypeVar
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
-from AutoregressiveVideo2WorldGeneration.misc import Color
+from .lazy_config_init import LazyDict
+from .misc import Color
 
 T = TypeVar("T")
 

config_helper.py

@@ -27,8 +27,8 @@ from hydra import compose, initialize
 from hydra.core.config_store import ConfigStore
 from omegaconf import DictConfig, OmegaConf
 
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.config import Config
+from .log import log
+from .config import Config
 
 
 def is_attrs_or_dataclass(obj) -> bool:

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 AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 
 class RotaryEmbedding3D(RotaryEmbedding):

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

@@ -34,10 +34,10 @@ from nemo.lightning import io
 from nemo.lightning.ckpt_utils import ckpt_to_context_subdir
 
 from cosmos1.models.autoregressive.nemo.utils import run_diffusion_decoder_model
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
-from AutoregressiveVideo2WorldGeneration.ar_utils_inference import load_vision_input
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from .ar_utils_inference import load_vision_input
 from AutoregressiveVideo2WorldGeneration import guardrail_common_presets as guardrail_presets
-from AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 torch._C._jit_set_texpr_fuser_enabled(False)
 

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

@@ -23,8 +23,8 @@ from huggingface_hub import snapshot_download
 from nemo.collections.nlp.data.language_modeling.megatron import indexed_dataset
 
 from cosmos1.models.autoregressive.nemo.utils import read_input_videos
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
-from AutoregressiveVideo2WorldGeneration import log
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from .log import log
 
 TOKENIZER_COMPRESSION_FACTOR = [8, 16, 16]
 DATA_RESOLUTION_SUPPORTED = [640, 1024]

cosmos1/models/autoregressive/nemo/utils.py

@@ -23,16 +23,16 @@ import torch
 import torchvision
 from huggingface_hub import snapshot_download
 
-from AutoregressiveVideo2WorldGeneration.ar_config_inference_inference_config import DiffusionDecoderSamplingConfig
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_inference import diffusion_decoder_process_tokens
-from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_model import LatentDiffusionDecoderModel
-from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import (
+from .ar_config_inference_inference_config import DiffusionDecoderSamplingConfig
+from .ar_diffusion_decoder_inference import diffusion_decoder_process_tokens
+from .ar_diffusion_decoder_model import LatentDiffusionDecoderModel
+from .df_inference_inference_utils import (
     load_network_model,
     load_tokenizer_model,
     skip_init_linear,
 )
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.config_helper import get_config_module, override
+from .log import log
+from .config_helper import get_config_module, override
 
 TOKENIZER_COMPRESSION_FACTOR = [8, 16, 16]
 DATA_RESOLUTION_SUPPORTED = [640, 1024]

cosmos1/models/diffusion/config/config.py

@@ -17,10 +17,10 @@ from typing import Any, List
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.df_config_base_model import DefaultModelConfig
-from AutoregressiveVideo2WorldGeneration.df_config_registry import register_configs
+from .df_config_base_model import DefaultModelConfig
+from .df_config_registry import register_configs
 from AutoregressiveVideo2WorldGeneration import config
-from AutoregressiveVideo2WorldGeneration.config_helper import import_all_modules_from_package
+from .config_helper import import_all_modules_from_package
 
 
 @attrs.define(slots=False)

cosmos1/models/diffusion/config/inference/cosmos-1-diffusion-text2world.py

@@ -15,7 +15,7 @@
 
 from hydra.core.config_store import ConfigStore
 
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .lazy_config_init import LazyDict
 
 Cosmos_1_0_Diffusion_Text2World_7B: LazyDict = LazyDict(
     dict(

cosmos1/models/diffusion/config/inference/cosmos-1-diffusion-video2world.py

@@ -16,8 +16,8 @@
 from hydra.core.config_store import ConfigStore
 
 from cosmos1.models.diffusion.networks.general_dit_video_conditioned import VideoExtendGeneralDIT
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 Cosmos_1_0_Diffusion_Video2World_7B: LazyDict = LazyDict(
     dict(

cosmos1/models/diffusion/inference/text2world.py

@@ -16,13 +16,13 @@
 import argparse
 import os
 
-from AutoregressiveVideo2WorldGeneration import misc
+from .misc import misc
 import torch
 
-from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import add_common_arguments, validate_args
+from .df_inference_inference_utils import add_common_arguments, validate_args
 from cosmos1.models.diffusion.inference.world_generation_pipeline import DiffusionText2WorldGenerationPipeline
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.io import read_prompts_from_file, save_video
+from .log import log
+from .io import read_prompts_from_file, save_video
 
 torch.enable_grad(False)
 

cosmos1/models/diffusion/inference/video2world.py

@@ -16,13 +16,13 @@
 import argparse
 import os
 
-from AutoregressiveVideo2WorldGeneration import misc
+from .misc import misc
 import torch
 
-from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import add_common_arguments, check_input_frames, validate_args
+from .df_inference_inference_utils import add_common_arguments, check_input_frames, validate_args
 from cosmos1.models.diffusion.inference.world_generation_pipeline import DiffusionVideo2WorldGenerationPipeline
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.io import read_prompts_from_file, save_video
+from .log import log
+from .io import read_prompts_from_file, save_video
 
 torch.enable_grad(False)
 

cosmos1/models/diffusion/inference/world_generation_pipeline.py

@@ -20,8 +20,8 @@ from typing import Any, Optional
 import numpy as np
 import torch
 
-from AutoregressiveVideo2WorldGeneration.base_world_generation_pipeline import BaseWorldGenerationPipeline
-from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import (
+from .base_world_generation_pipeline import BaseWorldGenerationPipeline
+from .df_inference_inference_utils import (
     generate_world_from_text,
     generate_world_from_video,
     get_condition_latent,
@@ -30,8 +30,8 @@ from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import (
     load_network_model,
     load_tokenizer_model,
 )
-from AutoregressiveVideo2WorldGeneration.df_model_model_t2w import DiffusionT2WModel
-from AutoregressiveVideo2WorldGeneration.df_model_model_v2w import DiffusionV2WModel
+from .df_model_model_t2w import DiffusionT2WModel
+from .df_model_model_v2w import DiffusionV2WModel
 from cosmos1.models.diffusion.prompt_upsampler.text2world_prompt_upsampler_inference import (
     create_prompt_upsampler,
     run_chat_completion,
@@ -43,7 +43,7 @@ from cosmos1.models.diffusion.prompt_upsampler.video2world_prompt_upsampler_infe
 from cosmos1.models.diffusion.prompt_upsampler.video2world_prompt_upsampler_inference import (
     run_chat_completion as run_chat_completion_vlm,
 )
-from AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 MODEL_NAME_DICT = {
     "Cosmos-1.0-Diffusion-7B-Text2World": "Cosmos_1_0_Diffusion_Text2World_7B",

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 AutoregressiveVideo2WorldGeneration 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

@@ -19,18 +19,18 @@ import imageio
 import numpy as np
 import torch
 
-from AutoregressiveVideo2WorldGeneration.ar_model import AutoRegressiveModel
+from .ar_model import AutoRegressiveModel
 from cosmos1.models.diffusion.prompt_upsampler.text2world_prompt_upsampler_inference import (
     create_prompt_upsampler,
     run_chat_completion,
 )
-from AutoregressiveVideo2WorldGeneration.guardrail_common_presets import (
+from .guardrail_common_presets import (
     create_text_guardrail_runner,
     create_video_guardrail_runner,
     run_text_guardrail,
     run_video_guardrail,
 )
-from AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 
 def get_upsampled_prompt(

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 AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 
 def get_parser():

cosmos1/models/diffusion/networks/general_dit_video_conditioned.py

@@ -19,10 +19,10 @@ import torch
 from einops import rearrange
 from torch import nn
 
-from AutoregressiveVideo2WorldGeneration.df_conditioner import DataType
-from AutoregressiveVideo2WorldGeneration.df_module_blocks import TimestepEmbedding, Timesteps
-from AutoregressiveVideo2WorldGeneration.df_network_general_dit import GeneralDIT
-from AutoregressiveVideo2WorldGeneration import log
+from .df_conditioner import DataType
+from .df_module_blocks import TimestepEmbedding, Timesteps
+from .df_network_general_dit import GeneralDIT
+from .log import log
 
 
 class VideoExtendGeneralDIT(GeneralDIT):

cosmos1/models/diffusion/prompt_upsampler/inference.py

@@ -17,9 +17,9 @@ from typing import List, Optional, TypedDict
 
 import torch
 
-from AutoregressiveVideo2WorldGeneration.ar_model import AutoRegressiveModel
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
-from AutoregressiveVideo2WorldGeneration.ar_tokenizer_text_tokenizer import TextTokenizer
+from .ar_model import AutoRegressiveModel
+from .ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
+from .ar_tokenizer_text_tokenizer import TextTokenizer
 
 
 class ChatPrediction(TypedDict, total=False):

cosmos1/models/diffusion/prompt_upsampler/text2world_prompt_upsampler_inference.py

@@ -23,11 +23,11 @@ import argparse
 import os
 import re
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_model_config import create_text_model_config
-from AutoregressiveVideo2WorldGeneration.ar_model import AutoRegressiveModel
+from .ar_config_base_model_config import create_text_model_config
+from .ar_model import AutoRegressiveModel
 from cosmos1.models.diffusion.prompt_upsampler.inference import chat_completion
 from AutoregressiveVideo2WorldGeneration import guardrail_common_presets as guardrail_presets
-from AutoregressiveVideo2WorldGeneration import log
+from .log import log
 
 
 def create_prompt_upsampler(checkpoint_dir: str) -> AutoRegressiveModel:

cosmos1/models/diffusion/prompt_upsampler/video2world_prompt_upsampler_inference.py

@@ -26,12 +26,12 @@ from math import ceil
 
 from PIL import Image
 
-from AutoregressiveVideo2WorldGeneration.ar_config_base_model_config import create_vision_language_model_config
-from AutoregressiveVideo2WorldGeneration.ar_model import AutoRegressiveModel
+from .ar_config_base_model_config import create_vision_language_model_config
+from .ar_model import AutoRegressiveModel
 from cosmos1.models.diffusion.prompt_upsampler.inference import chat_completion
 from AutoregressiveVideo2WorldGeneration import guardrail_common_presets as guardrail_presets
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.io import load_from_fileobj
+from .log import log
+from .io import load_from_fileobj
 
 
 def create_vlm_prompt_upsampler(

df_conditioner.py

@@ -23,9 +23,9 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 
-from AutoregressiveVideo2WorldGeneration.df_df_functional_batch_ops import batch_mul
-from AutoregressiveVideo2WorldGeneration import log
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import instantiate
+from .df_df_functional_batch_ops import batch_mul
+from .log import log
+from .lazy_config_init import instantiate
 
 
 class BaseConditionEntry(nn.Module):

df_config_base_conditioner.py

@@ -18,9 +18,9 @@ from typing import Dict, List, Optional
 import attrs
 import torch
 
-from AutoregressiveVideo2WorldGeneration.df_conditioner import BaseConditionEntry, TextAttr, VideoConditioner, VideoExtendConditioner
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .df_conditioner import BaseConditionEntry, TextAttr, VideoConditioner, VideoExtendConditioner
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 
 @attrs.define(slots=False)

df_config_base_model.py

@@ -17,7 +17,7 @@ from typing import List
 
 import attrs
 
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .lazy_config_init import LazyDict
 
 
 @attrs.define(slots=False)

df_config_base_net.py

@@ -15,9 +15,9 @@
 
 import copy
 
-from AutoregressiveVideo2WorldGeneration.df_network_general_dit import GeneralDIT
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
+from .df_network_general_dit import GeneralDIT
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
 
 FADITV2Config: LazyDict = L(GeneralDIT)(
     max_img_h=240,

df_config_base_tokenizer.py

@@ -15,8 +15,8 @@
 
 import omegaconf
 
-from AutoregressiveVideo2WorldGeneration.df_module_pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
-from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
+from .df_module_pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
+from .lazy_config_init import LazyCall as L
 
 TOKENIZER_OPTIONS = {}