Update app.py

app.py

@@ -21,6 +21,33 @@ is_spaces = True if "SPACE_ID" in os.environ else False
 is_shared_ui = False
 
 
+def process_image_from_binary(img_stream):
+    if img_stream is None:
+        print("no image binary")
+        return
+    image_data = base64.b64decode(img_stream)
+    image_bytes = BytesIO(image_data)
+    img = Image.open(image_bytes)
+    
+    return img
+
+def generate_img(concept, prompt):
+    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://{}:{}/generate'.format(myip, myport), 
+                             json={"concept": concept, "prompt": prompt},
+                             timeout=(10, 1200))
+    print(f"result: {response}")
+    image = None
+    if response.status_code == 200:
+        response_json = response.json()
+        print(response_json)
+        image = process_image_from_binary(response_json[‘img'])
+    else:
+        print(f"Request failed with status code {response.status_code}")
+    
+    return image
+
 with gr.Blocks() as demo:
     gr.HTML(TITLE)
     gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
@@ -37,6 +64,8 @@ with gr.Blocks() as demo:
         with gr.Column():
             start_button = gr.Button("AdvUnlearn",size='lg')
         with gr.Column(min_width=512):
-            img2 = gr.Image("images/cheetah.jpg",label="Image Generated with AdvUnlearn",width=512,show_share_button=False,show_download_button=False)
+            result_img = gr.Image(label="Image Gnerated by AdvUnlearn",width=512,show_share_button=False,show_download_button=False)
+
+    start_button.click(fn=generate_img, inputs=[drop_text, text_input], outputs=result_img, api_name="generate")
             
 demo.launch()