diff --git a/app/schemas/generate_image.py b/app/schemas/generate_image.py
index 48cfbd2..a26d5d6 100644
--- a/app/schemas/generate_image.py
+++ b/app/schemas/generate_image.py
@@ -36,7 +36,7 @@ class Flux2ToProductImgModel(BaseModel):
     tasks_id: str
     prompt: str
     image_path: str
-    infer_step: float | None = None
+    infer_step: int | None = None
 
 
 class GenerateRelightImageModel(BaseModel):