Update app.py

app.py

@@ -29,8 +29,8 @@ def pref_inpainting(image,
 
     pipe = AutoPipelineForInpainting.from_pretrained(
                 '/data1/kendong/joint-rl-diffusion/alignment_log/exp_reward_group_regression_all_1w_1.6boundary/iteration_2560', num_inference_steps=steps)
-    
-    pipe = pipe.to("cuda")
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    pipe = pipe.to(device)
 
    
     color, mask = outpainting_generator_rectangle(image, box_width_ratio/100, mask_random_start)
@@ -42,15 +42,15 @@ def pref_inpainting(image,
     mask_[mask < 125] = 0
     mask_[mask >= 125] = 1
     
-    color = torch.from_numpy(color).to("cuda")
-    mask = torch.from_numpy(mask).to("cuda")
+    color = torch.from_numpy(color).to(device)
+    mask = torch.from_numpy(mask).to(device)
 
 
     color, mask = transform(color), transform(mask)
     res = pipe(prompt='', image=color, mask_image=mask, eta=config.eta).images[0]
     
     
-    res.save(os.path.join('./',  'test.png'))
+    # res.save(os.path.join('./',  'test.png'))
     
     return color, res