Update app.py

app.py

@@ -12,12 +12,29 @@ MODELS_DIR = Path("models")
 INPUT_DIR.mkdir(exist_ok=True)
 OUTPUT_DIR.mkdir(exist_ok=True)
 
-# Load pre-trained models
+# Function to load pre-trained models
 def load_model(model_path, use_cpu=False):
-    # Explicitly set weights_only=False to load the full model object
-    model = torch.load(model_path, map_location="cpu" if use_cpu else None, weights_only=False)
-    if not use_cpu and torch.cuda.is_available():
-        model = model.cuda()
+    if not model_path.exists():
+        raise FileNotFoundError(f"Model file not found: {model_path}")
+
+    device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"
+
+    # Load state_dict if the model was saved that way
+    model_state = torch.load(model_path, map_location=device)
+
+    # If a full model object was saved, load it directly
+    if isinstance(model_state, torch.nn.Module):
+        model = model_state
+    else:
+        # If saved as state_dict, we need a model architecture (Assuming CNN or custom model)
+        model = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 64, kernel_size=3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 3, kernel_size=3, padding=1)
+        )
+        model.load_state_dict(model_state)
+
+    model.to(device)
     model.eval()
     return model
 
@@ -33,36 +50,24 @@ def process_image(input_path, tile_size=512, seamless=False, use_cpu=False):
 
     # Convert to tensor
     img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
-    if not use_cpu and torch.cuda.is_available():
-        img_tensor = img_tensor.cuda()
+    img_tensor = img_tensor.unsqueeze(0)  # Add batch dimension
+
+    device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"
+    img_tensor = img_tensor.to(device)
 
     # Generate maps
     with torch.no_grad():
-        # Normal map
-        normal_map = normal_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze()
-        # Franken map (contains Displacement and Roughness)
-        franken_map = franken_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze()
+        normal_map = normal_model(img_tensor).cpu().numpy().squeeze()
+        franken_map = franken_model(img_tensor).cpu().numpy().squeeze()
 
     # Post-process maps
-    # Normal map (RGB)
-    if normal_map.ndim == 3:
-        normal_map = normal_map.transpose(1, 2, 0)
-    else:
-        normal_map = np.stack([normal_map] * 3, axis=-1)  # Convert grayscale to RGB if needed
-    normal_map = (normal_map * 255).clip(0, 255).astype(np.uint8)
-
-    # Franken map: Extract Displacement (red) and Roughness (green)
-    if franken_map.ndim == 3:
-        franken_map = franken_map.transpose(1, 2, 0)
-    # Displacement map (red channel, grayscale)
-    disp_map = franken_map[:, :, 0]  # Red channel
-    disp_map = (disp_map * 255).clip(0, 255).astype(np.uint8)
-    disp_map = np.stack([disp_map] * 3, axis=-1)  # Convert to RGB for Gradio display
-
-    # Roughness map (green channel, grayscale)
-    rough_map = franken_map[:, :, 1]  # Green channel
-    rough_map = (rough_map * 255).clip(0, 255).astype(np.uint8)
-    rough_map = np.stack([rough_map] * 3, axis=-1)  # Convert to RGB for Gradio display
+    normal_map = (normal_map.transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
+    disp_map = (franken_map[0] * 255).clip(0, 255).astype(np.uint8)
+    rough_map = (franken_map[1] * 255).clip(0, 255).astype(np.uint8)
+
+    # Convert grayscale to RGB for Gradio display
+    disp_map = np.stack([disp_map] * 3, axis=-1)
+    rough_map = np.stack([rough_map] * 3, axis=-1)
 
     # Define output paths
     base_name = input_path.stem
@@ -114,4 +119,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()