pp

app.py

@@ -15,16 +15,21 @@ from beta_regressor import predict_betas
 from smpl_generator import generate_mesh
 from renderer import render_avatar
 
-SMPL_MODEL_PATH = os.getenv("SMPL_MODEL_PATH", "smpl/smpl")
-
-smpl_models_dir = Path("smpl/smpl/models")
-if smpl_models_dir.exists():
-    print(f"Found SMPL models in {smpl_models_dir}")
-    model_files = list(smpl_models_dir.glob("*.pkl"))
-    for f in model_files:
-        print(f"  - {f.name}")
-else:
-    print(f"Warning: SMPL models not found in {smpl_models_dir}")
+SMPL_MODEL_PATH = os.getenv("SMPL_MODEL_PATH", "smpl")
+
+smpl_models_dirs = [Path("smpl/smpl/models"), Path("smpl/models"), Path(SMPL_MODEL_PATH) / "models"]
+found_models = False
+for models_dir in smpl_models_dirs:
+    if models_dir.exists():
+        print(f"Found SMPL models in {models_dir}")
+        model_files = list(models_dir.glob("*.pkl"))
+        for f in model_files:
+            print(f"  - {f.name}")
+        found_models = True
+        break
+
+if not found_models:
+    print(f"Warning: SMPL models not found in expected locations")
     print(f"Looking in {SMPL_MODEL_PATH}...")
 
 app = FastAPI(

smpl_generator.py

@@ -7,20 +7,31 @@ import os
 
 
 class SMPLGenerator:
-    def __init__(self, model_path: str = "smpl/smpl", gender: str = "neutral", device: str = "cpu"):
+    def __init__(self, model_path: str = "smpl", gender: str = "neutral", device: str = "cpu"):
         self.device = torch.device(device)
         self.gender = gender
         
         model_path_obj = Path(model_path)
         
         if not model_path_obj.exists():
-            alt_path = Path("smpl/smpl")
-            if alt_path.exists():
-                model_path_obj = alt_path
-                print(f"Using alternative model path: {model_path_obj}")
+            alt_paths = [Path("smpl"), Path("smpl/smpl")]
+            for alt_path in alt_paths:
+                if alt_path.exists():
+                    model_path_obj = alt_path
+                    print(f"Using alternative model path: {model_path_obj}")
+                    break
             else:
                 model_path_obj.mkdir(parents=True, exist_ok=True)
         
+        models_source = model_path_obj / "smpl" / "models"
+        if not models_source.exists():
+            models_source = model_path_obj / "models"
+        
+        if models_source.exists() and (model_path_obj / "smpl").exists():
+            expected_models_dir = model_path_obj / "smpl" / "models"
+            if not expected_models_dir.exists() and models_source.exists():
+                print(f"Models found in {models_source}, smplx may need them in {expected_models_dir}")
+        
         self.model_path = model_path_obj
         model_path_str = str(self.model_path)
         
@@ -28,37 +39,72 @@ class SMPLGenerator:
             gender = "male"
             print("Note: Neutral gender not available, using male model")
         
-        try:
-            self.smpl_model = smplx.create(
-                model_path=model_path_str,
-                model_type='smpl',
-                gender=gender,
-                batch_size=1,
-                ext='npz'
-            ).to(self.device)
-        except Exception as e:
+        models_dir = model_path_obj / "models"
+        if not models_dir.exists():
+            models_dir = model_path_obj / "smpl" / "models"
+        
+        print(f"Looking for SMPL models in: {model_path_str}")
+        if models_dir.exists():
+            model_files = list(models_dir.glob("*.pkl"))
+            print(f"Found {len(model_files)} model files: {[f.name for f in model_files]}")
+        
+        model_paths_to_try = [
+            "smpl",
+            model_path_str,
+            str(model_path_obj / "smpl"),
+            str(models_dir.parent) if models_dir.exists() else None,
+            "smpl/smpl"
+        ]
+        model_paths_to_try = [p for p in model_paths_to_try if p is not None and Path(p).exists()]
+        
+        last_error = None
+        for try_path in model_paths_to_try:
+            print(f"Trying model path: {try_path}")
             try:
                 self.smpl_model = smplx.create(
-                    model_path=model_path_str,
+                    model_path=try_path,
                     model_type='smpl',
                     gender=gender,
                     batch_size=1,
-                    ext='pkl'
+                    ext='npz'
                 ).to(self.device)
-            except Exception as e2:
+                print(f"Successfully loaded model from: {try_path}")
+                break
+            except Exception as e:
+                last_error = e
                 try:
                     self.smpl_model = smplx.create(
-                        model_path=model_path_str,
+                        model_path=try_path,
                         model_type='smpl',
                         gender=gender,
-                        batch_size=1
+                        batch_size=1,
+                        ext='pkl'
                     ).to(self.device)
-                except Exception as e3:
-                    raise RuntimeError(
-                        f"Failed to load SMPL model. Error: {str(e3)}. "
-                        f"Model path: {model_path_str}. "
-                        f"Please ensure SMPL model files are in {model_path_str}/models/ directory."
-                    )
+                    print(f"Successfully loaded model from: {try_path}")
+                    break
+                except Exception as e2:
+                    last_error = e2
+                    try:
+                        self.smpl_model = smplx.create(
+                            model_path=try_path,
+                            model_type='smpl',
+                            gender=gender,
+                            batch_size=1
+                        ).to(self.device)
+                        print(f"Successfully loaded model from: {try_path}")
+                        break
+                    except Exception as e3:
+                        last_error = e3
+                        continue
+        else:
+            error_msg = str(last_error) if last_error else "Unknown error"
+            print(f"Error details: {error_msg}")
+            raise RuntimeError(
+                f"Failed to load SMPL model after trying paths: {model_paths_to_try}. "
+                f"Error: {error_msg}. "
+                f"Models should be in a 'models' subdirectory. "
+                f"Expected files: basicModel_f_lbs_*.pkl (female) or basicmodel_m_lbs_*.pkl (male)"
+            )
     
     def generate_mesh(
         self,
@@ -107,7 +153,7 @@ class SMPLGenerator:
 _generator_instance = None
 
 
-def get_generator(model_path: str = "smpl/smpl", gender: str = "neutral", device: str = "cpu") -> SMPLGenerator:
+def get_generator(model_path: str = "smpl", gender: str = "neutral", device: str = "cpu") -> SMPLGenerator:
     global _generator_instance
     if _generator_instance is None:
         _generator_instance = SMPLGenerator(model_path=model_path, gender=gender, device=device)
@@ -116,7 +162,7 @@ def get_generator(model_path: str = "smpl/smpl", gender: str = "neutral", device
 
 def generate_mesh(
     betas: np.ndarray,
-    model_path: str = "smpl/smpl",
+    model_path: str = "smpl",
     gender: str = "neutral",
     device: str = "cpu"
 ) -> Tuple[np.ndarray, np.ndarray]: