reorganize files

cosmos1/models/autoregressive/configs/base/model.py → ar_config_base_model.py

@@ -17,7 +17,7 @@ from typing import Optional
 
 import attrs
 
-from cosmos1.models.autoregressive.configs.base.tokenizer import TokenizerConfig
+from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import TokenizerConfig
 
 
 @attrs.define

cosmos1/models/autoregressive/configs/base/model_config.py → ar_config_base_model_config.py

@@ -16,17 +16,17 @@
 import copy
 from typing import Callable, List, Optional
 
-from cosmos1.models.autoregressive.configs.base.model import ModelConfig
-from cosmos1.models.autoregressive.configs.base.tokenizer import (
+from AutoregressiveVideo2WorldGeneration.ar_config_base_model import ModelConfig
+from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import (
     TextTokenizerConfig,
     TokenizerConfig,
     VideoTokenizerConfig,
     create_discrete_video_fsq_tokenizer_state_dict_config,
 )
-from cosmos1.models.autoregressive.tokenizer.image_text_tokenizer import ImageTextTokenizer
-from cosmos1.models.autoregressive.tokenizer.text_tokenizer import TextTokenizer
-from cosmos1.utils import log
-from cosmos1.utils.lazy_config import LazyCall as L
+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
 
 # Common architecture specifications
 BASE_CONFIG = {"n_kv_heads": 8, "norm_type": "rmsnorm", "norm_eps": 1e-5, "ffn_hidden_size": 14336}

cosmos1/models/autoregressive/configs/base/tokenizer.py → ar_config_base_tokenizer.py

@@ -17,10 +17,10 @@ from typing import Optional
 
 import attrs
 
-from cosmos1.models.autoregressive.tokenizer.discrete_video import DiscreteVideoFSQStateDictTokenizer
-from cosmos1.models.autoregressive.tokenizer.networks import CausalDiscreteVideoTokenizer
-from cosmos1.utils.lazy_config import LazyCall as L
-from cosmos1.utils.lazy_config import LazyDict
+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
 
 
 def create_discrete_video_fsq_tokenizer_state_dict_config(

cosmos1/models/autoregressive/configs/inference/inference_config.py → ar_config_inference_inference_config.py

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

cosmos1/models/autoregressive/diffusion_decoder/config/base/conditioner.py → ar_diffusion_decoder_config_base_conditioner.py

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

cosmos1/models/autoregressive/diffusion_decoder/config/config_latent_diffusion_decoder.py → ar_diffusion_decoder_config_config_latent_diffusion_decoder.py

@@ -17,11 +17,11 @@ from typing import Any, List
 
 import attrs
 
-from cosmos1.models.autoregressive.diffusion_decoder.config.registry import register_configs as register_dd_configs
-from cosmos1.models.diffusion.config.base.model import LatentDiffusionDecoderModelConfig
-from cosmos1.models.diffusion.config.registry import register_configs
-from cosmos1.utils import config
-from cosmos1.utils.config_helper import import_all_modules_from_package
+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
 
 
 @attrs.define(slots=False)

cosmos1/models/autoregressive/diffusion_decoder/config/inference/cosmos_diffusiondecoder_7b.py → ar_diffusion_decoder_config_inference_cosmos_diffusiondecoder_7b.py

@@ -15,9 +15,9 @@
 
 from hydra.core.config_store import ConfigStore
 
-from cosmos1.models.autoregressive.diffusion_decoder.network import DiffusionDecoderGeneralDIT
-from cosmos1.utils.lazy_config import LazyCall as L
-from cosmos1.utils.lazy_config import LazyDict
+from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_network import DiffusionDecoderGeneralDIT
+from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
+from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
 
 num_frames = 57
 Cosmos_DiffusionDecoder_7B_INFERENCE_ONLY: LazyDict = LazyDict(

cosmos1/models/autoregressive/diffusion_decoder/config/registry.py → ar_diffusion_decoder_config_registry.py

@@ -15,12 +15,12 @@
 
 from hydra.core.config_store import ConfigStore
 
-from cosmos1.models.autoregressive.diffusion_decoder.config.base.conditioner import (
+from AutoregressiveVideo2WorldGeneration.ar_diffusion_decoder_config_base_conditioner import (
     VideoLatentDiffusionDecoderConditionerConfig,
 )
-from cosmos1.models.autoregressive.tokenizer.discrete_video import DiscreteVideoFSQJITTokenizer
-from cosmos1.models.diffusion.module.pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
-from cosmos1.utils.lazy_config import LazyCall as L
+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
 
 
 def get_cosmos_video_discrete_tokenizer_comp8x16x16(

cosmos1/models/autoregressive/diffusion_decoder/inference.py → ar_diffusion_decoder_inference.py

@@ -19,10 +19,10 @@ from typing import List
 
 import torch
 
-from cosmos1.models.autoregressive.configs.inference.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 cosmos1.utils import log
+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
 
 
 def diffusion_decoder_process_tokens(

cosmos1/models/autoregressive/diffusion_decoder/model.py → ar_diffusion_decoder_model.py

@@ -19,11 +19,11 @@ from typing import Callable, Dict, Optional, Tuple
 import torch
 from torch import Tensor
 
-from cosmos1.models.diffusion.conditioner import BaseVideoCondition
-from cosmos1.models.diffusion.diffusion.functional.batch_ops import batch_mul
-from cosmos1.models.diffusion.diffusion.modules.res_sampler import COMMON_SOLVER_OPTIONS
-from cosmos1.models.diffusion.model.model_t2w import DiffusionT2WModel as VideoDiffusionModel
-from cosmos1.utils.lazy_config import instantiate as lazy_instantiate
+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
 
 
 @dataclass

cosmos1/models/autoregressive/diffusion_decoder/network.py → ar_diffusion_decoder_network.py

@@ -20,8 +20,8 @@ from einops import rearrange
 from torch import nn
 from torchvision import transforms
 
-from cosmos1.models.diffusion.module.blocks import PatchEmbed
-from cosmos1.models.diffusion.networks.general_dit import GeneralDIT
+from AutoregressiveVideo2WorldGeneration.df_module_blocks import PatchEmbed
+from AutoregressiveVideo2WorldGeneration.df_network_general_dit import GeneralDIT
 
 
 class DiffusionDecoderGeneralDIT(GeneralDIT):

cosmos1/models/autoregressive/model.py → ar_model.py

@@ -19,23 +19,24 @@ import time
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set
 
+from AutoregressiveVideo2WorldGeneration import misc
 import torch
 from safetensors.torch import load_file
 from torch.nn.modules.module import _IncompatibleKeys
 
-from cosmos1.models.autoregressive.configs.base.model import ModelConfig
-from cosmos1.models.autoregressive.configs.base.tokenizer import TokenizerConfig
-from cosmos1.models.autoregressive.modules.mm_projector import MultimodalProjector
-from cosmos1.models.autoregressive.networks.transformer import Transformer
-from cosmos1.models.autoregressive.networks.vit import VisionTransformer, get_vit_config
-from cosmos1.models.autoregressive.tokenizer.tokenizer import DiscreteMultimodalTokenizer, update_vocab_size
-from cosmos1.models.autoregressive.utils.checkpoint import (
+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 (
     get_partial_state_dict,
     process_state_dict,
     substrings_to_ignore,
 )
-from cosmos1.models.autoregressive.utils.sampling import decode_n_tokens, decode_one_token, prefill
-from cosmos1.utils import log, misc
+from AutoregressiveVideo2WorldGeneration.ar_utils_sampling import decode_n_tokens, decode_one_token, prefill
+from AutoregressiveVideo2WorldGeneration import log
 
 
 class AutoRegressiveModel(torch.nn.Module):

cosmos1/models/autoregressive/modules/attention.py → ar_module_attention.py

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

cosmos1/models/autoregressive/networks/transformer.py → ar_network_transformer.py

@@ -19,17 +19,17 @@ import torch
 import torch.nn as nn
 from torch.nn.modules.module import _IncompatibleKeys
 
-from cosmos1.models.autoregressive.modules.attention import Attention
-from cosmos1.models.autoregressive.modules.embedding import (
+from AutoregressiveVideo2WorldGeneration.ar_module_attention import Attention
+from AutoregressiveVideo2WorldGeneration.ar_module_embedding import (
     RotaryPositionEmbeddingPytorchV1,
     RotaryPositionEmbeddingPytorchV2,
     SinCosPosEmbAxisTE,
 )
-from cosmos1.models.autoregressive.modules.mlp import MLP
-from cosmos1.models.autoregressive.modules.normalization import create_norm
-from cosmos1.models.autoregressive.utils.checkpoint import process_state_dict, substrings_to_ignore
-from cosmos1.models.autoregressive.utils.misc import maybe_convert_to_namespace
-from cosmos1.utils import log
+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
 
 
 class TransformerBlock(nn.Module):

cosmos1/models/autoregressive/networks/vit.py → ar_network_vit.py

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

cosmos1/models/autoregressive/tokenizer/discrete_video.py → ar_tokenizer_discrete_video.py

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

cosmos1/models/autoregressive/tokenizer/image_text_tokenizer.py → 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 cosmos1.models.autoregressive.tokenizer.text_tokenizer import TextTokenizer
-from cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration.ar_tokenizer_text_tokenizer import TextTokenizer
+from AutoregressiveVideo2WorldGeneration import log
 
 # Configuration for different vision-language models
 IMAGE_CONFIGS = {

cosmos1/models/autoregressive/tokenizer/modules.py → ar_tokenizer_modules.py

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

cosmos1/models/autoregressive/tokenizer/networks.py → ar_tokenizer_networks.py

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

cosmos1/models/autoregressive/tokenizer/quantizers.py → ar_tokenizer_quantizers.py

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

cosmos1/models/autoregressive/tokenizer/text_tokenizer.py → ar_tokenizer_text_tokenizer.py

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

cosmos1/models/autoregressive/tokenizer/tokenizer.py → ar_tokenizer_tokenizer.py

@@ -19,8 +19,8 @@ from typing import Optional
 import torch
 from einops import rearrange
 
-from cosmos1.models.autoregressive.configs.base.tokenizer import TokenizerConfig
-from cosmos1.utils.lazy_config import instantiate as lazy_instantiate
+from AutoregressiveVideo2WorldGeneration.ar_config_base_tokenizer import TokenizerConfig
+from AutoregressiveVideo2WorldGeneration.lazy_config_init import instantiate as lazy_instantiate
 
 
 def update_vocab_size(

cosmos1/models/autoregressive/utils/inference.py → ar_utils_inference.py

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

cosmos1/models/autoregressive/utils/sampling.py → ar_utils_sampling.py

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

cosmos1/models/autoregressive/inference/base.py → base.py

@@ -19,9 +19,9 @@ import os
 import imageio
 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 cosmos1.utils import log
+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
 
 
 def parse_args():

cosmos1/models/common/base_world_generation_pipeline.py → base_world_generation_pipeline.py

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

config.json

@@ -0,0 +1,10 @@
+{
+    "architectures": [
+      "ARVideo2World"
+    ],
+    "auto_map": {
+      "AutoConfig": "video2world_hf.ARVideo2WorldConfig",
+      "AutoModel": "video2world_hf.ARVideo2World"
+    },
+    "model_type": "AutoModel"
+}

cosmos1/utils/config.py → config.py

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

cosmos1/utils/config_helper.py → 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 cosmos1.utils import log
-from cosmos1.utils.config import Config
+from AutoregressiveVideo2WorldGeneration import log
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration 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 cosmos1.models.autoregressive.tokenizer.discrete_video import DiscreteVideoFSQJITTokenizer
-from cosmos1.models.autoregressive.utils.inference import load_vision_input
-from cosmos1.models.guardrail.common import presets as guardrail_presets
-from cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from AutoregressiveVideo2WorldGeneration.ar_utils_inference import load_vision_input
+from AutoregressiveVideo2WorldGeneration import guardrail_common_presets as guardrail_presets
+from AutoregressiveVideo2WorldGeneration 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 cosmos1.models.autoregressive.tokenizer.discrete_video import DiscreteVideoFSQJITTokenizer
-from cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration.ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from AutoregressiveVideo2WorldGeneration 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 cosmos1.models.autoregressive.configs.inference.inference_config import DiffusionDecoderSamplingConfig
-from cosmos1.models.autoregressive.diffusion_decoder.inference import diffusion_decoder_process_tokens
-from cosmos1.models.autoregressive.diffusion_decoder.model import LatentDiffusionDecoderModel
-from cosmos1.models.diffusion.inference.inference_utils import (
+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 (
     load_network_model,
     load_tokenizer_model,
     skip_init_linear,
 )
-from cosmos1.utils import log
-from cosmos1.utils.config_helper import get_config_module, override
+from AutoregressiveVideo2WorldGeneration import log
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.models.diffusion.config.base.model import DefaultModelConfig
-from cosmos1.models.diffusion.config.registry import register_configs
-from cosmos1.utils import config
-from cosmos1.utils.config_helper import import_all_modules_from_package
+from AutoregressiveVideo2WorldGeneration.df_config_base_model import DefaultModelConfig
+from AutoregressiveVideo2WorldGeneration.df_config_registry import register_configs
+from AutoregressiveVideo2WorldGeneration import config
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.utils.lazy_config import LazyDict
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.utils.lazy_config import LazyCall as L
-from cosmos1.utils.lazy_config import LazyDict
+from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyCall as L
+from AutoregressiveVideo2WorldGeneration.lazy_config_init import LazyDict
 
 Cosmos_1_0_Diffusion_Video2World_7B: LazyDict = LazyDict(
     dict(

cosmos1/models/diffusion/inference/text2world.py

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

cosmos1/models/diffusion/inference/video2world.py

@@ -16,12 +16,13 @@
 import argparse
 import os
 
+from AutoregressiveVideo2WorldGeneration import misc
 import torch
 
-from cosmos1.models.diffusion.inference.inference_utils import add_common_arguments, check_input_frames, validate_args
+from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import add_common_arguments, check_input_frames, validate_args
 from cosmos1.models.diffusion.inference.world_generation_pipeline import DiffusionVideo2WorldGenerationPipeline
-from cosmos1.utils import log, misc
-from cosmos1.utils.io import read_prompts_from_file, save_video
+from AutoregressiveVideo2WorldGeneration import log
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.models.common.base_world_generation_pipeline import BaseWorldGenerationPipeline
-from cosmos1.models.diffusion.inference.inference_utils import (
+from AutoregressiveVideo2WorldGeneration.base_world_generation_pipeline import BaseWorldGenerationPipeline
+from AutoregressiveVideo2WorldGeneration.df_inference_inference_utils import (
     generate_world_from_text,
     generate_world_from_video,
     get_condition_latent,
@@ -30,8 +30,8 @@ from cosmos1.models.diffusion.inference.inference_utils import (
     load_network_model,
     load_tokenizer_model,
 )
-from cosmos1.models.diffusion.model.model_t2w import DiffusionT2WModel
-from cosmos1.models.diffusion.model.model_v2w import DiffusionV2WModel
+from AutoregressiveVideo2WorldGeneration.df_model_model_t2w import DiffusionT2WModel
+from AutoregressiveVideo2WorldGeneration.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 cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration 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 cosmos1.utils import log
+from AutoregressiveVideo2WorldGeneration import log
 
 EXAMPLE_PROMPT = (
     "The teal robot is cooking food in a kitchen. Steam rises from a simmering pot "