Update main.py

main.py

@@ -33,7 +33,6 @@ app.add_middleware(
 # Globals
 executor = ThreadPoolExecutor(max_workers=10)
 llm = None
-aura_sr = None
 client = MongoClient(f"mongodb+srv://hammad:{quote_plus('momimaad@123')}@cluster0.2a9yu.mongodb.net/")
 db = client["Flux"]
 collection = db["chat_histories"]
@@ -51,14 +50,20 @@ image_store = {}
 
 @app.on_event("startup")
 async def startup():
-    global llm, upscale_model
+    global llm, aura_sr
     llm = ChatGroq(
         model="llama-3.3-70b-versatile",
         temperature=0.7,
         max_tokens=1024,
         api_key="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",
     )
-    aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
+    # Ensure AuraSR is initialized properly
+    try:
+        aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
+    except Exception as e:
+        print(f"Error initializing AuraSR: {e}")
+        aura_sr = None
+
 
 @app.on_event("shutdown")
 def shutdown():
@@ -196,19 +201,23 @@ class UpscaleRequest(BaseModel):
 
 
 def process_image(image_url):
+    if aura_sr is None:
+        raise RuntimeError("aura_sr is None. Ensure it's initialized during startup.")
+
     response = requests.get(image_url)
     img = Image.open(BytesIO(response.content))
-    if aura_sr is None:
-        raise RuntimeError("aura_sr is None. Check initialization.")
+
     try:
         upscaled_image = aura_sr.upscale_4x_overlapped(img)
     except Exception as e:
         raise RuntimeError(f"Error during upscaling: {str(e)}")
+    
     filename = f"upscaled_{uuid.uuid4()}.png"
     filepath = save_image_locally(upscaled_image, filename)
     return filepath
 
 
+
 @app.post("/upscale-image", response_model=dict)
 async def upscale_image(request: UpscaleRequest, background_tasks: BackgroundTasks):
     # Add the task directly to the background