Add the init of SUPIR

GenVideo_app.py

@@ -28,6 +28,51 @@ hf_hub_download(repo_id="camenduru/SUPIR", filename="SUPIR-v0F.ckpt", local_dir=
 hf_hub_download(repo_id="camenduru/SUPIR", filename="SUPIR-v0Q.ckpt", local_dir="yushan777_SUPIR")
 hf_hub_download(repo_id="RunDiffusion/Juggernaut-XL-Lightning", filename="Juggernaut_RunDiffusionPhoto2_Lightning_4Steps.safetensors", local_dir="RunDiffusion_Juggernaut-XL-Lightning")
 
+parser = argparse.ArgumentParser()
+parser.add_argument("--opt", type=str, default='options/SUPIR_v0.yaml')
+parser.add_argument("--ip", type=str, default='127.0.0.1')
+parser.add_argument("--port", type=int, default='6688')
+parser.add_argument("--no_llava", action='store_true', default=False)
+parser.add_argument("--use_image_slider", action='store_true', default=False)
+parser.add_argument("--log_history", action='store_true', default=False)
+parser.add_argument("--loading_half_params", action='store_true', default=False)
+parser.add_argument("--use_tile_vae", action='store_true', default=False)
+parser.add_argument("--encoder_tile_size", type=int, default=512)
+parser.add_argument("--decoder_tile_size", type=int, default=64)
+parser.add_argument("--load_8bit_llava", action='store_true', default=False)
+args = parser.parse_args()
+server_ip = args.ip
+server_port = args.port
+use_llava = not args.no_llava
+
+if torch.cuda.device_count() > 0:
+    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:
+        SUPIR_device = 'cpu'
+        LLaVA_device = 'cpu'
+
+    # 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
+
 
 # Available adapters (replace with your actual adapter names)
 adapter_options = {