Update app.py

app.py

@@ -39,19 +39,35 @@ parser.add_argument("--decoder_tile_size", type=int, default=64)
 parser.add_argument("--load_8bit_llava", action='store_true', default=True)
 args = parser.parse_args()
 
-if torch.cuda.device_count() > 0:
+use_llava = not args.no_llava
+
+if torch.cuda.device_count() >= 2:
+    SUPIR_device = 'cuda:0'
+    LLaVA_device = 'cuda:1'
+elif torch.cuda.device_count() == 1:
     SUPIR_device = 'cuda:0'
+    LLaVA_device = 'cuda:0'
+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 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:
+    llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device, load_8bit=args.load_8bit_llava, load_4bit=False)
+else:
+    llava_agent = None
+
 
 def check_upload(input_image):
     if input_image is None: