Fixx

models/text_to_image.py

@@ -5,10 +5,10 @@ from typing import List, Optional
 class TextToImageRequest(BaseModel):
     prompt: str = Field(..., description="The prompt to generate an image from.")
     negative_prompt: Optional[List[str]] = Field(None, description="One or several prompts to guide what NOT to include in image generation.")
-    height: Optional[float] = Field(None, description="The height in pixels of the image to generate.")
-    width: Optional[float] = Field(None, description="The width in pixels of the image to generate.")
-    num_inference_steps: Optional[int] = Field(None, description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.")
-    guidance_scale: Optional[float] = Field(None, description="A higher guidance scale value encourages the model to generate images closely linked to the text prompt, but values too high may cause saturation and other artifacts.")
+    height: Optional[float] = Field(None, description="The height in pixels of the image to generate.", ge=64, le=2048)
+    width: Optional[float] = Field(None, description="The width in pixels of the image to generate.", ge=64, le=2048)
+    num_inference_steps: Optional[int] = Field(None, description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.", ge=1, le=500)
+    guidance_scale: Optional[float] = Field(None, description="A higher guidance scale value encourages the model to generate images closely linked to the text prompt, but values too high may cause saturation and other artifacts.", ge=1, le=20)
     model: Optional[str] = Field(None, description="The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended text-to-image model will be used.")
     scheduler: Optional[str] = Field(None, description="Override the scheduler with a compatible one.")
     #target_size: Optional[TextToImageTargetSize] = Field(None, description="The size in pixel of the output image")

routes/textToImage.py

@@ -22,4 +22,21 @@ async def text_to_image(t2i_body: TextToImageRequest):
     img_byte_arr = io.BytesIO()
     res.save(img_byte_arr, format="PNG")
     img_byte_arr.seek(0)
-    return Response(content=img_byte_arr.getvalue(), media_type="image/png")
+    try:
+        res = client.text_to_image(
+            prompt=t2i_body.prompt,
+            negative_prompt=t2i_body.negative_prompt,
+            height=t2i_body.height,
+            width=t2i_body.width,
+            num_inference_steps=t2i_body.num_inference_steps,
+            guidance_scale=t2i_body.guidance_scale,
+            scheduler=t2i_body.scheduler,
+            seed=t2i_body.seed
+        )
+        img_byte_arr = io.BytesIO()
+        res.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
+        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return {"error": str(e)}, 500