Update app.py

app.py

@@ -1,27 +1,27 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, HttpUrl, conint
 from typing import List
 
 from main import get_images
 
 app = FastAPI(title="Images API", version="1.0.0")
 
-# CORS (open by default; tighten if needed)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
 )
+
 class ImageItem(BaseModel):
-    ref_url: str
-    num: int
+    ref_url: HttpUrl
+    num: conint(ge=1)
 
 class GenerateRequest(BaseModel):
     items: List[ImageItem]
 
 class ImageResult(BaseModel):
-    ref_url: "str"
-    urls: List["str"]
+    ref_url: HttpUrl
+    urls: List[HttpUrl]
 
 class GenerateResponse(BaseModel):
     urls: List[ImageResult]
@@ -30,14 +30,15 @@ class GenerateResponse(BaseModel):
 def health():
     return {"ok": True}
 
-@app.post("/generate", response_model=GenerateResponse)
-def generate(req: GenerateRequest):
+def _generate_impl(req: GenerateRequest):
     try:
-        urls = get_images(req.items)
-        # Basic sanity check
+        urls = get_images(req.dict())  # get_images expects {"items":[...]}
         if not isinstance(urls, list):
-            raise ValueError("get_images must return a List[str]")
+            raise ValueError("get_images must return a List[ImageResult-like dicts]")
         return {"urls": urls}
     except Exception as e:
-        # Log in real app; return sanitized error
         raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+    return _generate_impl(req)