Got RIS and Metrics up and running

app.py

@@ -2,6 +2,8 @@ import sys
 import os
 import torch
 
+from metrics.metrics import ClipHair
+
 sys.path.append(".")
 
 from gradio_wrapper.gradio_options import GradioTestOptions
@@ -185,6 +187,9 @@ with gr.Blocks() as demo:
             result_batch = (x_hat, w_hat)
         return result_batch
     def run_metrics(base_img, edited_img):
+        #print(base_img.shape, edited_img.shape)
+        #base_img = base_img.unsqueeze(0)
+        #edited_img = edited_img.unqueeze(0)
         lpips_score = lpips_metric(base_img, edited_img)[0]
         ssim_score = ssim_metric(base_img, edited_img)[0]
         id_score = id_metric(base_img, edited_img)[0]
@@ -234,35 +239,36 @@ with gr.Blocks() as demo:
             hyperstyle_metrics_text = ''
             if 'Hyperstyle' in inverter_bools:
                 hyperstyle_batch, hyperstyle_latents, hyperstyle_deltas, _ = run_inversion(input_img.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)
-                if invert_bool:
-                    invert_hyperstyle = tensor2im(hyperstyle_batch[0])
-                else:
-                    invert_hyperstyle = None
+                invert_hyperstyle = tensor2im(hyperstyle_batch[0])
                 if mapper_bool:
                     mapped_hyperstyle, _ = map_latent(mapper, hyperstyle_latents, stylespace=False, weight_deltas=hyperstyle_deltas, strength=mapper_alpha)
-                    #clip_score = clip_text_metric(mapped_hyperstyle[0], mapper_args.description)
+                    clip_score = clip_text_metric(mapped_hyperstyle[0], mapper_args.description)
                     mapped_hyperstyle = tensor2im(mapped_hyperstyle[0])
-                    #lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), mapped_hyperstyle.resize(resize_to))
-                    #hyperstyle_metrics_text += f'Mapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
+                    lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), mapped_hyperstyle.resize(resize_to))
+                    hyperstyle_metrics_text += f'\nMapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
                 else:
                     mapped_hyperstyle = None
                 
                 if gd_bool:
                     gd_hyperstyle = edit_image(_, hyperstyle_latents[0], hyperstyle.decoder, direction_calculator, opts, hyperstyle_deltas)
-                    #clip_score = clip_text_metric(gd_hyperstyle[0], opts.target_text)
+                    clip_score = clip_text_metric(gd_hyperstyle[0], opts.target_text)
                     gd_hyperstyle = tensor2im(gd_hyperstyle[0])
-                    #lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), gd_hyperstyle.resize(resize_to))
-                    #hyperstyle_metrics_text += f'Global Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
+                    lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), gd_hyperstyle.resize(resize_to))
+                    hyperstyle_metrics_text += f'\nGlobal Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
                 else:
                     gd_hyperstyle = None
                 
                 if ris_bool:
 
                     ref_hyperstyle_batch, ref_hyperstyle_latents, ref_hyperstyle_deltas, _ = run_inversion(ref_input.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)
-                    blend_hyperstyle, blend_hyperstyle_latents = blend_latents(hyperstyle_latents, ref_hyperstyle_batch,
+                    blend_hyperstyle, blend_hyperstyle_latents = blend_latents(hyperstyle_latents, ref_hyperstyle_latents,
                                                     src_deltas=hyperstyle_deltas, ref_deltas=ref_hyperstyle_deltas,
                                                     generator=ris_gen, device=device)
-                    ris_hyperstyle = tensor2im(blend_hyperstyle)
+                    ris_hyperstyle = tensor2im(blend_hyperstyle[0])
+
+                    lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))
+                    clip_score = clip_hair(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))[1]
+                    hyperstyle_metrics_text += f'\nRIS Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Hair Score: \t{clip_score}'
                 else:
                     ris_hyperstyle=None
 
@@ -274,16 +280,13 @@ with gr.Blocks() as demo:
             if 'E4E' in inverter_bools:
                 e4e_batch, e4e_latents = hyperstyle.w_invert(input_img.unsqueeze(0))
                 e4e_deltas = None
-                if invert_bool:
-                    invert_e4e = tensor2im(e4e_batch[0])
-                else:
-                    invert_e4e = None
+                invert_e4e = tensor2im(e4e_batch[0])
                 if mapper_bool:
                     mapped_e4e, _ = map_latent(mapper, e4e_latents, stylespace=False, weight_deltas=e4e_deltas, strength=mapper_alpha)
-                    #clip_score = clip_text_metric(mapped_e4e[0], mapper_args.description)
+                    clip_score = clip_text_metric(mapped_e4e[0], mapper_args.description)
                     mapped_e4e = tensor2im(mapped_e4e[0])
-                    #lpips_score, ssim_score, id_score = run_metrics(invert_e4e, mapped_e4e)
-                    #e4e_metrics_text += f'Mapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
+                    lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), mapped_e4e.resize(resize_to))
+                    e4e_metrics_text += f'\nMapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
                 
                 else:
                     mapped_e4e = None
@@ -292,8 +295,8 @@ with gr.Blocks() as demo:
                     gd_e4e = edit_image(_, e4e_latents[0], hyperstyle.decoder, direction_calculator, opts, e4e_deltas)
                     clip_score = clip_text_metric(gd_e4e[0], opts.target_text)
                     gd_e4e = tensor2im(gd_e4e[0])
-                    lpips_score, ssim_score, id_score = run_metrics(invert_e4e, gd_e4e)
-                    e4e_metrics_text += f'Global Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
+                    lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), gd_e4e.resize(resize_to))
+                    e4e_metrics_text += f'\nGlobal Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'
                 
                 else:
                     gd_e4e = None
@@ -301,10 +304,14 @@ with gr.Blocks() as demo:
                 if ris_bool:
                     ref_e4e_batch, ref_e4e_latents, = hyperstyle.w_invert(ref_input.unsqueeze(0))
                     ref_e4e_deltas= None
-                    blend_e4e, blend_e4e_latents = blend_latents(e4e_latents, ref_e4e_batch,
+                    blend_e4e, blend_e4e_latents = blend_latents(e4e_latents, ref_e4e_latents,
                                                     src_deltas=None, ref_deltas=None,
                                                     generator=ris_gen, device=device)
-                    ris_e4e = tensor2im(blend_e4e)
+                    ris_e4e = tensor2im(blend_e4e[0])
+
+                    lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))
+                    clip_score = clip_hair(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))[1]
+                    e4e_metrics_text += f'\nRIS Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Hair Score: \t{clip_score}'
                 else:
                     ris_e4e=None
                 

ris/blend.py

@@ -111,7 +111,9 @@ def compute_M(w, generator, weights_deltas=None, device='cuda'):
 	return M
 
 def blend_latents (source_latent, ref_latent, generator, src_deltas=None, ref_deltas=None, device='cuda'):
+    #print(source_latent.shape)
     source = generator.get_latent(source_latent[0].unsqueeze(0), truncation=1, is_latent=True)
+    #print(ref_latent.shape)
     ref = generator.get_latent(ref_latent[0].unsqueeze(0), truncation=1, is_latent=True)
     source_M = compute_M(source, generator, weights_deltas=src_deltas, device='cpu')
     ref_M = compute_M(ref, generator, weights_deltas=ref_deltas, device='cpu')
@@ -127,5 +129,6 @@ def blend_latents (source_latent, ref_latent, generator, src_deltas=None, ref_de
     
     blend = style2list((add_direction(source, ref, part_M, 1.3)))
     blend_out, _ = generator(blend, weights_deltas=blend_deltas)
+    #print(blend_out.shape)
 
     return blend_out, blend

ris/model.py

@@ -160,13 +160,15 @@ class EqualLinear(nn.Module):
         self.lr_mul = lr_mul
 
     def forward(self, input):
+        weight = self.weight * self.scale
+        bias = self.bias * self.lr_mul
         if self.activation:
-            out = F.linear(input, self.weight * self.scale)
-            out = fused_leaky_relu(out, self.bias * self.lr_mul)
+            out = F.linear(input, weight)
+            out = fused_leaky_relu(out, bias)
 
         else:
             out = F.linear(
-                input, self.weight * self.scale, bias=self.bias * self.lr_mul
+                input, weight, bias=bias
             )
 
         return out