Update app.py

app.py

@@ -40,6 +40,26 @@ def process_image_from_binary(img_stream):
     
     return img
 
+def execute_prepare(diffusion_model_id, concept, steps, attack_id):
+    print(f"my IP is {myip}, my port is {myport}")
+    print(f"my input is diffusion_model_id: {diffusion_model_id}, concept: {concept}, steps: {steps}")
+    response = requests.post('http://{}:{}/prepare'.format(myip, myport), 
+                             json={"diffusion_model_id": diffusion_model_id, "concept": concept, "steps": steps, "attack_id": attack_id},
+                             timeout=(10, 1200))
+    print(f"result: {response}")
+    # result = result.text[1:-1]
+    prompt = ""
+    img = None
+    if response.status_code == 200:
+        response_json = response.json()
+        print(response_json)
+        prompt = response_json['input_prompt']
+        img = process_image_from_binary(response_json['no_attack_img'])
+    else:
+        print(f"Request failed with status code {response.status_code}")
+    
+    return prompt, img
+
 def excute_udiff(diffusion_model_id, concept, steps, attack_id):
     print(f"my IP is {myip}, my port is {myport}")
     print(f"my input is diffusion_model_id: {diffusion_model_id}, concept: {concept}, steps: {steps}")
@@ -48,21 +68,17 @@ def excute_udiff(diffusion_model_id, concept, steps, attack_id):
                              timeout=(10, 1200))
     print(f"result: {response}")
     # result = result.text[1:-1]
-    prompt1 = ""
-    prompt2 = ""
-    img1 = None
-    img2 = None
+    prompt = ""
+    img = None
     if response.status_code == 200:
         response_json = response.json()
         print(response_json)
-        prompt1 = response_json['input_prompt']
-        prompt2 = response_json['output_prompt']
-        img1 = process_image_from_binary(response_json['no_attack_img'])
-        img2 = process_image_from_binary(response_json['attack_img'])
+        prompt = response_json['output_prompt']
+        img = process_image_from_binary(response_json['attack_img'])
     else:
         print(f"Request failed with status code {response.status_code}")
     
-    return prompt1, prompt2, img1, img2
+    return prompt, img
 
 
 css = '''
@@ -111,17 +127,17 @@ with gr.Blocks(css=custom_css) as demo:
                             interactive=True,)
     with gr.Row() as attack:
         with gr.Column(min_width=512):
+            start_button = gr.Button("Attack prepare!",size='lg')
             text_input = gr.Textbox(label="Input Prompt")
             
             orig_img = gr.Image(label="Image Generated by Input Prompt",width=512,show_share_button=False,show_download_button=False)
         with gr.Column():
-            start_button = gr.Button("UnlearnDiffAtk!",size='lg')
-        with gr.Column(min_width=512):
+            attack_button = gr.Button("UnlearnDiffAtk!",size='lg')
             text_ouput = gr.Textbox(label="Prompt Genetated by UnlearnDiffAtk")
             result_img = gr.Image(label="Image Gnerated by Prompt of UnlearnDiffAtk",width=512,show_share_button=False,show_download_button=False)
             
-
-    start_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step, atk_idx], outputs=[text_input, text_ouput, orig_img, result_img], api_name="udiff")
+        start_button.click(fn=excute_prepare, inputs=[drop_model, drop, shown_columns_step, atk_idx], outputs=[text_input, orig_img], api_name="prepare")
+        attack_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step, atk_idx], outputs=[text_ouput, result_img], api_name="udiff")
 
 
 demo.queue().launch(server_name='0.0.0.0')