Added tensor dim expansion for edit directions.

styleclip/styleclip_global.py

@@ -89,6 +89,7 @@ imagenet_templates = [
     'a tattoo of the {}.',
 ]
 
+CONV_CODE_INDICES = [(0, 512), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)]
 FFHQ_CODE_INDICES = [(0, 512), (512, 1024), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)] + \
                     [(2048, 2560), (3584, 4096), (5120, 5632), (6656, 7168), (7936, 8192), (8576, 8704), (8896, 8960), (9056, 9088)]
 
@@ -107,6 +108,16 @@ def zeroshot_classifier(model, classnames, templates, device):
         zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
     return zeroshot_weights
 
+def expand_to_full_dim(partial_tensor):
+    full_dim_tensor = torch.zeros(size=(1, 9088))
+
+    start_idx = 0
+    for conv_start, conv_end in CONV_CODE_INDICES:
+        length = conv_end - conv_start
+        full_dim_tensor[:, conv_start:conv_end] = partial_tensor[:, start_idx:start_idx + length]
+        start_idx += length
+
+    return full_dim_tensor
 
 def get_direction(neutral_class, target_class, beta, di, clip_model=None):
 
@@ -157,6 +168,8 @@ def style_dict_to_style_tensor(style_dict, reference_generator):
 def project_code_with_styleclip(source_latent, source_class, target_class, alpha, beta, reference_generator, di, clip_model=None):
         edit_direction = get_direction(source_class, target_class, beta, di, clip_model)
 
+        edit_full_dim = expand_to_full_dim(edit_direction)
+
         source_s = style_dict_to_style_tensor(source_latent, reference_generator)
 
-        return source_s + alpha * edit_direction
+        return source_s + alpha * edit_full_dim