Remove Llava

app.py

@@ -1,13 +1,8 @@
 import os
-
 import gradio as gr
-from gradio_imageslider import ImageSlider
 import argparse
-from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype
 import numpy as np
 import torch
-from SUPIR.util import create_SUPIR_model, load_QF_ckpt
-from PIL import Image
 import einops
 import copy
 import math
@@ -15,6 +10,10 @@ import time
 import random
 import spaces
 import re
+
+from gradio_imageslider import ImageSlider
+from PIL import Image
+from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype, create_SUPIR_model, load_QF_ckpt
 from huggingface_hub import hf_hub_download
 
 hf_hub_download(repo_id="laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", filename="open_clip_pytorch_model.bin", local_dir="laion_CLIP-ViT-bigG-14-laion2B-39B-b160k")
@@ -36,20 +35,11 @@ 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()
-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'
+    SUPIR_device = 'cuda:0'
 
-    # load SUPIR
+    # 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()
@@ -59,7 +49,6 @@ if torch.cuda.device_count() > 0:
     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)
-    llava_agent = None
 
 def check_upload(input_image):
     if input_image is None:
@@ -349,10 +338,7 @@ def restore(
     LQ = LQ.round().clip(0, 255).astype(np.uint8)
     LQ = LQ / 255 * 2 - 1
     LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
-    if use_llava:
-        captions = [prompt]
-    else:
-        captions = ['']
+    captions = ['']
 
     model.ae_dtype = convert_dtype(ae_dtype)
     model.model.dtype = convert_dtype(diff_dtype)
@@ -390,7 +376,7 @@ def restore(
     print(information)
 
     # Only one image can be shown in the slider
-    return [input_image] + [results[0]], gr.update(label="Downloadable results in *." + output_format + " format", format = output_format, value = results), gr.update(value = information, visible = True)
+    return [noisy_image] + [results[0]], gr.update(label="Downloadable results in *." + output_format + " format", format = output_format, value = results), gr.update(value = information, visible = True)
 
 def load_and_reset(param_setting):
     print('load_and_reset ==>>')