Update app.py

app.py

@@ -3,7 +3,9 @@ import gradio as gr
 import os
 import requests
 import json
+import base64
 from huggingface_hub import login
+import Image
 
 
 # myip = os.environ["0.0.0.0"]
@@ -27,6 +29,13 @@ from about import (
 )
 
 
+def process_image_from_binary(img_stream):
+    image_data = base64.b64decode(img_stream)
+    image_bytes = BytesIO(image_data)
+    img = Image.open(image_bytes)
+    
+    return img
+
 def excute_udiff(diffusion_model_id, concept, steps):
     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}")
@@ -35,15 +44,19 @@ def excute_udiff(diffusion_model_id, concept, steps):
     # result = result.text[1:-1]
     prompt1 = ""
     prompt2 = ""
+    img1 = None
+    img2 = 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'])
     else:
         print(f"Request failed with status code {response.status_code}")
     
-    return prompt1, prompt2
+    return prompt1, prompt2, img1, img2
 
 
 css = '''
@@ -91,15 +104,16 @@ with gr.Blocks(css=custom_css) as demo:
     with gr.Row() as attack:
         with gr.Column(min_width=260):
             text_input = gr.Textbox(label="Input Prompt")
-            img1 = gr.Image("images/cheetah.jpg",label="Image Generated by Input Prompt",width=260,show_share_button=False,show_download_button=False)
+            
+            img1 = gr.Image(orig_img,label="Image Generated by Input Prompt",width=260,show_share_button=False,show_download_button=False)
         with gr.Column():
             start_button = gr.Button("UnlearnDiffAtk!",size='lg')
         with gr.Column(min_width=260):
             text_ouput = gr.Textbox(label="Prompt Genetated by UnlearnDiffAtk")
-            img2 = gr.Image("images/cheetah.jpg",label="Image Gnerated by Prompt of UnlearnDiffAtk",width=260,show_share_button=False,show_download_button=False)
+            img2 = gr.Image(result_img,label="Image Gnerated by Prompt of UnlearnDiffAtk",width=260,show_share_button=False,show_download_button=False)
             
 
-    start_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step], outputs=[text_input, text_ouput], api_name="udiff")
+    start_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step], outputs=[text_input, text_ouput, orig_img, result_img], api_name="udiff")
 
 
 demo.queue().launch(server_name='0.0.0.0')