Update main.py

main.py

@@ -191,10 +191,14 @@ async def generate_image(request: ImageRequest):
         "file_url": file_url,
     }
 
+class UpscaleRequest(BaseModel):
+    image_url: str
+
+
 @app.post("/upscale-image", response_model=dict)
-async def upscale_image(image_url: str, background_tasks: BackgroundTasks):
+async def upscale_image(request: UpscaleRequest, background_tasks: BackgroundTasks):
     def process_image():
-        response = requests.get(image_url)
+        response = requests.get(request.image_url)
         img = Image.open(BytesIO(response.content))
         upscaled_image = aura_sr.upscale_4x_overlapped(img)
         filename = f"upscaled_{uuid.uuid4()}.png"
@@ -206,6 +210,7 @@ async def upscale_image(image_url: str, background_tasks: BackgroundTasks):
     return {"status": "Processing"}
 
 
+
 @app.get("/")
 async def root():
     return {"message": "API is up and running!"}