Update api_v2.py

api_v2.py

@@ -445,7 +445,7 @@ async def tts_handle(req:dict):
                                 print(f"Removed temporary file: {temp_file}")
                         except Exception as e:
                             print(f"Error removing temporary file {temp_file}: {e}")
-                            
+            torch.cuda.empty_cache()
             return StreamingResponse(streaming_generator(tts_generator, media_type), media_type=f"audio/{media_type}")
         else:
             sr, audio_data = next(tts_generator)
@@ -459,7 +459,8 @@ async def tts_handle(req:dict):
                         print(f"Removed temporary file: {temp_file}")
                 except Exception as e:
                     print(f"Error removing temporary file {temp_file}: {e}")
-                    
+
+            torch.cuda.empty_cache()
             return Response(audio_data, media_type=f"audio/{media_type}")
     except Exception as e:
         # Clean up temporary files in case of error
@@ -470,7 +471,8 @@ async def tts_handle(req:dict):
                     print(f"Removed temporary file: {temp_file}")
             except Exception as cleanup_error:
                 print(f"Error removing temporary file {temp_file}: {cleanup_error}")
-                
+
+        torch.cuda.empty_cache()
         return JSONResponse(status_code=400, content={"message": f"tts failed", "Exception": str(e)})    
 
 
@@ -612,7 +614,8 @@ async def set_gpt_weights(weights_path: str = None):
         if temp_file and os.path.exists(temp_file):
             os.remove(temp_file)
             print(f"Removed temporary weights file: {temp_file}")
-            
+        
+        torch.cuda.empty_cache()    
         return JSONResponse(status_code=200, content={"message": "success"})
     except Exception as e:
         # Clean up temp file in case of error
@@ -622,7 +625,8 @@ async def set_gpt_weights(weights_path: str = None):
                 print(f"Removed temporary weights file: {temp_file}")
             except Exception as cleanup_error:
                 print(f"Error removing temporary file {temp_file}: {cleanup_error}")
-                
+
+        torch.cuda.empty_cache()
         return JSONResponse(status_code=400, content={"message": f"change gpt weight failed", "Exception": str(e)})
 
 @APP.get("/set_sovits_weights")
@@ -646,7 +650,8 @@ async def set_sovits_weights(weights_path: str = None):
         if temp_file and os.path.exists(temp_file):
             os.remove(temp_file)
             print(f"Removed temporary weights file: {temp_file}")
-            
+
+        torch.cuda.empty_cache()
         return JSONResponse(status_code=200, content={"message": "success"})
     except Exception as e:
         # Clean up temp file in case of error
@@ -657,6 +662,7 @@ async def set_sovits_weights(weights_path: str = None):
             except Exception as cleanup_error:
                 print(f"Error removing temporary file {temp_file}: {cleanup_error}")
                 
+        torch.cuda.empty_cache()       
         return JSONResponse(status_code=400, content={"message": f"change sovits weight failed", "Exception": str(e)})