update app

app.py

@@ -1,29 +1,186 @@
-import cv2
-import torch
-import clip
 import gradio as gr
+import legacy
+import dnnlib
 import numpy as np
+import torch
+
+from find_direction import find_direction
+
+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) 
+
+
+DESCRIPTION = '''# <a href="https://github.com/catlab-team/stylemc"> StyleMC:</a> Multi-Channel Based Fast Text-Guided Image Generation and Manipulation
+'''
+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')
 
-# Use GPU if available
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-else:
-    device = torch.device("cpu")
+#             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)
 
-def manipulate_image(img, text_queries, score_threshold):
-    return img
+        # 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,
+    )
 
-description = """
-"""
 
-demo = gr.Interface(
-    manipulate_image, 
-    inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)], 
-    outputs="image",
-    title="Text-guided image manipulation with StyleMC",
-    description=description,
-)
-demo.launch()
+if __name__ == '__main__':
+    main()

find_direction.py

@@ -72,7 +72,7 @@ def unravel_index(index, shape):
     return tuple(reversed(out))
 
 def find_direction(
-    network_pkl: str,
+    G,
     text_prompt: str,
     truncation_psi: float = 0.7,
     noise_mode: str = "const",
@@ -82,10 +82,7 @@ def find_direction(
     seeds=np.random.randint(0, 1000, 128)
 
     batch_size=1
-    print('Loading networks from "%s"...' % network_pkl)
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    with dnnlib.util.open_url(network_pkl) as f:
-        G = legacy.load_network_pkl(f)['G_ema'].to(device)  # type: ignore
 
     # Labels
     class_idx=None