Update app.py

app.py

@@ -6,6 +6,9 @@ import torch
 import torch.nn.functional as F
 from transformers import CLIPImageProcessor
 from huggingface_hub import hf_hub_download
+from pynvml import *
+nvmlInit()
+gpu_h = nvmlDeviceGetHandleByIndex(0)
 
 ctx_limit = 3500
 title = 'ViusualRWKV-v5'
@@ -14,12 +17,12 @@ vision_remote_path = "rwkv1b5-vitl336p14-577token_mix665k_visual.pth"
 vision_tower_name = 'openai/clip-vit-large-patch14-336'
 
 os.environ["RWKV_JIT_ON"] = '1'
-os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)
+os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
 
 from modeling_vision import VisionEncoder, VisionEncoderConfig
 from modeling_rwkv import RWKV
 model_path = hf_hub_download(repo_id="howard-hou/visualrwkv-5", filename=rwkv_remote_path)
-model = RWKV(model=model_path, strategy='cpu fp32')
+model = RWKV(model=model_path, strategy='cuda fp16')
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
@@ -32,6 +35,8 @@ vision_local_path = hf_hub_download(repo_id="howard-hou/visualrwkv-5", filename=
 vision_state_dict = torch.load(vision_local_path, map_location='cpu')
 visual_encoder.load_state_dict(vision_state_dict)
 image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
+if torch.cuda.is_available():
+    visual_encoder = visual_encoder.cuda()
 ##########################################################################
 def generate_prompt(instruction):
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
@@ -83,9 +88,13 @@ def generate(
             yield out_str.strip()
             out_last = i + 1
 
+    gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    print(f'{timestamp} - vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
     del out
     del state
     gc.collect()
+    torch.cuda.empty_cache()
     yield out_str.strip()
 
 
@@ -157,7 +166,7 @@ with gr.Blocks(title=title) as demo:
         with gr.Column():
             output = gr.Textbox(label="Output", lines=10)
     data = gr.Dataset(components=[image, prompt], samples=examples, label="Examples", headers=["Image", "Prompt"])
-    submit.click(chatbot, [image, prompt], [output], concurrency_limit=1)
+    submit.click(chatbot, [image, prompt], [output])
     clear.click(lambda: None, [], [output])
     data.click(lambda x: x, [data], [image, prompt])