Update app.py

app.py

@@ -1,10 +1,17 @@
 import gradio as gr
 import os, gc
 from datetime import datetime
+from transformers import CLIPImageProcessor
 from huggingface_hub import hf_hub_download
+from typing import List, Dict
+from dataclasses import dataclass
+DEFAULT_IMAGE_TOKEN = "<image>"
+
 
 ctx_limit = 3500
+num_image_embeddings = 4096
 title = "rwkv1b5-vitl336p14-577token_mix665k_rwkv"
+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)
@@ -17,32 +24,22 @@ pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
 ##########################################################################
 from modeling import VisualEncoder, EmbeddingMixer, VisualEncoderConfig
-emb_mixer = EmbeddingMixer(model.w["emb.weight"], num_image_embeddings=4096)
+emb_mixer = EmbeddingMixer(model.w["emb.weight"], 
+                           num_image_embeddings=num_image_embeddings)
 config = VisualEncoderConfig(n_embd=model.args.n_embd, 
-                             vision_tower_name='openai/clip-vit-large-patch14-336', 
+                             vision_tower_name=vision_tower_name, 
                              grid_size=-1)
 visual_encoder = VisualEncoder(config)
+image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
 ##########################################################################
-def generate_prompt(instruction, input=""):
+def generate_prompt(instruction):
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
     input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
-    if input:
-        return f"""Instruction: {instruction}
-
-Input: {input}
-
-Response:"""
-    else:
-        return f"""User: hi
-
-Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
+    return f"\n{instruction}\n\nAssistant:"
 
-User: {instruction}
-
-Assistant:"""
-
-def evaluate(
+def generate(
     ctx,
+    image_ids,
     token_count=200,
     temperature=1.0,
     top_p=0.7,
@@ -61,7 +58,11 @@ def evaluate(
     occurrence = {}
     state = None
     for i in range(int(token_count)):
-        out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:] if i == 0 else [token], state)
+        if i == 0:
+            input_ids = (image_ids + pipeline.encode(ctx))[-ctx_limit:]
+        else:
+            input_ids = [token]
+        out, state = model.forward(input_ids, state)
         for n in occurrence:
             out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
 
@@ -101,8 +102,13 @@ examples = [
     ]
 ]
 def test(image, question):
-    print(image, question)
-    return question
+    image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
+    image_features = visual_encoder.encode_images(image.unsqueeze(0))
+    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):
+        yield output
+
 demo = gr.Interface(fn=test, 
                     inputs=[gr.Image(type='pil'), "text"], 
                     outputs="text",