Fix checkpoint_name bug

futureworld_hf.py

@@ -112,13 +112,13 @@ class AutoregressiveFutureWorld(PreTrainedModel):
         other_args.pop("config")
         config.update(other_args)
         model_sizes = []
-        if "4B" in config.checkpoint_name:
+        if "4B" in config.ar_model_dir:
             model_sizes.append("4B")
-        if "5B" in config.checkpoint_name:
+        if "5B" in config.ar_model_dir:
             model_sizes.append("5B")
-        if "12B" in config.checkpoint_name:
+        if "12B" in config.ar_model_dir:
             model_sizes.append("12B")
-        if "13B" in config.checkpoint_name:
+        if "13B" in config.ar_model_dir:
             model_sizes.append("13B")
         download_autoregressive(model_sizes, config.checkpoint_dir)
         model = cls(config)

test_hf.py

@@ -1,5 +1,5 @@
 from transformers import AutoModel
-
+'''
 model = AutoModel.from_pretrained("EthanZyh/DiffusionText2WorldGeneration",
                                     offload_text_encoder_model = True,
                                     offload_diffusion_transformer = True,
@@ -8,7 +8,12 @@ model = AutoModel.from_pretrained("EthanZyh/DiffusionText2WorldGeneration",
                                     offload_tokenizer = True,
                                     offload_guardrail_models = True,
                                     trust_remote_code = True)
+'''
+
+#prompt = "Video of humanoid robot running in a field"
+
+#model(prompt)
 
-prompt = "Video of humanoid robot running in a field"
+model = AutoModel.from_pretrained("sheldonl/AutoregressiveFutureWorld",
+                                  trust_remote_code = True)
 
-model(prompt)