Update app.py

app.py

@@ -46,6 +46,7 @@ class ImageGenerationQueue:
         self.queue = asyncio.Queue()
         self.queue_size = 0
         self.queue_lock = threading.Lock()
+        self.next_id = 0
 
     def add_task(self, task):
         asyncio.run_coroutine_threadsafe(self._add_task(task), loop=asyncio.get_event_loop())
@@ -262,6 +263,23 @@ app.add_middleware(
     allow_headers=["*"],  # Allow all headers
 )
 
+@app.post("/get_image")
+async def get_image(
+    job_id: int = Form(...),    
+):
+    image_path = f"/tmp/{job_id}_output_{image.filename}"
+    with open(image_path, "rb") as file:
+        generated_image = file.read()
+
+    # Save the generated image as binary data
+    output_image_io = io.BytesIO()
+    generated_image.save(output_image_io, format="PNG")
+    output_image_io.seek(0)
+    
+    # Return the image as a streaming response
+    return StreamingResponse(content=output_image_io, media_type="image/png")
+    
+
 @app.post("/generate_image")
 async def generate_image(
     prompt: str = Form(...),
@@ -274,10 +292,10 @@ async def generate_image(
     image: UploadFile = File(...),
     background_tasks: BackgroundTasks = BackgroundTasks()
 ):
-    async def generate_image_task(position_in_queue):
+    async def generate_image_task(job_id):
         try:
             # Save the uploaded image to a temporary file
-            temp_image_path = f"/tmp/{position_in_queue}_{image.filename}"
+            temp_image_path = f"/tmp/{job_id}_{image.filename}"
             with open(temp_image_path, "wb") as temp_image:
                 temp_image.write(image.file.read())
 
@@ -289,25 +307,30 @@ async def generate_image(
             if generated_image is None:
                 return "Failed to generate image"
 
-            # Save the generated image as binary data
-            output_image_io = io.BytesIO()
-            generated_image.save(output_image_io, format="PNG")
-            output_image_io.seek(0)
+            output_image_path = f"/tmp/{job_id}_output_{image.filename}"
+            with open(output_image_path, "wb") as output_image:
+                output_image.write(generated_image)
+
+            # # Save the generated image as binary data
+            # output_image_io = io.BytesIO()
+            # generated_image.save(output_image_io, format="PNG")
+            # output_image_io.seek(0)
             
-            # Return the image as a streaming response
-            return StreamingResponse(content=output_image_io, media_type="image/png")
+            # # Return the image as a streaming response
+            # return StreamingResponse(content=output_image_io, media_type="image/png")
 
         except Exception as e:
             logger.error("Error occurred during image generation: %s", str(e))
             return "Failed to generate image"
 
     try:
-        position_in_queue = queue_manager.queue.qsize() + 1
-        queue_manager.add_task(lambda _: generate_image_task(position_in_queue))
+        id = queue_manager.next_id++
+        queue_manager.add_task(lambda _: generate_image_task(id))
+        position_in_queue = queue_manager.queue.qsize()
         # Total queue size is still async
         total_queue_size = await queue_manager.get_total_queue_size()  # Implement this function
 
-        return {"position_in_queue": position_in_queue, "total_queue_size": total_queue_size}
+        return {"job_id": id, "position_in_queue": position_in_queue, "total_queue_size": total_queue_size}
 
     except Exception as e:
         logger.error("Error occurred during image generation: %s", str(e))