Update app.py

app.py

@@ -76,8 +76,8 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     pipe = pipe.to("cpu")
     torch.cuda.empty_cache()
     
-    # Return the image in the correct format for ChatInterface
-    return {"files": [image], "text": None}
+    # Return the PIL Image directly - ChatInterface will handle it properly
+    return image
 
 # --- UI Definition using gr.ChatInterface ---
 
@@ -102,6 +102,7 @@ demo = gr.ChatInterface(
     <br>
     Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>.
     </p>""",
+    multimodal=True,  # This is important for MultimodalTextbox to work
     textbox=gr.MultimodalTextbox(
         file_types=["image"],
         placeholder="Type a prompt and/or upload an image...",