autocast

app.py

@@ -531,18 +531,20 @@ def infer(
         print(f"Rewritten Prompt: {prompt}")
     
 
-    # Generate the image
+    # Enable Autocast for better results.
-    image = pipe(
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
-        image=pil_images if len(pil_images) > 0 else None,
+        # Generate the image
-        prompt=prompt,
+        image = pipe(
-        height=height,
+            image=pil_images if len(pil_images) > 0 else None,
-        width=width,
+            prompt=prompt,
-        negative_prompt=negative_prompt,
+            height=height,
-        num_inference_steps=num_inference_steps,
+            width=width,
-        generator=generator,
+            negative_prompt=negative_prompt,
-        true_cfg_scale=true_guidance_scale,
+            num_inference_steps=num_inference_steps,
-        num_images_per_prompt=num_images_per_prompt,
+            generator=generator,
-    ).images
+            true_cfg_scale=true_guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+        ).images
 
     _maintain_session_state(
         primary_buffers=image,