app changes

.gitignore

@@ -0,0 +1 @@
+__pycache__/

app.py

@@ -40,38 +40,26 @@ mapper_dict = {
     'wavy':'./pretrained_models/styleCLIP_mappers/wavy_hairstyle.pt'
 }
 
+mapper_descs = {
+    'afro':'A face with an afro',
+    'bob':'A face with a bob-cut hairstyle',
+    'bowl':'A face with a bowl cut hairstyle',
+    'buzz':'A face with a buzz cut hairstyle',
+    'caesar':'A face with a caesar cut hairstyle',
+    'crew':'A face with a crew cut hairstyle',
+    'pixie':'A face with a pixie cut hairstyle',
+    'straight':'A face with a straight hair hairstyle',
+    'undercut':'A face with a undercut hairstyle',
+    'wavy':'A face with a wavy hair hairstyle',
+}
+
+
 predictor = dlib.shape_predictor("./pretrained_models/hyperstyle/shape_predictor_68_face_landmarks.dat")
 hyperstyle, hyperstyle_args = load_model(opts.hyperstyle_checkpoint_path, update_opts=opts)
 resize_amount = (256, 256) if hyperstyle_args.resize_outputs else (hyperstyle_args.output_size, hyperstyle_args.output_size)
 im2tensor_transforms = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
 direction_calculator = load_direction_calculator(opts)
 
-ckpt = torch.load(mapper_dict['afro'], map_location='cpu')
-opts.checkpoint_path = mapper_dict['afro']
-mapper_args = ckpt['opts']
-mapper_args.update(vars(opts))
-mapper_args = Namespace(**mapper_args)
-mapper = StyleCLIPMapper(mapper_args)
-mapper.eval()
-mapper.cuda()
-
-def change_mapper(desc):
-    global mapper
-    global mapper_args
-    mapper = None
-    ckpt = None
-    mapper_args = None
-    torch.cuda.empty_cache()
-    opts.checkpoint_path = mapper_dict[desc]
-    ckpt = torch.load(mapper_dict[desc], map_location='cpu')
-    mapper_args = ckpt['opts']
-    mapper_args.update(vars(opts))
-    mapper_args = Namespace(**mapper_args)
-    mapper = StyleCLIPMapper(mapper_args)
-    mapper.eval()
-    mapper.cuda()
-
-
 
 with gr.Blocks() as demo:    
     with gr.Row() as row:
@@ -79,17 +67,17 @@ with gr.Blocks() as demo:
             source = gr.Image(label="Image to Map", type='filepath')
             align = gr.Checkbox(True, label='Align Image')
             inverter_bools = gr.CheckboxGroup(["Hyperstyle", "E4E"], value=['Hyperstyle'], label='Inverter Choices')
-            n_hyperstyle_iterations = gr.Number(3, label='Number of Iterations For Hyperstyle', precision=0)
+            n_hyperstyle_iterations = gr.Number(5, label='Number of Iterations For Hyperstyle', precision=0)
             with gr.Box():
                 mapper_bool = gr.Checkbox(True, label='Output Mapper Result')
                 with gr.Box() as mapper_opts:
-                    mapper_choice = gr.Dropdown(['afro', 'bob', 'bowl', 'buzz', 'caesar', 'crew', 'pixie', 'straight', 'undercut', 'wavy'], value='afro', label='What Hairstyle Mapper to Use?')
-                    mapper_alpha = gr.Slider(minimum=-0.5, maximum=0.5, value=0.01, step=0.1, label='Strength of Mapper Alpha',)
+                    mapper_choice = gr.Dropdown(list(mapper_dict.keys()), value='afro', label='What Hairstyle Mapper to Use?')
+                    mapper_alpha = gr.Slider(minimum=-0.5, maximum=0.5, value=0.1, step=0.01, label='Strength of Mapper Alpha',)
             with gr.Box():
                 gd_bool = gr.Checkbox(False, label='Output Global Direction Result')
                 with gr.Box(visible=False) as gd_opts:
                     neutral_text = gr.Text(value='A face with hair', label='Neutral Text')
-                    target_text = gr.Text(value=mapper_args.description, label='Target Text')
+                    target_text = gr.Text(value=mapper_descs['afro'], label='Target Text')
                     alpha = gr.Slider(minimum=-10.0, maximum=10.0, value=4.1, step=0.1, label="Alpha for Global Direction")
                     beta = gr.Slider(minimum=0.0, maximum=0.30, value=0.15, step=0.01, label="Beta for Global Direction")
             submit_button = gr.Button("Edit Image")
@@ -100,9 +88,13 @@ with gr.Blocks() as demo:
             with gr.Row(visible=False) as e4e_images:
                 output_e4e_mapper = gr.Image(type='pil', label="E4E Mapper")
                 output_e4e_gd = gr.Image(type='pil', label="E4E Global Directions", visible=False)
+    def n_iter_change(number):
+        if number < 0:
+            return 0
+        else:
+            return number
     def mapper_change(new_mapper):
-        change_mapper(new_mapper)
-        return mapper_args.description
+        return mapper_descs[new_mapper]
     def inverter_toggles(bools):
         e4e_bool = 'E4E' in bools
         hyperstyle_bool = 'Hyperstyle' in bools
@@ -125,11 +117,12 @@ with gr.Blocks() as demo:
             output_e4e_gd: gr.update(visible=bool)
             }
     
+    n_hyperstyle_iterations.change(n_iter_change, n_hyperstyle_iterations, n_hyperstyle_iterations)
     mapper_choice.change(mapper_change, mapper_choice, [target_text])
     inverter_bools.change(inverter_toggles, inverter_bools, [hyperstyle_images, e4e_images, n_hyperstyle_iterations])
     mapper_bool.change(mapper_toggles, mapper_bool, [mapper_opts, output_hyperstyle_mapper, output_e4e_mapper])
     gd_bool.change(gd_toggles, gd_bool, [gd_opts, output_hyperstyle_gd, output_e4e_gd])
-    def map_latent(inputs, stylespace=False, weight_deltas=None, strength=0.1):
+    def map_latent(mapper, inputs, stylespace=False, weight_deltas=None, strength=0.1):
         w = inputs.cuda()
         with torch.no_grad():
             if stylespace:
@@ -150,6 +143,14 @@ with gr.Blocks() as demo:
         gd_bool, neutral_text, target_text, alpha, beta,
         ):
         torch.cuda.empty_cache()
+        opts.checkpoint_path = mapper_dict[mapper_choice]
+        ckpt = torch.load(mapper_dict[mapper_choice], map_location='cpu')
+        mapper_args = ckpt['opts']
+        mapper_args.update(vars(opts))
+        mapper_args = Namespace(**mapper_args)
+        mapper = StyleCLIPMapper(mapper_args)
+        mapper.eval()
+        mapper.cuda()
         with torch.no_grad():
             output_imgs = []
             if align_img:
@@ -167,7 +168,7 @@ with gr.Blocks() as demo:
             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 mapper_bool:
-                    mapped_hyperstyle, _ = map_latent(hyperstyle_latents, stylespace=False, weight_deltas=hyperstyle_deltas, strength=mapper_alpha)
+                    mapped_hyperstyle, _ = map_latent(mapper, hyperstyle_latents, stylespace=False, weight_deltas=hyperstyle_deltas, strength=mapper_alpha)
                     mapped_hyperstyle = tensor2im(mapped_hyperstyle[0])
                 else:
                     mapped_hyperstyle = None
@@ -186,7 +187,7 @@ with gr.Blocks() as demo:
                 e4e_batch, e4e_latents = hyperstyle.w_invert(input_img.unsqueeze(0))
                 e4e_deltas = None
                 if mapper_bool:
-                    mapped_e4e, _ = map_latent(e4e_latents, stylespace=False, weight_deltas=e4e_deltas, strength=mapper_alpha)
+                    mapped_e4e, _ = map_latent(mapper, e4e_latents, stylespace=False, weight_deltas=e4e_deltas, strength=mapper_alpha)
                     mapped_e4e = tensor2im(mapped_e4e[0])
                 else:
                     mapped_e4e = None