move df_config out

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

@@ -18,8 +18,8 @@ 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 df_base_model import LatentDiffusionDecoderModelConfig
+from df_config_registry import register_configs
 from . import config
 from config_helper import import_all_modules_from_package
 

cosmos1/models/autoregressive/diffusion_decoder/network.py

@@ -21,7 +21,7 @@ from torch import nn
 from torchvision import transforms
 
 from blocks import PatchEmbed
-from cosmos1.models.diffusion.networks.general_dit import GeneralDIT
+from general_dit import GeneralDIT
 
 
 class DiffusionDecoderGeneralDIT(GeneralDIT):

cosmos1/models/autoregressive/inference/world_generation_pipeline.py

@@ -40,7 +40,7 @@ from inference_utils import (
     load_network_model,
     load_tokenizer_model,
 )
-from . import misc
+from .misc import misc, Color, timer
 
 
 def detect_model_size_from_ckpt_path(ckpt_path: str) -> str:

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

@@ -36,7 +36,7 @@ from nemo.lightning.ckpt_utils import ckpt_to_context_subdir
 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 .presets import presets as guardrail_presets
 from .log import log
 
 torch._C._jit_set_texpr_fuser_enabled(False)

cosmos1/models/diffusion/networks/general_dit_video_conditioned.py

@@ -21,7 +21,7 @@ from torch import nn
 
 from conditioner import DataType
 from blocks import TimestepEmbedding, Timesteps
-from cosmos1.models.diffusion.networks.general_dit import GeneralDIT
+from general_dit import GeneralDIT
 from .log import log
 
 

cosmos1/models/diffusion/config/base/model.py → df_base_model.py

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

cosmos1/models/diffusion/config/base/net.py → df_config_base_net.py

@@ -15,9 +15,9 @@
 
 import copy
 
-from cosmos1.models.diffusion.networks.general_dit import GeneralDIT
-from lazy_config_init import LazyCall as L
-from lazy_config_init import LazyDict
+from .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,

cosmos1/models/diffusion/config/base/tokenizer.py → df_config_base_tokenizer.py

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

cosmos1/models/diffusion/config/config.py → df_config_config.py

@@ -17,14 +17,14 @@ 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 . import config
-from config_helper import import_all_modules_from_package
+from .df_base_model import DefaultModelConfig
+from .df_config_registry import register_configs
+from .config import Config as ori_Config
+from .config_helper import import_all_modules_from_package
 
 
 @attrs.define(slots=False)
-class Config(config.Config):
+class Config(ori_Config):
     # default config groups that will be used unless overwritten
     # see config groups in registry.py
     defaults: List[Any] = attrs.field(

cosmos1/models/diffusion/config/registry.py → df_config_registry.py

@@ -15,13 +15,13 @@
 
 from hydra.core.config_store import ConfigStore
 
-from config_base_conditioner import (
+from .config_base_conditioner import (
     BaseVideoConditionerConfig,
     VideoConditionerFpsSizePaddingConfig,
     VideoExtendConditionerConfig,
 )
-from cosmos1.models.diffusion.config.base.net import FADITV2_14B_Config, FADITV2Config
-from cosmos1.models.diffusion.config.base.tokenizer import get_cosmos_diffusion_tokenizer_comp8x8x8
+from .df_config_base_net import FADITV2_14B_Config, FADITV2Config
+from .df_config_base_tokenizer import get_cosmos_diffusion_tokenizer_comp8x8x8
 
 
 def register_net(cs):

cosmos1/models/diffusion/networks/general_dit.py → general_dit.py

@@ -24,16 +24,16 @@ from einops import rearrange
 from torch import nn
 from torchvision import transforms
 
-from conditioner import DataType
-from attention import get_normalization
-from blocks import (
+from .conditioner import DataType
+from .attention import get_normalization
+from .blocks import (
     FinalLayer,
     GeneralDITTransformerBlock,
     PatchEmbed,
     TimestepEmbedding,
     Timesteps,
 )
-from position_embedding import LearnablePosEmbAxis, VideoRopePosition3DEmb
+from .position_embedding import LearnablePosEmbAxis, VideoRopePosition3DEmb
 from .log import log
 
 

text2world_hf.py

@@ -8,6 +8,7 @@ from .world_generation_pipeline import DiffusionText2WorldGenerationPipeline
 from .log import log
 from .misc import misc, Color, timer
 from .utils_io import read_prompts_from_file, save_video
+from .df_config_config import attrs  # this makes huggingface to download the file
 
 
 # custom config class

world_generation_pipeline.py

@@ -140,7 +140,7 @@ class DiffusionText2WorldGenerationPipeline(BaseWorldGenerationPipeline):
     def _load_model(self):
         self.model = load_model_by_config(
             config_job_name=self.model_name,
-            config_file="cosmos1/models/diffusion/config/config.py",
+            config_file="df_config_config.py",
             model_class=DiffusionT2WModel,
         )
 
@@ -468,7 +468,7 @@ class DiffusionVideo2WorldGenerationPipeline(DiffusionText2WorldGenerationPipeli
     def _load_model(self):
         self.model = load_model_by_config(
             config_job_name=self.model_name,
-            config_file="cosmos1/models/diffusion/config/config.py",
+            config_file="df_config_config.py",
             model_class=DiffusionV2WModel,
         )