Update myapp.py

myapp.py

@@ -6,6 +6,7 @@ import os
 from threading import RLock
 from huggingface_hub import InferenceClient
 from all_models import models  # Importing models from all_models
+from io import BytesIO  # For converting image to bytes
 
 myapp = Flask(__name__)
 CORS(myapp)  # Enable CORS for all routes
@@ -40,9 +41,16 @@ async def infer(client, prompt, seed=1, timeout=inference_timeout, model="prompt
     
     if task.done() and result is not None:
         with lock:
+            # Convert image result to bytes
+            image_bytes = BytesIO()
+            result.save(image_bytes, format='PNG')  # Save the image to a BytesIO object
+            image_bytes.seek(0)  # Go to the start of the byte stream
+            
+            # Save the result image as a temporary file
             temp_image = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
             with open(temp_image.name, "wb") as f:
-                f.write(result)  # Save the result image as a temporary file
+                f.write(image_bytes.read())  # Write the bytes to the temp file
+
         return temp_image.name  # Return the path to the saved image
     return None
 
@@ -79,4 +87,4 @@ def generate_api():
 
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":
-    myapp.run(host='0.0.0.0', port=7860)  # Run directly if needed for testin
+    myapp.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing