update

src/model_processing.py

@@ -11,7 +11,7 @@ def download_ckpt_yaml(model_path, model_name, ckpt_path, yaml_url=None):
         with open(save_path, 'wb') as f:
             f.write(response.content)
 
-    # os.makedirs(model_path, exist_ok=True)
+    os.makedirs(model_path, exist_ok=True)
     local_dir = os.path.join(model_path, model_name)
     os.makedirs(local_dir, exist_ok=True)
 
@@ -34,6 +34,7 @@ def download_ckpt_yaml(model_path, model_name, ckpt_path, yaml_url=None):
 
 def get_model(model_path, model_name):
     model = None
+    huggingface_token = os.environ.get("HUGGINGFACE_TOKEN", None)
     data_params = {
         "target_image_size": (512, 512),
         "lock_ratio": True,
@@ -253,7 +254,10 @@ def get_model(model_path, model_name):
         else:
             raise Exception(f"Unsupported model: {model_name}")
 
-        model = TiTok.from_pretrained(ckpt_path)
+        model = TiTok.from_pretrained(
+            ckpt_path,
+            token=huggingface_token
+        )
         data_params = {
             "target_image_size": (256, 256),
             "lock_ratio": True,
@@ -268,7 +272,8 @@ def get_model(model_path, model_name):
         import types
 
         model = MultiModalityCausalLM.from_pretrained(
-            "deepseek-ai/Janus-Pro-7B", trust_remote_code=True
+            "deepseek-ai/Janus-Pro-7B", trust_remote_code=True,
+            token=huggingface_token
         ).gen_vision_model
         model.forward = types.MethodType(forward, model)
         data_params = {
@@ -355,7 +360,8 @@ def get_model(model_path, model_name):
         import types
 
         model = AutoencoderKL.from_pretrained(
-            "huaweilin/stable-diffusion-3.5-large-vae", subfolder="vae"
+            "huaweilin/stable-diffusion-3.5-large-vae", subfolder="vae",
+            token=huggingface_token
         )
         model.forward = types.MethodType(forward, model)
         data_params = {
@@ -372,7 +378,8 @@ def get_model(model_path, model_name):
         import types
 
         model = AutoencoderKL.from_pretrained(
-            "black-forest-labs/FLUX.1-dev", subfolder="vae"
+            "black-forest-labs/FLUX.1-dev", subfolder="vae",
+            token=huggingface_token
         )
         model.forward = types.MethodType(forward, model)
         data_params = {