Update app.py

app.py

@@ -1,32 +1,34 @@
 import gradio as gr
 import os, gc
-from datetime import datetime
+import torch
 from transformers import CLIPImageProcessor
 from huggingface_hub import hf_hub_download
-DEFAULT_IMAGE_TOKEN = "<image>"
-
 
 ctx_limit = 3500
 num_image_embeddings = 4096
-title = "rwkv1b5-vitl336p14-577token_mix665k_rwkv"
+title = 'ViusualRWKV-v5'
+rwkv_remote_path = "rwkv1b5-vitl336p14-577token_mix665k_rwkv.pth"
+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)
 
-from modeling import UpdatableRWKV, VisualEncoder, EmbeddingMixer, VisualEncoderConfig
-model_path = hf_hub_download(repo_id="howard-hou/visualrwkv-5", filename=f"{title}.pth")
-model = UpdatableRWKV(model=model_path, strategy='cpu fp32')
+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')
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
 ##########################################################################
-emb_mixer = EmbeddingMixer(model.w["emb.weight"], 
-                           num_image_embeddings=num_image_embeddings)
-config = VisualEncoderConfig(n_embd=model.args.n_embd, 
+config = VisionEncoderConfig(n_embd=model.args.n_embd, 
                              vision_tower_name=vision_tower_name, 
                              grid_size=-1)
-visual_encoder = VisualEncoder(config)
+visual_encoder = VisionEncoder(config)
+vision_local_path = hf_hub_download(repo_id="howard-hou/visualrwkv-5", filename=vision_remote_path)
+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)
 ##########################################################################
 def generate_prompt(instruction):
@@ -35,7 +37,7 @@ def generate_prompt(instruction):
 
 def generate(
     ctx,
-    image_ids,
+    image_features,
     token_count=200,
     temperature=1.0,
     top_p=0.7,
@@ -52,14 +54,15 @@ def generate(
     out_last = 0
     out_str = ''
     occurrence = {}
-    state = None
-    print(model.w["emb.weight"].shape)
     for i in range(int(token_count)):
         if i == 0:
-            input_ids = (image_ids + pipeline.encode(ctx))[-ctx_limit:]
+            input_ids = pipeline.encode(ctx)
+            text_embs = model.w['emb.weight'][input_ids]
+            input_embs = torch.cat((image_features, text_embs), dim=0)[-ctx_limit:]
+            out, state = model.forward(embs=input_embs, state=None)
         else:
             input_ids = [token]
-        out, state = model.forward(input_ids, state)
+            out, state = model.forward(input_ids, state)
         for n in occurrence:
             out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
 
@@ -100,14 +103,9 @@ examples = [
 ]
 def chatbot(image, question):
     image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
-    image_features = visual_encoder.encode_images(image.unsqueeze(0))
-    emb_mixer.set_image_embeddings(image_features.squeeze(0))
-    model.update_emb_weight(emb_mixer.get_input_embeddings())
-    print(emb_mixer.get_input_embeddings().shape)
-    print(model.w["emb.weight"].shape)
-    image_ids = [i for i in range(emb_mixer.image_start_index, emb_mixer.image_start_index + len(image_features))]
+    image_features = visual_encoder.encode_images(image.unsqueeze(0)).squeeze(0) # [L, D]
     input_text = generate_prompt(question)
-    for output in generate(input_text, image_ids):
+    for output in generate(input_text, image_features):
         yield output
 
 with gr.Blocks(title=title) as demo: