Update main.py

main.py

@@ -41,9 +41,8 @@ llms = {
     
 #Fast API
 app = FastAPI()
-llm = None
 
-def clean_memory():
+def clean_memory(llm):
     llm = None
     gc.collect()
     
@@ -63,6 +62,5 @@ async def stream(item: dict):
 
     prompt = f"{prefix}{user.replace('{prompt}', item['prompt'])}{suffix}"
     result = llm(prompt, max_tokens=max_tokens)
-    thread = threading.Thread(target=clean_memory)
-    thread.start()
+    thread = threading.Thread(target=clean_memory, args=(llm)).start()
     return result