update app

app.py

@@ -4,11 +4,27 @@ import dnnlib
 import numpy as np
 import torch
 
-from find_direction import find_direction
+import find_direction
+import generator
+import psp_wrapper
+
+
+psp_encoder_path = "./pretrained/e4e_ffhq_encode.pt"
+landmarks_path = "./pretrained/shape_predictor_68_face_landmarks.dat"
+e4e_embedder = psp_wrapper.psp_encoder(psp_encoder_path, landmarks_path)
+G_ffhq_path = "./pretrained/ffhq.pkl"
+G_metfaces_path = "./pretrained/metfaces.pkl"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-with dnnlib.util.open_url("./pretrained/ffhq.pkl") as f:
-    G = legacy.load_network_pkl(f)['G_ema'].to(device) 
+
+with dnnlib.util.open_url(G_ffhq_path) as f:
+    G_ffhq = legacy.load_network_pkl(f)['G_ema'].to(device)
+
+with dnnlib.util.open_url(G_metfaces_path) as f:
+    G_metfaces = legacy.load_network_pkl(f)['G_ema'].to(device)
+
+G_dict = {"FFHQ": G_ffhq, "MetFaces": G_metfaces}
+
 
 
 DESCRIPTION = '''# <a href="https://github.com/catlab-team/stylemc"> StyleMC:</a> Multi-Channel Based Fast Text-Guided Image Generation and Manipulation
@@ -16,171 +32,76 @@ DESCRIPTION = '''# <a href="https://github.com/catlab-team/stylemc"> StyleMC:</a
 FOOTER = 'This space is built by <a href = "https://github.com/catlab-team">Catlab Team</a>.'
 
 
-def main():
-    with gr.Blocks(css='style.css') as demo:
-        gr.Markdown(DESCRIPTION)
-
-        with gr.Box():
-            gr.Markdown('''## Step 1 (Finding a global manipulation direction)
-- Please enter the target **text prompt** and **identity loss weight** to find global manipulation direction:
-- Hit the **Find Direction** button.
-''')
-            with gr.Row():
-                with gr.Column():
-                    with gr.Row():
-                        text = gr.Textbox(
-                            label="Enter your prompt",
-                            show_label=False,
-                            max_lines=1,
-                            placeholder="Enter your prompt",
-                        ).style(
-                            container=False,
-                        )
-                        identity_loss_weight = gr.Slider(0.1,
-                                                10,
-                                                value=0.5,
-                                                step=0.1,
-                                                label='Identity Loss Weight',
-                                                interactive=True)
-                        btn = gr.Button("Find Direction").style(full_width=False)
-
-        with gr.Box():
-            gr.Markdown('''## Step 2 (Manipulation)
-- Please upload an image for manipulation:
-    - You can also select the **previous directions** and determine the **manipulation strength**.
-- Hit the **Generate** button.
-''')
-            with gr.Row():
-                identity_loss_weight = gr.Slider(0.1,
-                                                  100,
-                                                  value=50,
-                                                  step=0.1,
-                                                  label='Manipulation Strength',
-                                                  interactive=True)
-            with gr.Row():
-                with gr.Column():
-                    with gr.Row():
-                        input_image = gr.Image(label='Input Image',
-                                               type='filepath')
-                    with gr.Row():
-                        generate_button = gr.Button('Generate')
-                with gr.Column():
-                    with gr.Row():
-                        generated_image = gr.Image(label='Generated Image',
-                                                type='numpy',
-                                                interactive=False)
-                        
-
-
-
-#         with gr.Box():
-#             gr.Markdown('''## Step 2 (Select Style Image)
-# - Select **Style Type**.
-# - Select **Style Image Index** from the image table below.
-# ''')
-#             with gr.Row():
-#                 with gr.Column():
-#                     style_type = gr.Radio(model.style_types,
-#                                           label='Style Type')
-#                     text = get_style_image_markdown_text('cartoon')
-#                     style_image = gr.Markdown(value=text)
-#                     style_index = gr.Slider(0,
-#                                             316,
-#                                             value=26,
-#                                             step=1,
-#                                             label='Style Image Index')
-
-#             with gr.Row():
-#                 example_styles = gr.Dataset(
-#                     components=[style_type, style_index],
-#                     samples=[
-#                         ['cartoon', 26],
-#                         ['caricature', 65],
-#                         ['arcane', 63],
-#                         ['pixar', 80],
-#                     ])
-
-#         with gr.Box():
-#             gr.Markdown('''## Step 3 (Generate Style Transferred Image)
-# - Adjust **Structure Weight** and **Color Weight**.
-#     - These are weights for the style image, so the larger the value, the closer the resulting image will be to the style image.
-# - Hit the **Generate** button.
-# ''')
-#             with gr.Row():
-#                 with gr.Column():
-#                     with gr.Row():
-#                         structure_weight = gr.Slider(0,
-#                                                      1,
-#                                                      value=0.6,
-#                                                      step=0.1,
-#                                                      label='Structure Weight')
-#                     with gr.Row():
-#                         color_weight = gr.Slider(0,
-#                                                  1,
-#                                                  value=1,
-#                                                  step=0.1,
-#                                                  label='Color Weight')
-#                     with gr.Row():
-#                         structure_only = gr.Checkbox(label='Structure Only')
-#                     with gr.Row():
-#                         generate_button = gr.Button('Generate')
-
-#                 with gr.Column():
-#                     result = gr.Image(label='Result')
-
-#             with gr.Row():
-#                 example_weights = gr.Dataset(
-#                     components=[structure_weight, color_weight],
-#                     samples=[
-#                         [0.6, 1.0],
-#                         [0.3, 1.0],
-#                         [0.0, 1.0],
-#                         [1.0, 0.0],
-#                     ])
-
-        gr.Markdown(FOOTER)
-
-        # preprocess_button.click(fn=model.detect_and_align_face,
-        #                         inputs=input_image,
-        #                         outputs=aligned_face)
-        # aligned_face.change(fn=model.reconstruct_face,
-        #                     inputs=aligned_face,
-        #                     outputs=[
-        #                         reconstructed_face,
-        #                         instyle,
-        #                     ])
-        # style_type.change(fn=update_slider,
-        #                   inputs=style_type,
-        #                   outputs=style_index)
-        # style_type.change(fn=update_style_image,
-        #                   inputs=style_type,
-        #                   outputs=style_image)
-        # generate_button.click(fn=model.generate,
-        #                       inputs=[
-        #                           style_type,
-        #                           style_index,
-        #                           structure_weight,
-        #                           color_weight,
-        #                           structure_only,
-        #                           instyle,
-        #                       ],
-        #                       outputs=result)
-        # example_images.click(fn=set_example_image,
-        #                      inputs=example_images,
-        #                      outputs=example_images.components)
-        # example_styles.click(fn=set_example_styles,
-        #                      inputs=example_styles,
-        #                      outputs=example_styles.components)
-        # example_weights.click(fn=set_example_weights,
-        #                       inputs=example_weights,
-        #                       outputs=example_weights.components)
-
-    demo.launch(
-        # enable_queue=args.enable_queue,
-        # server_port=args.port,
-        # share=args.share,
-    )
+def add_direction(prompt, stylegan_type, id_loss_w):
+    new_dir_name = prompt+" "+stylegan_type+" w_id_loss"+str(id_loss_w)
+    if (prompt != None) and (new_dir_name not in direction_list):
+        print("adding direction with id:", new_dir_name)
+        direction = find_direction.find_direction(G_dict[stylegan_type], prompt)
+        print(f"new direction calculated with {stylegan_type} and id loss weight = {id_loss_w}")
+        direction_list.append(new_dir_name)
+        direction_map[new_dir_name]={"direction":direction, "stylegan_type":stylegan_type}
+
+    return gr.Radio.update(choices=direction_list, value=None, visible=True)
 
 
-if __name__ == '__main__':
-    main()
+def generate_output_image(image_path, direction_id, change_power):
+    direction = direction_map[direction_id]["direction"]
+    G=G_dict["FFHQ"]
+
+    w = e4e_embedder.get_w(image_path) # numpy array
+    s = generator.w_to_s(GIn=G, wsIn=w)
+    output_image = generator.generate_from_style(
+        GIn=G, 
+        styles=s, 
+        styles_direction=direction,
+        change_power=change_power,
+        outdir='.'
+    )
+    return output_image
+  
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Box():
+        gr.Markdown('''## Step 1 (Finding a global manipulation direction) - Please enter the target **text prompt** and **identity loss weight** to find global manipulation direction: - Hit the **Find Direction** button.''')
+        with gr.Row():
+            with gr.Column():
+                style_gan_type = gr.Radio(["FFHQ", "MetFaces"], value = "FFHQ", label="StyleGAN Type", interactive=True)
+            with gr.Column():
+                identity_loss_weight = gr.Slider(
+                    0.1, 10, value=0.5, step=0.1,label="Identity Loss Weight",interactive=True
+                )
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    text = gr.Textbox(
+                        label="Enter your prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt"
+                    ).style(container=False)
+
+                  find_direction_btn = gr.Button("Find Direction").style(full_width=False)
+
+    with gr.Box():
+        gr.Markdown('''## Step 2 (Manipulation) - Please upload an image for manipulation: - You can also select the **previous directions** and determine the **manipulation strength**. - Hit the **Generate** button.''')
+        with gr.Row():
+            direction_radio = gr.Radio(direction_list, label="List of Directions")
+        with gr.Row():
+            manipulation_strength = gr.Slider(
+                0.1, 100, value=25, step=0.1, label="Manipulation Strength",interactive=True
+            )
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    input_image = gr.Image(label="Input Image", type="filepath")
+                with gr.Row():
+                    generate_btn = gr.Button("Generate")
+            with gr.Column():
+                with gr.Row():
+                    generated_image = gr.Image(label="Generated Image",type="pil",interactive=False)
+
+    find_direction_btn.click(add_direction, inputs=[text, style_gan_type, identity_loss_weight], outputs=direction_radio)
+    generate_btn.click(generate_output_image, inputs=[input_image, direction_radio,manipulation_strength], outputs=generated_image)
+
+demo.launch(debug=True)

find_direction.py

@@ -22,6 +22,7 @@ from torch_utils.ops import upfirdn2d
 import id_loss
 from copy import deepcopy
 
+
 def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
         w_iter = iter(ws.unbind(dim=1))

generator.py

@@ -31,7 +31,7 @@ from torch import linalg as LA
 import torch.nn.functional as F
 
 
-def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
+def block_forward(self, x, img, ws, shapes, force_fp32=True, fused_modconv=None, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
         w_iter = iter(ws.unbind(dim=1))
         dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
@@ -74,7 +74,7 @@ def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None
         return x, img
 
 
-def block_forward_from_style(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
+def block_forward_from_style(self, x, img, ws, shapes, force_fp32=True, fused_modconv=None, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
         w_iter = iter(ws.unbind(dim=1))
         dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32

psp_wrapper.py

@@ -31,15 +31,16 @@ EXPERIMENT_ARGS['transform'] = transforms.Compose([
 
 class psp_encoder:
     def __init__(self, model_path: str, shape_predictor_path: str):
-        self.ckpt = torch.load(model_path, map_location='cpu')
-        self.opts = self.ckpt['opts']
+        self.ckpt = torch.load(model_path, map_location="cpu")
+        self.opts = self.ckpt["opts"]
         # update the training options
-        self.opts['checkpoint_path'] = model_path
+        self.opts["checkpoint_path"] = model_path
         self.opts= Namespace(**self.opts)
         self.net = pSp(self.opts)
         self.net.eval()
         self.net.cuda()
         self.shape_predictor = dlib.shape_predictor(shape_predictor_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     def get_w(self, image_path):
         original_image = Image.open(image_path)
@@ -47,8 +48,10 @@ class psp_encoder:
         input_image = align_face(filepath=image_path, predictor=self.shape_predictor)
         resize_dims = (256, 256)
         input_image.resize(resize_dims)
-        img_transforms = EXPERIMENT_ARGS['transform']
+        img_transforms = EXPERIMENT_ARGS["transform"]
         transformed_image = img_transforms(input_image)
+
+
         with torch.no_grad():
-            _, latents = self.net(transformed_image.unsqueeze(0).to("cuda").float(), randomize_noise=False, return_latents=True)
+            _, latents = self.net(transformed_image.unsqueeze(0).to(device).float(), randomize_noise=False, return_latents=True)
             return latents.cpu().numpy()