Finished rearranging model and config files

.gitignore

@@ -0,0 +1 @@
+huggingface.txt

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

@@ -19,9 +19,9 @@ from typing import List
 
 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 .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
 
 

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 conditioner import BaseVideoCondition
-from batch_ops import batch_mul
-from res_sampler import COMMON_SOLVER_OPTIONS
-from model_t2w import DiffusionT2WModel as VideoDiffusionModel
-from lazy_config_init import instantiate as lazy_instantiate
+from .conditioner import BaseVideoCondition
+from .batch_ops import batch_mul
+from .res_sampler import COMMON_SOLVER_OPTIONS
+from .model_t2w import DiffusionT2WModel as VideoDiffusionModel
+from .lazy_config_init import instantiate as lazy_instantiate
 
 
 @dataclass

cosmos1/models/autoregressive/inference/world_generation_pipeline.py

@@ -22,7 +22,7 @@ import numpy as np
 import torch
 from einops import rearrange
 
-from model_config import create_video2world_model_config
+from ar_configs_model_config import create_video2world_model_config
 from ar_config_tokenizer import TokenizerConfig
 from inference_config import (
     DataShapeConfig,
@@ -30,8 +30,8 @@ from inference_config import (
     InferenceConfig,
     SamplingConfig,
 )
-from cosmos1.models.autoregressive.diffusion_decoder.inference import diffusion_decoder_process_tokens
-from cosmos1.models.autoregressive.diffusion_decoder.model import LatentDiffusionDecoderModel
+from cosmos1.models.autoregressive.diffusion_decoder.ar_diffusion_decoder_inference import diffusion_decoder_process_tokens
+from cosmos1.models.autoregressive.diffusion_decoder.ar_diffusion_decoder_model import LatentDiffusionDecoderModel
 from ar_model import AutoRegressiveModel
 from cosmos1.models.autoregressive.utils.inference import _SUPPORTED_CONTEXT_LEN, prepare_video_batch_for_saving
 from base_world_generation_pipeline import BaseWorldGenerationPipeline

cosmos1/models/autoregressive/nemo/utils.py

@@ -24,8 +24,8 @@ import torchvision
 from huggingface_hub import snapshot_download
 
 from 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.autoregressive.diffusion_decoder.ar_diffusion_decoder_inference import diffusion_decoder_process_tokens
+from cosmos1.models.autoregressive.diffusion_decoder.ar_diffusion_decoder_model import LatentDiffusionDecoderModel
 from inference_utils import (
     load_network_model,
     load_tokenizer_model,

futureworld_hf.py

@@ -19,15 +19,23 @@ class AutoregressiveFutureWorldConfig(PretrainedConfig):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.checkpoint_dir = kwargs.get("checkpoint_dir", "checkpoints")
-        self.checkpoint_name = kwargs.get("checkpoint_name", "Cosmos-1.0-Autoregressive-4B")
+        self.ar_model_dir = kwargs.get("ar_model_dir", "Cosmos-1.0-Autoregressive-4B")
         self.disable_diffusion_decoder = kwargs.get("disable_diffusion_decoder", False)
         self.offload_guardrail_models = kwargs.get("offload_guardrail_models", False)
         self.offload_diffusion_decoder = kwargs.get("offload_diffusion_decoder", False)
-        self.offload_network = kwargs.get("offload_network", False)
+        self.offload_ar_model = kwargs.get("offload_ar_model", False)
         self.offload_tokenizer = kwargs.get("offload_tokenizer", False)
         self.video_save_name = kwargs.get("video_save_name", "output")
         self.video_save_folder = kwargs.get("video_save_folder", "outputs/")
-        self.seed = kwargs.get()
+        self.seed = kwargs.get("seed", 0)
+        self.temperature = kwargs.get("temperature", 1.0)
+        self.top_p = kwargs.get("top_p", 0.8)
+        self.input_type = None
+        self.batch_input_path = None
+        self.input_image_or_video_path = None
+        self.data_resolution = None
+        self.num_input_frames = None
+
 
 # custom model class
 class AutoregressiveFutureWorld(PreTrainedModel):
@@ -37,17 +45,16 @@ class AutoregressiveFutureWorld(PreTrainedModel):
         super().__init__(config)
         torch._C._jit_set_texpr_fuser_enabled(False)
         self.config = config
-        inference_type = "base"
-        sampling_config = validate_args(config, inference_type)
+        self.inference_type = "base"
         self.pipeline = ARBaseGenerationPipeline(
-            inference_type=inference_type,
-            checkpoint_dir=self.checkpoint_dir,
-            checkpoint_name=self.ar_model_dir,
-            disable_diffusion_decoder=self.disable_diffusion_decoder,
-            offload_guardrail_models=self.offload_guardrail_models,
-            offload_diffusion_decoder=self.offload_diffusion_decoder,
-            offload_network=self.offload_ar_model,
-            offload_tokenizer=self.offload_tokenizer,
+            inference_type=self.inference_type,
+            checkpoint_dir=self.config.checkpoint_dir,
+            checkpoint_name=self.config.ar_model_dir,
+            disable_diffusion_decoder=self.config.disable_diffusion_decoder,
+            offload_guardrail_models=self.config.offload_guardrail_models,
+            offload_diffusion_decoder=self.config.offload_diffusion_decoder,
+            offload_network=self.config.offload_ar_model,
+            offload_tokenizer=self.config.offload_tokenizer,
         )
 
     # modifed from text2world.py demo function
@@ -63,6 +70,12 @@ class AutoregressiveFutureWorld(PreTrainedModel):
             data_resolution=data_resolution,
             num_input_frames=num_input_frames,
         )
+        self.config.input_type = input_type
+        self.config.batch_input_path = batch_input_path
+        self.config.input_image_or_video_path = input_image_or_video_path
+        self.config.data_resolution = data_resolution
+        self.config.num_input_frames = num_input_frames
+        sampling_config = validate_args(self.config, self.inference_type)
 
         for idx, input_filename in enumerate(input_videos):
             inp_vid = input_videos[input_filename]
@@ -71,7 +84,7 @@ class AutoregressiveFutureWorld(PreTrainedModel):
             out_vid = self.pipeline.generate(
                 inp_vid=inp_vid,
                 num_input_frames=num_input_frames,
-                seed=args.seed,
+                seed=self.config.seed,
                 sampling_config=sampling_config,
             )
             if out_vid is None:
@@ -80,9 +93,9 @@ class AutoregressiveFutureWorld(PreTrainedModel):
 
             # Save video
             if input_image_or_video_path:
-                out_vid_path = os.path.join(args.video_save_folder, f"{args.video_save_name}.mp4")
+                out_vid_path = os.path.join(self.config.video_save_folder, f"{self.config.video_save_name}.mp4")
             else:
-                out_vid_path = os.path.join(args.video_save_folder, f"{idx}.mp4")
+                out_vid_path = os.path.join(self.config.video_save_folder, f"{idx}.mp4")
 
             imageio.mimsave(out_vid_path, out_vid, fps=25)
 

text2world_prompt_upsampler_inference.py

@@ -23,7 +23,7 @@ import argparse
 import os
 import re
 
-from .model_config import create_text_model_config
+from .ar_configs_model_config import create_text_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
 from .presets import presets as guardrail_presets

video2world_prompt_upsampler_inference.py

@@ -26,7 +26,7 @@ from math import ceil
 
 from PIL import Image
 
-from .model_config import create_vision_language_model_config
+from .ar_configs_model_config import create_vision_language_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
 from .presets import presets as guardrail_presets

world_generation_pipeline.py

@@ -21,25 +21,26 @@ import numpy as np
 import torch
 from einops import rearrange
 
-from cosmos1.models.autoregressive.configs.base.model_config import create_video2world_model_config
-from cosmos1.models.autoregressive.configs.base.tokenizer import TokenizerConfig
-from cosmos1.models.autoregressive.configs.inference.inference_config import (
+from .ar_configs_model_config import create_video2world_model_config
+from .ar_config_tokenizer import TokenizerConfig
+from .ar_configs_inference import (
     DataShapeConfig,
     DiffusionDecoderSamplingConfig,
     InferenceConfig,
     SamplingConfig,
 )
-from cosmos1.models.autoregressive.diffusion_decoder.inference import diffusion_decoder_process_tokens
-from cosmos1.models.autoregressive.diffusion_decoder.model import LatentDiffusionDecoderModel
-from cosmos1.models.autoregressive.model import AutoRegressiveModel
-from cosmos1.models.autoregressive.utils.inference import _SUPPORTED_CONTEXT_LEN, prepare_video_batch_for_saving
-from cosmos1.models.common.base_world_generation_pipeline import BaseWorldGenerationPipeline
-from cosmos1.models.diffusion.inference.inference_utils import (
+from .ar_diffusion_decoder_inference import diffusion_decoder_process_tokens
+from .ar_diffusion_decoder_model import LatentDiffusionDecoderModel
+from .ar_model import AutoRegressiveModel
+from .ar_utils_inference import _SUPPORTED_CONTEXT_LEN, prepare_video_batch_for_saving
+from .base_world_generation_pipeline import BaseWorldGenerationPipeline
+from .inference_utils import (
     load_model_by_config,
     load_network_model,
     load_tokenizer_model,
 )
-from cosmos1.utils import log, misc
+from .log import log
+from .misc import misc
 
 
 def detect_model_size_from_ckpt_path(ckpt_path: str) -> str: