Show the interface on CPU

gradio_demo.py

@@ -31,32 +31,33 @@ server_ip = args.ip
 server_port = args.port
 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:
-    SUPIR_device = 'cpu'
-    LLaVA_device = 'cpu'
+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 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
+    # 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 stage1_process(input_image, gamma_correction):
     if torch.cuda.device_count() == 0: