Update app.py

app.py

@@ -267,7 +267,10 @@ app.add_middleware(
 async def get_image(
     job_id: int = Form(...)
 ):
-    image_path = f"/tmp/{job_id}_output_{image.filename}"
+    image_path = f"/tmp/{job_id}_output.png"
+    if !os.path.isfile(image_path):
+        return None
+        
     with open(image_path, "rb") as file:
         generated_image = file.read()
 
@@ -307,7 +310,7 @@ async def generate_image(
             if generated_image is None:
                 return "Failed to generate image"
 
-            output_image_path = f"/tmp/{job_id}_output_{image.filename}"
+            output_image_path = f"/tmp/{job_id}_output.png"
             with open(output_image_path, "wb") as output_image:
                 output_image.write(generated_image)