diff --git a/app/schemas/generate_image.py b/app/schemas/generate_image.py
index 3e150d5..48cfbd2 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
-    image_strength: int
+    infer_step: float | None = None
 
 
 class GenerateRelightImageModel(BaseModel):