Update app.py

app.py

@@ -35,18 +35,18 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     if image:
-        # Convert image to base64
+        # Convert image(s) to base64
         if isinstance(image, Image.Image):
             image_b64 = image_to_base64(image)
             messages.append({"role": "user", "content": "Image uploaded", "image": image_b64})
         else:
-            # Handle multiple images if necessary
             for img in image:
                 image_b64 = image_to_base64(img)
                 messages.append({"role": "user", "content": "Image uploaded", "image": image_b64})
 
     # Call Hugging Face model for response
     try:
+        responses = []
         for response in client.chat_completion(
             messages,
             max_tokens=max_tokens,
@@ -54,11 +54,20 @@ def respond(
             temperature=temperature,
             top_p=top_p,
         ):
-            if response.choices:
-                token = response.choices[0].delta.content
-                yield token
+            token = response.choices[0].delta.content
+            responses.append(token)
+
+            # Check if the response contains an image to be displayed
+            if response.choices[0].delta.image:
+                image_b64 = response.choices[0].delta.image
+                image_data = base64.b64decode(image_b64)
+                image = Image.open(io.BytesIO(image_data))
+                yield responses, image
+            else:
+                yield responses, None
+
     except Exception as e:
-        yield str(e)
+        yield [str(e)], None
 
 # Debugging print statements
 print("Starting Gradio interface setup...")
@@ -68,9 +77,12 @@ try:
         fn=respond,
         inputs=[
             gr.Textbox(label="Message"),
-            gr.Image(label="Upload Medical Image", type="pil")
+            gr.Image(label="Upload Medical Image", type="pil", optional=True)
+        ],
+        outputs=[
+            gr.Textbox(label="Response", placeholder="Model response will appear here..."),
+            gr.Image(label="Generated Image", type="pil", output=True)
         ],
-        outputs=gr.Textbox(label="Response", placeholder="Model response will appear here..."),
         title="LLAVA Model - Medical Image and Question",
         description="Upload a medical image and ask a specific question about the image for a medical description.",
         additional_inputs=[