Update main.py

main.py

@@ -219,10 +219,36 @@ def process_image(image_url):
 
 
 @app.post("/upscale-image", response_model=dict)
-async def upscale_image(request: UpscaleRequest, background_tasks: BackgroundTasks):
-    # Add the task directly to the background
-    background_tasks.add_task(process_image, request.image_url)
-    return {"status": "Processing"}
+async def upscale_image(request: UpscaleRequest):
+    if aura_sr is None:
+        raise HTTPException(status_code=500, detail="Upscaling model not initialized.")
+
+    try:
+        # Fetch the image from the provided URL
+        response = requests.get(request.image_url)
+        if response.status_code != 200:
+            raise HTTPException(status_code=400, detail="Failed to fetch the image from the provided URL.")
+        
+        img = Image.open(BytesIO(response.content))
+
+        # Perform upscaling
+        upscaled_image = aura_sr.upscale_4x_overlapped(img)
+
+        # Save the upscaled image locally
+        filename = f"upscaled_{uuid.uuid4()}.png"
+        filepath = save_image_locally(upscaled_image, filename)
+
+        # Generate a public URL for the upscaled image
+        file_url = f"/images/{filename}"
+
+        return {
+            "status": "Upscaling successful",
+            "file_path": filepath,
+            "file_url": file_url,
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during upscaling: {str(e)}")
+