From c764e9b33bd4992417ea13ce18e4c9588513f4d5 Mon Sep 17 00:00:00 2001 From: SAAD EJAZ Date: Wed, 17 Jun 2026 21:49:17 +0200 Subject: [PATCH] [fix] - bugs in shape decoration + normalization should be along channel dimensions + bg_mask is the inverse of inp_mask --- src/features/decorate_shape.py | 6 +++--- src/init/features.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/features/decorate_shape.py b/src/features/decorate_shape.py index f9aec85..7240683 100644 --- a/src/features/decorate_shape.py +++ b/src/features/decorate_shape.py @@ -71,7 +71,7 @@ def extract_controlnet_dino_features( aligned_unet = F.grid_sample(ft, grid.float(), mode=extr_mode, align_corners=False).reshape( 1, ft_dim, -1 ) - aligned_unet = F.normalize(aligned_unet, dim=-1) + aligned_unet = F.normalize(aligned_unet, dim=1) if dino_from_original: orig_dino_features = get_dino_features( @@ -80,14 +80,14 @@ def extract_controlnet_dino_features( orig_dino_features = F.grid_sample( orig_dino_features, grid, mode=extr_mode, align_corners=False ).reshape(1, 768, -1) - aligned_dino_features = F.normalize(orig_dino_features, dim=-1) + aligned_dino_features = F.normalize(orig_dino_features, dim=1) else: # ! 2. DINOv2 on ControlNet-textured image aligned_dino_features = get_dino_features(device, dino_model, diffusion_output[1][0]) aligned_dino_features = F.grid_sample( aligned_dino_features, grid, mode=extr_mode, align_corners=False ).reshape(1, 768, -1) - aligned_dino_features = F.normalize(aligned_dino_features, dim=-1) + aligned_dino_features = F.normalize(aligned_dino_features, dim=1) # ! 3. Fuse features aligned_fused_features = torch.hstack( diff --git a/src/init/features.py b/src/init/features.py index 570c1c0..017d4fe 100644 --- a/src/init/features.py +++ b/src/init/features.py @@ -87,7 +87,7 @@ def decorate_2d( inp_normals = normals_from_depth_sobel( inp_depth.cpu().numpy(), - bg_mask=inp_mask.squeeze().cpu().numpy(), + bg_mask=~inp_mask.squeeze().cpu().numpy(), ) # * Mask out background