Update app.py

app.py

@@ -357,16 +357,17 @@ async def process_message(message: str):
             eos_found = False
             
             start_time = time.time()
+            
+            executor = ThreadPoolExecutor()
             while current_inputs and not eos_found:
-                with ThreadPoolExecutor() as executor:
-                    futures = [
-                        executor.submit(generate_model_response, model, current_inputs, max_tokens=max_token_limit)
-                        for model in global_data['models'].values()
-                    ]
-                    responses = [
-                        {'model': model_name, 'response': future.result()}
-                        for model_name, future in zip(global_data['models'].keys(), as_completed(futures))
-                    ]
+                futures = [
+                    executor.submit(generate_model_response, model, current_inputs, max_tokens=max_token_limit)
+                    for model in global_data['models'].values()
+                ]
+                responses = [
+                    {'model': model_name, 'response': future.result()}
+                    for model_name, future in zip(global_data['models'].keys(), as_completed(futures))
+                ]
                 unique_responses = remove_repetitive_responses(responses)
                 formatted_response = next(iter(unique_responses.values()))
                 
@@ -397,6 +398,7 @@ async def process_message(message: str):
                 current_inputs = formatted_response if len(formatted_response.split()) > 0 else ""
             
             end_time = time.time()
+            executor.shutdown(wait=True) # waits for all threads to finish
             print(f"Total time taken to process response {end_time-start_time}")
             
     return StreamingResponse(stream_response(inputs), media_type="text/plain")
@@ -419,13 +421,57 @@ async def generate_image(prompt: str):
              
             image = Image.open(io.BytesIO(image_bytes))
             print("Image generated successfully.")
-            return image
+            
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            image_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+            return JSONResponse(content={"image": image_base64})
         except Exception as e:
            print(f"Error generating image: {e}")
-           return None
+           return JSONResponse(content={"error": str(e)})
     else:
          print("No image model loaded.")
-         return None
+         return JSONResponse(content={"error": "No image model loaded"})
+
+def release_resources():
+    try:
+        torch.cuda.empty_cache()
+        gc.collect()
+    except Exception as e:
+        print(f"Failed to release resources: {e}")
+
+def resource_manager():
+    MAX_RAM_PERCENT = 10
+    MAX_CPU_PERCENT = 10
+    MAX_GPU_PERCENT = 10
+    MAX_RAM_MB = 1024 # 1GB
+
+    while True:
+        try:
+            virtual_mem = psutil.virtual_memory()
+            current_ram_percent = virtual_mem.percent
+            current_ram_mb = virtual_mem.used / (1024 * 1024)  # Convert to MB
+
+            if current_ram_percent > MAX_RAM_PERCENT or current_ram_mb > MAX_RAM_MB:
+                release_resources()
+
+            current_cpu_percent = psutil.cpu_percent()
+            if current_cpu_percent > MAX_CPU_PERCENT:
+               print("CPU usage too high, attempting to reduce nice")
+               p = psutil.Process(os.getpid())
+               p.nice(1)
+
+            if torch.cuda.is_available():
+                gpu = torch.cuda.current_device()
+                gpu_mem = torch.cuda.memory_percent(gpu)
+
+                if gpu_mem > MAX_GPU_PERCENT:
+                    release_resources()
+
+            time.sleep(10) # Check every 10 seconds
+        except Exception as e:
+            print(f"Error in resource manager: {e}")
     
 
 app = FastAPI()
@@ -440,19 +486,10 @@ async def generate(request: ChatRequest):
 
 @app.post("/generate_image")
 async def generate_image_endpoint(request: ImageRequest):
-    try:
-        image = await generate_image(request.prompt)
-        if image:
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            image_base64 = base64.b64encode(buffered.getvalue()).decode()
-
-            return JSONResponse(content={"image": image_base64})
-        else:
-            return JSONResponse(content={"error": "Image generation failed or no model loaded"})
-    except Exception as e:
-        return JSONResponse(content={"error": str(e)})
-
+   try:
+       return await generate_image(request.prompt)
+   except Exception as e:
+       return JSONResponse(content={"error": str(e)})
 
 def run_uvicorn():
     try:
@@ -460,7 +497,7 @@ def run_uvicorn():
     except Exception as e:
         print(f"Error al ejecutar uvicorn: {e}")
 
-
 if __name__ == "__main__":
     Thread(target=run_uvicorn).start()
+    Thread(target=resource_manager, daemon=True).start()  # Run resource manager in background
     asyncio.get_event_loop().run_forever()