Update app.py

app.py

@@ -71,7 +71,8 @@ transform = transforms.Compose([
 
 def predict(radio, dict, word_mask, prompt=""):
     if(radio == "draw a mask above"):
-        with autocast(device): #"cuda"
+        #with autocast(device): #"cuda"
+        with autocast(enable=(False if device=='cpu' else True)):
             init_image = dict["image"].convert("RGB").resize((imgRes, imgRes))
             mask = dict["mask"].convert("RGB").resize((imgRes, imgRes))
     elif(radio == "type what to keep"):
@@ -137,7 +138,8 @@ def predict(radio, dict, word_mask, prompt=""):
         os.remove(filename)
         
         #with autocast(device): #"cuda"
-        with autocast(device_type="cpu", dtype=torch.bfloat16):
+        with autocast(enable=(False if device=='cpu' else True)):
+        #with autocast(device_type="cpu", dtype=torch.bfloat16):
             images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
         return images[0]