Update app.py

app.py

@@ -26,6 +26,7 @@ from transformers import AutoFeatureExtractor, CLIPFeatureExtractor
 import random
 import time
 import tempfile
+import threading
 
 logger = logging.getLogger(__name__)
 # Set the logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
@@ -43,16 +44,32 @@ logger.addHandler(stream_handler)
 class ImageGenerationQueue:
     def __init__(self):
         self.queue = asyncio.Queue()
+        self.queue_size = 0
+        self.queue_lock = threading.Lock()
 
-    async def add_task(self, task):
+    def add_task(self, task):
+        asyncio.run_coroutine_threadsafe(self._add_task(task), loop=asyncio.get_event_loop())
+
+    async def _add_task(self, task):
         await self.queue.put(task)
+        # Update queue size in a thread-safe manner
+        with self.queue_lock:
+            self.queue_size = self.queue.qsize()
 
     async def process_queue(self):
         while True:
             task = await self.queue.get()
             await task()
+            # Update queue size in a thread-safe manner
+            with self.queue_lock:
+                self.queue_size = self.queue.qsize()
             self.queue.task_done()
 
+    async def get_total_queue_size(self):
+        # Return the queue size in a thread-safe manner
+        with self.queue_lock:
+            return self.queue_size
+
 app = FastAPI()
 queue_manager = ImageGenerationQueue()
 
@@ -230,13 +247,13 @@ def generate_image_from_parameters(prompt, guidance_scale, controlnet_scale, con
         output_image_binary = output_image_io.read()
 
         # Return the generated image binary data
-        logger.debug("Output Values: generated_image=<binary data>")
         return output_image_binary
 
     except Exception as e:
         # Handle exceptions and return an error message if something goes wrong
         return str(e)
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # You can replace ["*"] with specific origins if needed
@@ -256,7 +273,7 @@ async def generate_image(
     sampler_type: str = Form(...),
     image: UploadFile = File(...)
 ):
-    async def generate_image_task():
+    def generate_image_task():
         try:
             # Save the uploaded image to a temporary file
             temp_image_path = f"/tmp/{int(time.time())}_{image.filename}"
@@ -284,9 +301,11 @@ async def generate_image(
             return "Failed to generate image"
 
     try:
-        await queue_manager.add_task(generate_image_task)
+        queue_manager.add_task(generate_image_task)
         position_in_queue = queue_manager.queue.qsize()
-        total_queue_size = queue_manager.queue.qsize()  # Implement this function
+        # Total queue size is still async
+        total_queue_size = await get_total_queue_size(queue_manager.queue)  # Implement this function
+
         return {"position_in_queue": position_in_queue, "total_queue_size": total_queue_size}
 
     except Exception as e: