Upload app.py

app.py

@@ -64,7 +64,7 @@ parser.add_argument('--model_ckpt',
                     default="/home/sf/weights/sd_stage2/mp_rank_00_model_states_0628.pt",
                     type=str)
 parser.add_argument('--output_path', type=str, default="./output_ipa_control_resampler")
-parser.add_argument('--device', type=str, default="cuda:0")
+# parser.add_argument('--device', type=str, default="cuda:0")
 args = parser.parse_args()
 
 # svae path
@@ -73,6 +73,8 @@ output_path = args.output_path
 if not os.path.exists(output_path):
     os.makedirs(output_path)
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+args.device = device
 
 generator = torch.Generator(device=args.device).manual_seed(42)
 vae = AutoencoderKL.from_pretrained(args.pretrained_vae_model_path).to(dtype=torch.float16, device=args.device)