Update app.py

app.py

@@ -44,23 +44,34 @@ use_llava = not args.no_llava
 if torch.cuda.device_count() >= 2:
     SUPIR_device = 'cuda:0'
     LLaVA_device = 'cuda:1'
+    # Load SUPIR
+    model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
+    if args.loading_half_params:
+        model = model.half()
+    if args.use_tile_vae:
+        model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
+    model = model.to(SUPIR_device)
+    model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
+    model.current_model = 'v0-Q'
+    ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
 elif torch.cuda.device_count() == 1:
     SUPIR_device = 'cuda:0'
     LLaVA_device = 'cuda:0'
+    # Load SUPIR
+    model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
+    if args.loading_half_params:
+        model = model.half()
+    if args.use_tile_vae:
+        model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
+    model = model.to(SUPIR_device)
+    model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
+    model.current_model = 'v0-Q'
+    ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
 else:
     raise ValueError('Currently support CUDA only.')
 
     
-# Load SUPIR
-model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
-if args.loading_half_params:
-    model = model.half()
-if args.use_tile_vae:
-    model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
-model = model.to(SUPIR_device)
-model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
-model.current_model = 'v0-Q'
-ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
+
 
 # load LLaVA
 if use_llava: