Ai-Instagram-Caption-Generator/utils.py at main · rachit23tech/Ai-Instagram-Caption-Generator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
from transformers import BlipProcessor, BlipForConditionalGeneration
from sentence_transformers import SentenceTransformer, util
from groq import Groq
import streamlit as st
import os
import json
import random


# -----------------------------
# Load BLIP Vision Model
# -----------------------------
@st.cache_resource
def load_blip():
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
    return processor, model


# -----------------------------
# Load CLIP Model
# -----------------------------
@st.cache_resource
def load_clip():
    return SentenceTransformer("clip-ViT-B-32")


processor, model = load_blip()
clip_model = load_clip()


# -----------------------------
# Groq Client
# -----------------------------
api_key = os.getenv("GROQ_API_KEY")

if not api_key:
    raise ValueError("GROQ_API_KEY not found. Set it in environment variables.")

client = Groq(api_key=api_key)


# -----------------------------
# Generate image description
# -----------------------------
def generate_image_description(image):

    inputs = processor(image, return_tensors="pt")

    output = model.generate(**inputs, max_length=40)

    description = processor.decode(output[0], skip_special_tokens=True)

    return description


# -----------------------------
# Generate caption package (1 API call)
# -----------------------------
def generate_caption_package(description, style="Instagram", num_captions=3):

    prompt = f"""
Image description:
{description}

Generate {num_captions} {style} Instagram captions.

Caption style guidelines:
Instagram → trendy and modern
Funny → humorous and playful
Savage → bold and confident
Aesthetic → minimal and stylish
Motivational → inspiring
Poetic → artistic

For each caption return:
- caption
- emoji_caption
- Punjabi translation
- viral_score (1-10)

Return JSON list only.
"""

    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.1-8b-instant",
    )

    text = response.choices[0].message.content

    # Extract JSON safely
    try:
        start = text.find("[")
        end = text.rfind("]") + 1
        json_text = text[start:end]
        data = json.loads(json_text)
    except:
        data = []

    return data


# -----------------------------
# Caption ranking using CLIP
# -----------------------------
def rank_captions(description, captions):

    if not captions:
        return []

    caption_texts = [c["caption"] for c in captions]

    desc_embedding = clip_model.encode(description, convert_to_tensor=True)
    caption_embeddings = clip_model.encode(caption_texts, convert_to_tensor=True)

    scores = util.cos_sim(desc_embedding, caption_embeddings)[0]

    ranked = sorted(
        zip(captions, scores),
        key=lambda x: x[1],
        reverse=True
    )

    ranked_captions = [c[0] for c in ranked]

    return ranked_captions


# -----------------------------
# Simple hashtag generator
# -----------------------------
def generate_hashtags(caption):

    words = caption.split()

    tags = []

    for word in words:
        if len(word) > 4:
            tags.append("#" + word.lower())

    extra_tags = [
        "#style",
        "#fashion",
        "#instagood",
        "#photography",
        "#dailylook",
        "#trending",
        "#aesthetic"
    ]

    tags.extend(random.sample(extra_tags, 3))

    return list(set(tags))[:6]