Upload 3 files

app.py

@@ -0,0 +1,211 @@
+import streamlit as st
+import torch
+import onnxruntime as ort
+from PIL import Image
+import requests
+import numpy as np
+from transformers import AutoTokenizer, AutoProcessor
+import os
+
+if not os.path.exists("vision_encoder_q4f16.onnx"):
+    os.system('wget https://huggingface.co/llava-hf/llava-interleave-qwen-0.5b-hf/resolve/main/onnx/vision_encoder_q4f16.onnx')
+if not os.path.exists("decoder_model_merged_q4f16.onnx"):
+    os.system('wget https://huggingface.co/llava-hf/llava-interleave-qwen-0.5b-hf/resolve/main/onnx/decoder_model_merged_q4f16.onnx')
+if not os.path.exists("embed_tokens_q4f16.onnx"):
+    os.system('wget https://huggingface.co/llava-hf/llava-interleave-qwen-0.5b-hf/resolve/main/onnx/embed_tokens_q4f16.onnx')
+
+# Load the tokenizer and processor
+tokenizer = AutoTokenizer.from_pretrained("llava-hf/llava-interleave-qwen-0.5b-hf")
+processor = AutoProcessor.from_pretrained("llava-hf/llava-interleave-qwen-0.5b-hf")
+
+vision_encoder_session = ort.InferenceSession("vision_encoder_q4f16.onnx")
+decoder_session = ort.InferenceSession("decoder_model_merged_q4f16.onnx")
+embed_tokens_session = ort.InferenceSession("embed_tokens_q4f16.onnx")
+
+def merge_input_ids_with_image_features(image_features, inputs_embeds, input_ids, attention_mask,pad_token_id,special_image_token_id):
+    num_images, num_image_patches, embed_dim = image_features.shape
+    batch_size, sequence_length = input_ids.shape
+    left_padding = not np.sum(input_ids[:, -1] == pad_token_id)
+    # 1. Create a mask to know where special image tokens are
+    special_image_token_mask = input_ids == special_image_token_id
+    num_special_image_tokens = np.sum(special_image_token_mask, axis=-1)
+    # Compute the maximum embed dimension
+    max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length
+    batch_indices, non_image_indices = np.where(input_ids != special_image_token_id)
+
+    # 2. Compute the positions where text should be written
+    # Calculate new positions for text tokens in merged image-text sequence.
+    # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens.
+    # `np.cumsum` computes how each image token shifts subsequent text token positions.
+    # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
+    new_token_positions = np.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1
+    nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
+    if left_padding:
+        new_token_positions += nb_image_pad[:, None]  # offset for left padding
+    text_to_overwrite = new_token_positions[batch_indices, non_image_indices]
+
+    # 3. Create the full embedding, already padded to the maximum position
+    final_embedding = np.zeros((batch_size, max_embed_dim, embed_dim), dtype=np.float32)
+    final_attention_mask = np.zeros((batch_size, max_embed_dim), dtype=np.int64)
+
+    # 4. Fill the embeddings based on the mask. If we have ["hey" "<image>", "how", "are"]
+    # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features
+    final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices]
+    final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices]
+    # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835)
+    image_to_overwrite = np.full((batch_size, max_embed_dim), True)
+    image_to_overwrite[batch_indices, text_to_overwrite] = False
+    image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None]
+
+    final_embedding[image_to_overwrite] = image_features.reshape(-1, embed_dim)
+    final_attention_mask = np.logical_or(final_attention_mask, image_to_overwrite).astype(final_attention_mask.dtype)
+    position_ids = final_attention_mask.cumsum(axis=-1) - 1
+    position_ids = np.where(final_attention_mask == 0, 1, position_ids)
+
+    # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.
+    batch_indices, pad_indices = np.where(input_ids == pad_token_id)
+    indices_to_mask = new_token_positions[batch_indices, pad_indices]
+    final_embedding[batch_indices, indices_to_mask] = 0
+
+    return final_embedding, final_attention_mask, position_ids
+
+# Load model and processor
+
+def describe_image(image):
+    if(image.mode != 'RGB'):
+      image = image.convert('RGB')
+    conversation = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant who describes image."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image in about 200 words and explain each and every element in full detail"},
+                {"type": "image"},
+            ],
+        },
+    ]
+
+    # Apply chat template
+    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+    # Preprocess the image and text
+    inputs = processor(images=image, text=prompt, return_tensors="np")
+    vision_input_name = vision_encoder_session.get_inputs()[0].name
+    vision_output_name = vision_encoder_session.get_outputs()[0].name
+    vision_features = vision_encoder_session.run([vision_output_name], {vision_input_name: inputs["pixel_values"]})[0]
+
+    # print('Total Time for Image Features Making ', time.time() - start)
+
+    # Tokens for the prompt
+    input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
+
+    # Prepare inputs
+    sequence_length = input_ids.shape[1]
+    batch_size = 1
+    num_layers = 24
+    head_dim = 64
+    num_heads = 16
+    pad_token_id = tokenizer.pad_token_id
+    past_sequence_length = 0  # Set to 0 for the initial pass
+    special_image_token_id = 151646
+
+    # Position IDs
+    position_ids = np.arange(sequence_length, dtype=np.int64).reshape(1, -1)
+
+    # Past Key Values
+    past_key_values = {
+        f"past_key_values.{i}.key": np.zeros((batch_size, num_heads, past_sequence_length, head_dim), dtype=np.float32)
+        for i in range(num_layers)
+    }
+    past_key_values.update({
+        f"past_key_values.{i}.value": np.zeros((batch_size, num_heads, past_sequence_length, head_dim), dtype=np.float32)
+        for i in range(num_layers)
+    })
+
+    # Run embed tokens
+    embed_input_name = embed_tokens_session.get_inputs()[0].name
+    embed_output_name = embed_tokens_session.get_outputs()[0].name
+    token_embeddings = embed_tokens_session.run([embed_output_name], {embed_input_name: input_ids})[0]
+
+    # Combine token embeddings and vision features
+    combined_embeddings, attention_mask, position_ids = merge_input_ids_with_image_features(vision_features, token_embeddings, input_ids, attention_mask,pad_token_id,special_image_token_id)
+    combined_len = combined_embeddings.shape[1]
+
+    # Combine all inputs
+    decoder_inputs = {
+        "attention_mask": attention_mask,
+        "position_ids": position_ids,
+        "inputs_embeds": combined_embeddings,
+        **past_key_values
+    }
+
+    # Print input shapes
+    for name, value in decoder_inputs.items():
+        print(f"{name} shape: {value.shape} dtype {value.dtype}")
+
+    # Run the decoder
+    decoder_input_names = [input.name for input in decoder_session.get_inputs()]
+    decoder_output_name = decoder_session.get_outputs()[0].name
+    names = [n.name for n in decoder_session.get_outputs()]
+    outputs = decoder_session.run(names, {name: decoder_inputs[name] for name in decoder_input_names if name in decoder_inputs})
+
+    # ... (previous code remains the same until after the decoder run)
+    # print(f"Outputs shape: {outputs[0].shape}")
+    # print(f"Outputs type: {outputs[0].dtype}")
+
+    # Process outputs (decode tokens to text)
+    generated_tokens = []
+    eos_token_id = tokenizer.eos_token_id
+    max_new_tokens = 2048
+
+    for i in range(max_new_tokens):
+        logits = outputs[0]
+        past_kv = outputs[1:]
+        logits_next_token = logits[:, -1]
+        token_id = np.argmax(logits_next_token)
+
+        if token_id == eos_token_id:
+            break
+
+        generated_tokens.append(token_id)
+
+        # Prepare input for next token generation
+        new_input_embeds = embed_tokens_session.run([embed_output_name], {embed_input_name: np.array([[token_id]])})[0]
+
+        past_key_values = {name.replace("present", "past_key_values"): value for name, value in zip(names[1:], outputs[1:])}
+
+        attention_mask = np.ones((1, combined_len + i + 1), dtype=np.int64)
+        position_ids = np.arange(combined_len + i + 1, dtype=np.int64).reshape(1, -1)[:, -1:]
+
+        decoder_inputs = {
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "inputs_embeds": new_input_embeds,
+            **past_key_values
+        }
+
+        outputs = decoder_session.run(names, {name: decoder_inputs[name] for name in decoder_input_names if name in decoder_inputs})
+
+    # Convert to list of integers
+    token_ids = [int(token) for token in generated_tokens]
+
+    print(f"Generated token IDs: {token_ids}")
+
+    # Decode tokens one by one
+    decoded_tokens = [tokenizer.decode([token]) for token in token_ids]
+    print(f"Decoded tokens: {decoded_tokens}")
+
+    # Full decoded output
+    decoded_output = tokenizer.decode(token_ids, skip_special_tokens=True)
+    return decoded_output
+
+# Streamlit app
+st.title("Image Description Generator")
+uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+if uploaded_image is not None:
+    image = Image.open(uploaded_image)
+    description = describe_image(image)
+    st.text_area("Description", description, height=300)

packages.txt

@@ -0,0 +1 @@
+wget

requirements.txt

@@ -0,0 +1,9 @@
+onnxruntime
+onnx
+gradio
+Pillow
+torch
+transformers
+numpy
+flask
+streamlit