Update app.py

app.py

@@ -80,24 +80,35 @@ def save_image_to_s3(image):
     url = f"https://{S3_BUCKET}.s3.{S3_REGION}.amazonaws.com/{filename}"
     return url
 
+def process_image(height, width, steps, scales, prompt, seed):
+    global pipe
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), Timer("inference"):
+        return pipe(
+            prompt=[prompt],
+            generator=torch.Generator().manual_seed(int(seed)),
+            num_inference_steps=int(steps),
+            guidance_scale=float(scales),
+            height=int(height),
+            width=int(width),
+            max_sequence_length=256
+        ).images[0]
+
 @app.post("/infer")
 async def infer(request: InferenceRequest):
     if request.randomize_seed:
         seed = random.randint(0, MAX_SEED)
     else:
         seed = request.seed
-    generator = torch.Generator().manual_seed(seed)
     
     try:
-        with Timer("Image generation"):
-            image = pipe(
-                prompt=request.prompt, 
-                width=request.width,
-                height=request.height,
-                num_inference_steps=request.num_inference_steps, 
-                generator=generator,
-                guidance_scale=request.guidance_scale
-            ).images[0]
+        image = process_image(
+            height=request.height,
+            width=request.width,
+            steps=request.num_inference_steps,
+            scales=request.guidance_scale,
+            prompt=request.prompt,
+            seed=seed
+        )
         
         image_url = save_image_to_s3(image)