Create app.py

app.py

@@ -0,0 +1,51 @@
+import torch
+from transformers import AutoTokenizer
+from model import CustomClipPhi2
+
+clip_model_name = "openai/clip-vit-base-patch32"
+phi_model_name  = "microsoft/phi-2"
+
+tokenizer  = AutoTokenizer.from_pretrained(phi_model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+IMAGE_TOKEN_ID = 23903 # token for word Comments
+device = "cuda" if torch.cuda.is_available() else "cpu"
+max_tokens = 30
+model = CustomClipPhi2(tokenizer, phi2_model_name, clip_model_name, clip_embed=768, phi_embed=2560)
+
+def generate(images):
+        clip_outputs = model.clip_model(**images)
+        # remove cls token
+        images = clip_outputs.last_hidden_state[:, 1:, :]
+        image_embeddings = model.projection_layer(images).to(torch.float16)
+
+        batch_size = images.size()[0]
+        predicted_caption = torch.full((batch_size, max_tokens), model.EOS_TOKEN_ID, dtype=torch.long, device=device)
+        img_token_tensor = torch.tensor(IMAGE_TOKEN_ID).repeat(batch_size, 1)
+        img_token_embeds = model.phi2_model.model.embed_tokens(img_token_tensor.to(image_embeddings.device))
+        combined_embeds  = torch.cat([image_embeddings, img_token_embeds], dim=1)
+
+        for pos in range(max_tokens - 1):
+            model_output_logits = model.phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+            predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
+            predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1)
+            predicted_caption[:, pos] = predicted_word_token.view(1,-1).to('cpu')
+            next_token_embeds = model.phi2_model.model.embed_tokens(predicted_word_token)
+            combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
+        return predicted_caption
+
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=generate,  # Function to be called on user input
+    inputs=gr.Image(
+            width=416, height=416, 
+            type="pil", image_mode='RGB', label="Upload Image"
+        ),
+    outputs=gr.Textbox(
+            label="Response from AI Model: ", 
+        ),
+    examples = ['car.jpg']
+)
+
+# Launch the Gradio app
+iface.launch()