try add . in import

text2world_hf.py

@@ -3,11 +3,11 @@ import argparse
 import torch
 from transformers import PreTrainedModel, PretrainedConfig
 
-from cosmos1.models.diffusion.inference.inference_utils import add_common_arguments, validate_args
-from cosmos1.models.diffusion.inference.world_generation_pipeline import DiffusionText2WorldGenerationPipeline
-import cosmos1.utils.log as log 
-import cosmos1.utils.misc as misc
-from cosmos1.utils.io import read_prompts_from_file, save_video
+from .cosmos1.models.diffusion.inference.inference_utils import add_common_arguments, validate_args
+from .cosmos1.models.diffusion.inference.world_generation_pipeline import DiffusionText2WorldGenerationPipeline
+import .cosmos1.utils.log as log 
+import .cosmos1.utils.misc as misc
+from .cosmos1.utils.io import read_prompts_from_file, save_video
 
 class DiffusionText2WorldConfig(PretrainedConfig):
     model_type = "DiffusionText2World" 
@@ -46,8 +46,28 @@ class DiffusionText2World(PreTrainedModel):
         torch.enable_grad(False)   # TODO: do we need this?
         self.config = config
         inference_type = "text2world"
-        validate_args(argparse.Namespace(**config), inference_type)
-        self.pipeline = DiffusionText2WorldGenerationPipeline(config)
+        config.prompt = 1          # TODO: this is to hack args validation, maybe find a better way 
+        validate_args(config, inference_type) 
+        del config.prompt
+        self.pipeline = DiffusionText2WorldGenerationPipeline(
+            inference_type=inference_type,
+            checkpoint_dir=config.checkpoint_dir,
+            checkpoint_name=config.diffusion_transformer_dir,
+            prompt_upsampler_dir=config.prompt_upsampler_dir,
+            enable_prompt_upsampler=not config.disable_prompt_upsampler,
+            offload_network=config.offload_diffusion_transformer,
+            offload_tokenizer=config.offload_tokenizer,
+            offload_text_encoder_model=config.offload_text_encoder_model,
+            offload_prompt_upsampler=config.offload_prompt_upsampler,
+            offload_guardrail_models=config.offload_guardrail_models,
+            guidance=config.guidance,
+            num_steps=config.num_steps,
+            height=config.height,
+            width=config.width,
+            fps=config.fps,
+            num_video_frames=config.num_video_frames,
+            seed=config.seed,
+        )
 
     def forward(self, prompt):
         cfg = self.config