Update app.py

app.py

@@ -16,6 +16,10 @@ from PIL import Image
 from packaging import version
 from PIL.PngImagePlugin import PngInfo
 
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
 with httpimport.remote_repo(os.getenv("MODULE_URL")):
   import pipeline
   pipe, pipe2, pipe_img2img, pipe2_img2img = pipeline.get_pipeline_initialize()
@@ -171,19 +175,20 @@ def zero_inference_api(prompt, radio="model-v2", preset=PRESET_Q, h=1216, w=832,
     generator = torch.Generator(device="cuda").manual_seed(seed)
     if inference_steps > 50:
         inference_steps = 50
-    
-    if not do_img2img:
-        if radio == "model-v2":
-            image = pipe(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]    
-        else:
-            image = pipe2(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]           
-    else:
-        init_image = Image.fromarray(init_image)
-        if radio == "model-v2":
-            image = pipe_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]    
+
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        if not do_img2img:
+            if radio == "model-v2":
+                image = pipe(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]    
+            else:
+                image = pipe2(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]           
         else:
-            image = pipe2_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]
-    
+            init_image = Image.fromarray(init_image)
+            if radio == "model-v2":
+                image = pipe_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]    
+            else:
+                image = pipe2_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]
+        
     naifix = prompt[:40].replace(":", "_").replace("\\", "_").replace("/", "_") + f" s-{seed}-"
     with tempfile.NamedTemporaryFile(prefix=naifix, suffix=".png", delete=False) as tmpfile:
         parameters = {