Update app.py

app.py

@@ -11,7 +11,7 @@ import io
 import zipfile
 import uuid
 import traceback
-from huggingface_hub import snapshot_download, login
+from huggingface_hub import snapshot_download, login, HfFileSystem
 from flask_cors import CORS
 import numpy as np
 import trimesh
@@ -35,6 +35,8 @@ os.makedirs(RESULTS_FOLDER, exist_ok=True)
 os.makedirs(CACHE_DIR, exist_ok=True)
 
 os.environ['HF_HOME'] = CACHE_DIR
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
+os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 
@@ -92,7 +94,7 @@ def remove_background(image_path):
         result = remove(img_data)
         img = Image.open(io.BytesIO(result)).convert("RGBA")
         
-        # Check if image is fully transparent (no object)
+        # Check if image is fully transparent
         img_array = np.array(img)
         if np.all(img_array[:, :, 3] == 0):
             print(f"Warning: Image {image_path} is fully transparent or no object detected")
@@ -107,7 +109,6 @@ def remove_background(image_path):
         raise
 
 def preprocess_image(image_path):
-    # Remove background and add black background
     img = remove_background(image_path)
     if img is None:
         raise ValueError("Image is fully transparent or no object detected")
@@ -157,26 +158,32 @@ def load_models():
             print("Warning: HF_TOKEN not found in environment")
         
         dpt_model_name = "Intel/dpt-large"
-        max_retries = 3
-        retry_delay = 5
-        for attempt in range(max_retries):
-            try:
-                print(f"Attempting to download {dpt_model_name}, attempt {attempt+1}")
-                snapshot_download(
-                    repo_id=dpt_model_name,
-                    cache_dir=CACHE_DIR,
-                    resume_download=True,
-                    token=hf_token
-                )
-                print(f"Successfully downloaded {dpt_model_name}")
-                break
-            except Exception as e:
-                if attempt < max_retries - 1:
-                    print(f"DPT download attempt {attempt+1} failed: {str(e)}. Retrying after {retry_delay}s...")
-                    time.sleep(retry_delay)
-                    retry_delay *= 2
-                else:
-                    raise
+        fs = HfFileSystem(token=hf_token)
+        model_cached = os.path.exists(os.path.join(CACHE_DIR, "hub", "models--Intel--dpt-large"))
+        
+        if not model_cached:
+            max_retries = 3
+            retry_delay = 5
+            for attempt in range(max_retries):
+                try:
+                    print(f"Attempting to download {dpt_model_name}, attempt {attempt+1}")
+                    snapshot_download(
+                        repo_id=dpt_model_name,
+                        cache_dir=CACHE_DIR,
+                        resume_download=True,
+                        token=hf_token
+                    )
+                    print(f"Successfully downloaded {dpt_model_name}")
+                    break
+                except Exception as e:
+                    if attempt < max_retries - 1:
+                        print(f"DPT download attempt {attempt+1} failed: {str(e)}. Retrying after {retry_delay}s...")
+                        time.sleep(retry_delay)
+                        retry_delay *= 2
+                    else:
+                        raise
+        else:
+            print(f"{dpt_model_name} already cached in {CACHE_DIR}")
         
         dpt_estimator = pipeline(
             "depth-estimation",
@@ -188,29 +195,34 @@ def load_models():
         print("DPT-Large loaded")
         gc.collect()
         
-        da_model_name = "depth-anything/Depth-Anything-V2-Tiny-hf"
-        for attempt in range(max_retries):
-            try:
-                print(f"Attempting to download {da_model_name}, attempt {attempt+1}")
-                snapshot_download(
-                    repo_id=da_model_name,
-                    cache_dir=CACHE_DIR,
-                    resume_download=True,
-                    token=hf_token
-                )
-                print(f"Successfully downloaded {da_model_name}")
-                break
-            except Exception as e:
-                if attempt < max_retries - 1:
-                    print(f"Depth Anything download attempt {attempt+1} failed: {str(e)}. Retrying after {retry_delay}s...")
-                    time.sleep(retry_delay)
-                    retry_delay *= 2
-                else:
-                    print(f"Failed to load Depth Anything: {str(e)}. Falling back to DPT-Large only.")
-                    depth_anything_model = None
-                    depth_anything_processor = None
-                    model_loaded = True
-                    return dpt_estimator, None, None
+        da_model_name = "LiheYoung/depth-anything-v2-small"
+        da_model_cached = os.path.exists(os.path.join(CACHE_DIR, "hub", "models--LiheYoung--depth-anything-v2-small"))
+        
+        if not da_model_cached:
+            for attempt in range(max_retries):
+                try:
+                    print(f"Attempting to download {da_model_name}, attempt {attempt+1}")
+                    snapshot_download(
+                        repo_id=da_model_name,
+                        cache_dir=CACHE_DIR,
+                        resume_download=True,
+                        token=hf_token
+                    )
+                    print(f"Successfully downloaded {da_model_name}")
+                    break
+                except Exception as e:
+                    if attempt < max_retries - 1:
+                        print(f"Depth Anything download attempt {attempt+1} failed: {str(e)}. Retrying after {retry_delay}s...")
+                        time.sleep(retry_delay)
+                        retry_delay *= 2
+                    else:
+                        print(f"Failed to load Depth Anything: {str(e)}. Falling back to DPT-Large only.")
+                        depth_anything_model = None
+                        depth_anything_processor = None
+                        model_loaded = True
+                        return dpt_estimator, None, None
+        else:
+            print(f"{da_model_name} already cached in {CACHE_DIR}")
         
         depth_anything_processor = AutoImageProcessor.from_pretrained(
             da_model_name,
@@ -288,8 +300,8 @@ def enhance_depth_map(depth_map, detail_level='medium'):
     else:
         enhanced_depth = gaussian_filter(enhanced_depth, sigma=0.7)
     
-    fused_depth = np.clip(fused_depth, 0, 1)
-    return fused_depth
+    enhanced_depth = np.clip(enhanced_depth, 0, 1)
+    return enhanced_depth
 
 def depth_to_mesh(depth_map, image, resolution=80, detail_level='medium', view_angle=0):
     enhanced_depth = enhance_depth_map(depth_map, detail_level)
@@ -319,7 +331,6 @@ def depth_to_mesh(depth_map, image, resolution=80, detail_level='medium', view_a
     y_grid = (y_grid / h - 0.5) * 2.0
     vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
     
-    # Rotate vertices based on view angle (in radians)
     if view_angle != 0:
         rotation_matrix = trimesh.transformations.rotation_matrix(view_angle, [0, 1, 0])
         vertices = trimesh.transform_points(vertices, rotation_matrix)
@@ -398,11 +409,8 @@ def combine_meshes(meshes):
     
     combined_mesh = trimesh.Trimesh(vertices=combined_vertices, faces=combined_faces)
     
-    # Stitch overlapping vertices
     combined_mesh = combined_mesh.subdivide_to_size(max_edge=0.05)
     combined_mesh = combined_mesh.smoothed(method='laplacian', iterations=2)
-    
-    # Ensure watertight mesh
     combined_mesh.fill_holes()
     combined_mesh.fix_normals()
     
@@ -532,11 +540,9 @@ def convert_image_to_3d():
                     view_angles = {'front': 0, 'back': np.pi, 'left': np.pi/2, 'right': -np.pi/2}
                     with torch.no_grad():
                         for view, image in images.items():
-                            # DPT-Large
                             dpt_result = dpt_model(image)
                             dpt_depth = dpt_result["depth"]
                             
-                            # Depth Anything (if loaded)
                             if da_model and da_processor:
                                 inputs = da_processor(images=image, return_tensors="pt")
                                 inputs = {k: v.to("cpu") for k, v in inputs.items()}
@@ -760,7 +766,7 @@ def index():
             "detail_level": "low, medium, or high",
             "texture_quality": "low, medium, or high"
         },
-        "description": "Creates high-quality 3D models from multiple 2D images (front, back, left, right) using DPT-Large and Depth Anything."
+        "description": "Creates high-quality 3D models from multiple 2D images using DPT-Large and Depth Anything."
     }), 200
 
 if __name__ == '__main__':