Update app.py

app.py

@@ -14,14 +14,13 @@ 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 rwkv.model import RWKV
+from modeling import UpdatableRWKV, VisualEncoder, EmbeddingMixer, VisualEncoderConfig
 model_path = hf_hub_download(repo_id="howard-hou/visualrwkv-5", filename=f"{title}.pth")
-model = RWKV(model=model_path, strategy='cpu fp32')
+model = UpdatableRWKV(model=model_path, strategy='cpu fp32')
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
 ##########################################################################
-from modeling import VisualEncoder, EmbeddingMixer, VisualEncoderConfig
 emb_mixer = EmbeddingMixer(model.w["emb.weight"], 
                            num_image_embeddings=num_image_embeddings)
 config = VisualEncoderConfig(n_embd=model.args.n_embd, 
@@ -102,9 +101,7 @@ 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))
-    global model.w["emb.weight"]
-    model.w["emb.weight"] = emb_mixer.get_input_embeddings()
-    print(model.w["emb.weight"].shape)
+    model.update_emb_weight(emb_mixer.get_input_embeddings())
     image_ids = [i for i in range(emb_mixer.image_start_index, emb_mixer.image_start_index + len(image_features))]
     input_text = generate_prompt(question)
     for output in generate(input_text, image_ids):