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",