Update app.py

app.py

@@ -52,7 +52,7 @@ 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, device):
+        with autocast(enabled=True, dtype=torch.bfloat16):
             init_image = dict["image"].convert("RGB").resize((imgRes, imgRes))
             mask = dict["mask"].convert("RGB").resize((imgRes, imgRes))
     else:
@@ -71,7 +71,7 @@ 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, device):
+    with autocast(enabled=True, dtype=torch.bfloat16):
         images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
     return images[0]