Switch to gpt2 and blip models

app/services/ai_service.py

@@ -1,9 +1,14 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import  AutoTokenizer, AutoModelForCausalLM
 
-MODEL_NAME = "bickett/meme-llama"
+# Load GPT-2 small (text generation)
+# -----------------------------
+GPT2_MODEL_NAME = "gpt2"
+gpt2_tokenizer = AutoTokenizer.from_pretrained(GPT2_MODEL_NAME)
+gpt2_model = AutoModelForCausalLM.from_pretrained(GPT2_MODEL_NAME)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+# tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+# model = AutoModelWithLMHead.from_pretrained(MODEL_NAME)
+# model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 
 # def generate_meme_caption(image_path: str, max_length: int = 40, num_return_sequences: int = 3):
@@ -45,8 +50,8 @@ def generate_captions(prompt: str, max_length: int = 50, num_return_sequences: i
     """
     Generate AI meme captions given a prompt using Meme-LLaMA
     """
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(
+    inputs = gpt2_tokenizer(prompt, return_tensors="pt")
+    outputs = gpt2_model.generate(
         **inputs,
         max_length=max_length,
         num_return_sequences=num_return_sequences,
@@ -55,7 +60,7 @@ def generate_captions(prompt: str, max_length: int = 50, num_return_sequences: i
         temperature=0.8
     )
 
-    captions = [tokenizer.decode(out, skip_special_tokens=True) for out in outputs]
+    captions = [gpt2_tokenizer.decode(out, skip_special_tokens=True) for out in outputs]
     return captions
 
 # def generate_captions(prompt: str):

app/services/image_service.py

@@ -0,0 +1,15 @@
+from PIL import Image
+from transformers import AutoProcessor, BlipForConditionalGeneration
+
+
+BLIP_MODEL_NAME = "Salesforce/blip-image-captioning-base"
+blip_processor = AutoProcessor.from_pretrained(BLIP_MODEL_NAME)
+blip_model = BlipForConditionalGeneration.from_pretrained(BLIP_MODEL_NAME)
+
+
+def generate_image_caption(image_path: str) -> str:
+    image = Image.open(image_path).convert("RGB")
+    inputs = blip_processor(images=image, return_tensors="pt")
+    out = blip_model.generate(**inputs)
+    caption = blip_processor.decode(out[0], skip_special_tokens=True)
+    return caption