Update warp_design_on_dress.py

warp_design_on_dress.py

@@ -27,10 +27,10 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
     dress_tensor = tf(dress_img).unsqueeze(0).to(device)
     design_tensor = tf(design_img).unsqueeze(0).to(device)
     
-    # Create design mask (assuming design covers entire area)
+    # Create design mask
     design_mask = torch.ones_like(design_tensor[:, :1, :, :]).to(device)
 
-    # Prepare agnostic input (dress image)
+    # Prepare agnostic input
     agnostic = pad_to_22_channels(dress_tensor).to(device)
 
     # GMM Processing
@@ -62,19 +62,23 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
     with torch.no_grad():
         # Create simplified feature inputs
         gray = dress_tensor.mean(dim=1, keepdim=True)
-        edges = torch.abs(F.conv2d(
-            gray, 
-            torch.tensor([[[[-1,-1,-1],[-1,8,-1],[-1,-1,-1]]], dtype=torch.float32, device=device),
-            padding=1
-        ))
         
-        # Combine inputs (agnostic + warped design + warped mask + features)
+        # Create edge detection kernel
+        kernel = torch.tensor([
+            [[[-1, -1, -1],
+              [-1,  8, -1],
+              [-1, -1, -1]]], dtype=torch.float32, device=device)
+        
+        # Calculate edges
+        edges = torch.abs(F.conv2d(gray, kernel, padding=1))
+        
+        # Combine inputs
         tom_input = torch.cat([
             dress_tensor,          # 3 channels
-            warped_design,          # 3 channels
-            warped_mask,            # 1 channel
-            gray,                   # 1 channel
-            edges,                  # 1 channel
+            warped_design,         # 3 channels
+            warped_mask,           # 1 channel
+            gray,                  # 1 channel
+            edges,                 # 1 channel
             torch.zeros_like(dress_tensor)[:, :17]  # 17 dummy channels
         ], dim=1)  # Total: 3+3+1+1+1+17 = 26 channels