Update hy3dgen/texgen/pipelines.py

hy3dgen/texgen/pipelines.py

@@ -1,18 +1,17 @@
 # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
 # except for the third-party components listed below.
 # Hunyuan 3D does not impose any additional limitations beyond what is outlined
-# in the repsective licenses of these third-party components.
+# in the repsective licenses of third-party components.
 # Users must comply with all terms and conditions of original licenses of these third-party
 # components and must ensure that the usage of the third party components adheres to
 # all relevant laws and regulations.
 
 # For avoidance of doubts, Hunyuan 3D means the large language models and
-# their software and algorithms, including trained model weights, parameters (including
-# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
+# their software and algorithms, including trained model weights, parameters including
+# optimizer states, machine-learning model code, inference-enabling code, training-enabling code,
 # fine-tuning enabling code and other elements of the foregoing made publicly available
 # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
 
-
 import logging
 import numpy as np
 import os
@@ -49,35 +48,34 @@ class Hunyuan3DTexGenConfig:
 class Hunyuan3DPaintPipeline:
     @classmethod
     def from_pretrained(cls, model_path):
+        from huggingface_hub import snapshot_download
+
         original_model_path = model_path
+
         if not os.path.exists(model_path):
-            # try local path
             base_dir = os.environ.get('HY3DGEN_MODELS', '~/.cache/hy3dgen')
-            model_path = os.path.expanduser(os.path.join(base_dir, model_path))
-
+            model_path = os.path.expanduser(os.path.join(base_dir, original_model_path))
+
+        delight_model_path = os.path.join(model_path, 'hunyuan3d-delight-v2-0')
+        multiview_model_path = os.path.join(model_path, 'hunyuan3d-paint-v2-0')
+
+        if not os.path.exists(delight_model_path) or not os.path.exists(multiview_model_path):
+            model_path = snapshot_download(
+                repo_id=original_model_path,
+                allow_patterns=[
+                    "hunyuan3d-delight-v2-0/*",
+                    "hunyuan3d-paint-v2-0/*",
+                ],
+            )
             delight_model_path = os.path.join(model_path, 'hunyuan3d-delight-v2-0')
             multiview_model_path = os.path.join(model_path, 'hunyuan3d-paint-v2-0')
 
-            if not os.path.exists(delight_model_path) or not os.path.exists(multiview_model_path):
-                try:
-                    import huggingface_hub
-                    # download from huggingface
-                    model_path = huggingface_hub.snapshot_download(repo_id=original_model_path,
-                                                                   allow_patterns=["hunyuan3d-delight-v2-0/*"])
-                    model_path = huggingface_hub.snapshot_download(repo_id=original_model_path,
-                                                                   allow_patterns=["hunyuan3d-paint-v2-0/*"])
-                    delight_model_path = os.path.join(model_path, 'hunyuan3d-delight-v2-0')
-                    multiview_model_path = os.path.join(model_path, 'hunyuan3d-paint-v2-0')
-                    return cls(Hunyuan3DTexGenConfig(delight_model_path, multiview_model_path))
-                except ImportError:
-                    logger.warning(
-                        "You need to install HuggingFace Hub to load models from the hub."
-                    )
-                    raise RuntimeError(f"Model path {model_path} not found")
-            else:
-                return cls(Hunyuan3DTexGenConfig(delight_model_path, multiview_model_path))
-
-        raise FileNotFoundError(f"Model path {original_model_path} not found and we could not find it at huggingface")
+        if not os.path.exists(delight_model_path):
+            raise FileNotFoundError(f"Texture delight model folder not found: {delight_model_path}")
+        if not os.path.exists(multiview_model_path):
+            raise FileNotFoundError(f"Texture paint model folder not found: {multiview_model_path}")
+
+        return cls(Hunyuan3DTexGenConfig(delight_model_path, multiview_model_path))
 
     def __init__(self, config):
         self.config = config