add some comments

text2world_hf.py

@@ -3,12 +3,15 @@ import argparse
 import torch
 from transformers import PreTrainedModel, PretrainedConfig
 
+# TODO: This is a bug to fix. Huggingface cannot download .cosmos1.models.diffusion.inference.inference_utils because it's in a subfolder.
 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
 
+
+# custom config class
 class DiffusionText2WorldConfig(PretrainedConfig):
     model_type = "DiffusionText2World" 
     def __init__(self, **kwargs):
@@ -38,6 +41,7 @@ class DiffusionText2WorldConfig(PretrainedConfig):
         self.offload_guardrail_models = kwargs.get("offload_guardrail_models", False)
 
 
+# custom model calss
 class DiffusionText2World(PreTrainedModel):
     config_class = DiffusionText2WorldConfig
 
@@ -69,6 +73,7 @@ class DiffusionText2World(PreTrainedModel):
             seed=config.seed,
         )
 
+    # modifed from text2world.py demo function
     def forward(self, prompt):
         cfg = self.config
         # Handle multiple prompts if prompt file is provided
@@ -118,7 +123,7 @@ class DiffusionText2World(PreTrainedModel):
             log.info(f"Saved prompt to {prompt_save_path}")
     
     def save_pretrained(self, save_directory, **kwargs): 
-        # We don't save anything
+        # We don't save anything, but need this function to override
         pass
 
     @classmethod