Update app.py

app.py

@@ -11,7 +11,7 @@ client = Client("Geek7/mdztxi2")
 @app.route('/generate-image', methods=['POST'])
 def generate_image():
     data = request.get_json()
-    
+
     # Check for required fields in the request data
     model_str = data.get('model_str')
     prompt = data.get('prompt')
@@ -22,6 +22,7 @@ def generate_image():
 
     # Make a prediction request
     try:
+        # Making sure to capture the actual result returned
         result = client.predict(
             model_str=model_str,
             prompt=prompt,
@@ -29,10 +30,11 @@ def generate_image():
             api_name="/predict"
         )
 
-        # Check if result is a file path or image bytes
-        # If the result is an image path, adjust accordingly
-        if isinstance(result, str) and os.path.exists(result):  # Assuming result is the path to the image
+        # If the result is a file path, send it; if it's image bytes, handle accordingly
+        if isinstance(result, str) and os.path.exists(result):  # If the result is a path
             return send_file(result, mimetype='image/png')  # Send the generated image file
+        elif isinstance(result, bytes):  # If the result is in bytes
+            return send_file(io.BytesIO(result), mimetype='image/png')  # Return image bytes as a file
         else:
             return jsonify({"error": "Image generation failed, result not valid."}), 500
     except Exception as e: