Update app.py

app.py

@@ -52,7 +52,8 @@ def predict(radio, dict, word_mask, prompt=""):
     if(radio == "draw a mask above"):
         #with autocast("cuda"):
         #with autocast(device): #enable=(False if device=='cpu' else True)):
-        with autocast(enabled=True, dtype=torch.bfloat16):
+        #with autocast(enabled=True, dtype=torch.bfloat16):
+        with torch.cuda.amp.autocast(True):
             init_image = dict["image"].convert("RGB").resize((imgRes, imgRes))
             mask = dict["mask"].convert("RGB").resize((imgRes, imgRes))
     else:
@@ -71,7 +72,8 @@ def predict(radio, dict, word_mask, prompt=""):
         os.remove(filename)
     #with autocast("cuda"):
     #with autocast(device): #enable=(False if device=='cpu' else True)):
-    with autocast(enabled=True, dtype=torch.bfloat16):
+    #with autocast(enabled=True, dtype=torch.bfloat16):
+    with torch.cuda.amp.autocast(True):
         images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
     return images[0]