added queue

app.py

@@ -40,7 +40,21 @@ stream_handler.setFormatter(formatter)
 logger.addHandler(file_handler)
 logger.addHandler(stream_handler)
 
+class ImageGenerationQueue:
+    def __init__(self):
+        self.queue = asyncio.Queue()
+
+    async def add_task(self, task):
+        await self.queue.put(task)
+
+    async def process_queue(self):
+        while True:
+            task = await self.queue.get()
+            await task()
+            self.queue.task_done()
+
 app = FastAPI()
+queue_manager = ImageGenerationQueue()
 
 BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"
 
@@ -242,7 +256,8 @@ async def generate_image(
     sampler_type: str = Form(...),
     image: UploadFile = File(...)
 ):
-    try:
+    async def generate_image_task():
+        try:
         # Save the uploaded image to a temporary file
         temp_image_path = f"/tmp/{int(time.time())}_{image.filename}"
         with open(temp_image_path, "wb") as temp_image:
@@ -268,6 +283,19 @@ async def generate_image(
         logger.error("Error occurred during image generation: %s", str(e))
         return "Failed to generate image"
 
+        pass
+    try:
+        await queue_manager.add_task(generate_image_task)
+        position_in_queue = queue_manager.queue.qsize()
+        total_queue_size = queue_manager.queue.qsize()  # Implement this function
+        return {"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))
+        return "Failed to add task to the queue"        
+
+
+
 async def start_fastapi():
     # Get internal IP address
     internal_ip = socket.gethostbyname(socket.gethostname())
@@ -280,6 +308,10 @@ async def start_fastapi():
 
     print(f"Internal URL: http://{internal_ip}:7860")
     print(f"Public URL: http://{public_ip}:7860")
+    queue_manager = ImageGenerationQueue()
+
+    # Start processing the image generation queue
+    queue_processing_task = asyncio.create_task(queue_manager.process_queue())
 
     # Run FastAPI using hypercorn
     config = uvicorn.Config(app="app:app", host="0.0.0.0", port=7860, reload=True)