From 6114e0b719846b4a97c6c70884a6c8c8c02cbeec Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 15:17:28 +0000 Subject: [PATCH 1/6] add requirements.txt --- requirements.txt | 99 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8974b28 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,99 @@ +aiohappyeyeballs==2.4.6 +aiohttp==3.11.12 +aiosignal==1.3.2 +annotated-types==0.7.0 +anyio==4.8.0 +appnope==0.1.4 +asttokens==3.0.0 +attrs==25.1.0 +certifi==2025.1.31 +charset-normalizer==3.4.1 +click==8.1.8 +comm==0.2.2 +contourpy==1.3.1 +cycler==0.12.1 +datasets==3.3.0 +debugpy==1.8.12 +decorator==5.1.1 +dill==0.3.8 +docker-pycreds==0.4.0 +executing==2.2.0 +faiss-cpu==1.10.0 +fastapi==0.115.8 +filelock==3.17.0 +fonttools==4.56.0 +frozenlist==1.5.0 +fsspec==2024.12.0 +gitdb==4.0.12 +GitPython==3.1.44 +h11==0.14.0 +huggingface-hub==0.28.1 +idna==3.10 +ipykernel==6.29.5 +ipython==8.32.0 +jedi==0.19.2 +Jinja2==3.1.5 +joblib==1.4.2 +jupyter_client==8.6.3 +jupyter_core==5.7.2 +kiwisolver==1.4.8 +MarkupSafe==3.0.2 +matplotlib==3.10.0 +matplotlib-inline==0.1.7 +mpmath==1.3.0 +multidict==6.1.0 +multiprocess==0.70.16 +nest-asyncio==1.6.0 +networkx==3.4.2 +nltk==3.9.1 +numpy==2.2.3 +packaging==24.2 +pandas==2.2.3 +parso==0.8.4 +pexpect==4.9.0 +pillow==11.1.0 +platformdirs==4.3.6 +prompt_toolkit==3.0.50 +propcache==0.2.1 +protobuf==5.29.3 +psutil==7.0.0 +ptyprocess==0.7.0 +pure_eval==0.2.3 +pyarrow==19.0.0 +pydantic==2.10.6 +pydantic_core==2.27.2 +Pygments==2.19.1 +pyparsing==3.2.1 +python-dateutil==2.9.0.post0 +python-multipart==0.0.20 +pytz==2025.1 +PyYAML==6.0.2 +pyzmq==26.2.1 +regex==2024.11.6 +requests==2.32.3 +safetensors==0.5.2 +seaborn==0.13.2 +sentry-sdk==2.22.0 +setproctitle==1.3.4 +setuptools==75.8.0 +six==1.17.0 +smmap==5.0.2 +sniffio==1.3.1 +stack-data==0.6.3 +starlette==0.45.3 +sympy==1.13.1 +tokenizers==0.21.0 +torch==2.6.0 +torchvision==0.21.0 +tornado==6.4.2 +tqdm==4.67.1 +traitlets==5.14.3 +transformers==4.48.3 +typing_extensions==4.12.2 +tzdata==2025.1 +urllib3==2.3.0 +uvicorn==0.34.0 +wandb==0.19.6 +wcwidth==0.2.13 +xxhash==3.5.0 +yarl==1.18.3 From 9ca65202c7b91bc9ebfb4eda72af1d4a1db31d82 Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 17:35:40 +0000 Subject: [PATCH 2/6] streamlit added --- requirements.txt | 13 ++- streamlit_app.py | 299 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 306 insertions(+), 6 deletions(-) create mode 100644 streamlit_app.py diff --git a/requirements.txt b/requirements.txt index 8974b28..7603b8e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -38,7 +38,7 @@ jupyter_client==8.6.3 jupyter_core==5.7.2 kiwisolver==1.4.8 MarkupSafe==3.0.2 -matplotlib==3.10.0 +matplotlib==3.8.2 matplotlib-inline==0.1.7 mpmath==1.3.0 multidict==6.1.0 @@ -46,12 +46,12 @@ multiprocess==0.70.16 nest-asyncio==1.6.0 networkx==3.4.2 nltk==3.9.1 -numpy==2.2.3 +numpy==1.26.3 packaging==24.2 pandas==2.2.3 parso==0.8.4 pexpect==4.9.0 -pillow==11.1.0 +pillow==10.2.0 platformdirs==4.3.6 prompt_toolkit==3.0.50 propcache==0.2.1 @@ -83,12 +83,12 @@ stack-data==0.6.3 starlette==0.45.3 sympy==1.13.1 tokenizers==0.21.0 -torch==2.6.0 -torchvision==0.21.0 +torch==2.2.0 +torchvision==0.17.0 tornado==6.4.2 tqdm==4.67.1 traitlets==5.14.3 -transformers==4.48.3 +transformers==4.37.2 typing_extensions==4.12.2 tzdata==2025.1 urllib3==2.3.0 @@ -97,3 +97,4 @@ wandb==0.19.6 wcwidth==0.2.13 xxhash==3.5.0 yarl==1.18.3 +streamlit==1.31.1 diff --git a/streamlit_app.py b/streamlit_app.py new file mode 100644 index 0000000..c0fa86e --- /dev/null +++ b/streamlit_app.py @@ -0,0 +1,299 @@ +import streamlit as st +import torch +import torch.nn as nn +from PIL import Image +import io +import transformers +from Transformer import Transformer, Decoder, DecoderLayer +from Attention import MultiHeadAttention +import matplotlib.pyplot as plt +import json +import logging + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +# Set page config +st.set_page_config( + page_title="Amirs Image Caption Generator", + page_icon="🖼️", + layout="wide" +) + +# Add custom CSS +st.markdown(""" + +""", unsafe_allow_html=True) + +@st.cache_resource +def load_model(): + device = torch.device("cpu") + CLIP = transformers.CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + processor = transformers.CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") + + tokenizer = processor.tokenizer + vocab_size = tokenizer.vocab_size + tokenizer.add_special_tokens({"pad_token": "<<>>"}) + vocab = tokenizer.get_vocab() + tokenizer = processor.tokenizer + vocab_size = tokenizer.vocab_size + vocab_size_final = vocab_size + 1 #add 1 here + reverse_vocab = {idx: token for token, idx in vocab.items()} + + # clip_text_model = CLIP.text_model + + # Model parameters + d_model = 512 + text_dimension_embedding = 512 + image_encoder_output_dim = 768 + n_loops = 6 + num_heads = 8 + hidden_dim = 2048 + + self_attn_layer = MultiHeadAttention(encoder_output_dim=d_model, decoder_dim=d_model, d_model=d_model, num_heads=num_heads) + cross_attn_layer = MultiHeadAttention(encoder_output_dim=image_encoder_output_dim, decoder_dim=text_dimension_embedding, d_model=d_model, num_heads=num_heads) + feed_forward = nn.Sequential(nn.Linear(d_model, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, d_model)) + + text_model = CLIP.text_model + text_embedder = text_model.embeddings + + decoder_layer = DecoderLayer( + input_dim=text_dimension_embedding, + tgt_vocab_size=vocab_size_final, + intermediate_attn_dim=d_model, + n_loops=n_loops, + feed_forward=feed_forward, + self_attn_layer=self_attn_layer, + cross_attn_layer=cross_attn_layer + ) + + decoder = Decoder( + vocab_size_final, + pad_token=tokenizer.pad_token_id, + embedding_layer=text_embedder, + layer=decoder_layer, + n_loops=n_loops, + d_model=d_model + ) + + transformer = Transformer( + d_model=d_model, + text_encoder=text_embedder, + image_encoder=CLIP.vision_model, + decoder=decoder, + tgt_vocab_size=vocab_size_final, + pad_token=tokenizer.pad_token_id + ) + + checkpoint = torch.load("checkpoints/best_model_sat.pt", map_location=torch.device("cpu")) + transformer.load_state_dict(checkpoint['model_state_dict']) + + return transformer, processor, CLIP, reverse_vocab, tokenizer + +def process_image(image): + if image.mode in ('RGBA', 'LA'): + background = Image.new('RGB', image.size, (255, 255, 255)) + if image.mode == 'RGBA': + background.paste(image, mask=image.split()[3]) + else: + background.paste(image, mask=image.split()[1]) + image = background + elif image.mode != 'RGB': + image = image.convert('RGB') + return image + +def generate_captions(transformer, data, k, reverse_vocab): + logging.info(f"Starting generate_captions with k={k}") + transformer.eval() + MAX_SEQ_LENGTH = 35 + START_TOKEN_ID = 49406 + BEAM_WIDTH = k + + captions = [] # Store captions in a list instead of yielding + # Initialize beam with start token + current_sequences = torch.full((BEAM_WIDTH, 1), START_TOKEN_ID) + + # Generate first token with beam search + print(data, "data") + output = transformer.forward(data['pixel_values'].repeat(BEAM_WIDTH, 1, 1, 1), current_sequences) + logging.info('eyo2') + output_probabilities = torch.softmax(output, dim=2) + top_k_values, top_k_indices = torch.topk(output_probabilities[0, 0], k=BEAM_WIDTH) + print('does it reach') + # Create BEAM_WIDTH different sequences + for beam_idx in range(BEAM_WIDTH): + current_sequence = torch.full((1, 1), START_TOKEN_ID) + predicted_indices = torch.zeros((1, 1)) + + # Use the beam_idx-th best first token + first_token = top_k_indices[beam_idx].item() + current_sequence = torch.cat((current_sequence, torch.tensor([first_token]).unsqueeze(0)), dim=1) + predicted_indices[0, 0] = first_token + + # Continue generating remaining tokens + for pos in range(1, MAX_SEQ_LENGTH): + output = transformer.forward(data['pixel_values'], current_sequence) + output_probabilities = torch.softmax(output, dim=2) + predicted_digit = torch.argmax(output_probabilities[0, pos]) + + if pos < MAX_SEQ_LENGTH - 1: + current_sequence = torch.cat((current_sequence, torch.tensor([predicted_digit.item()]).unsqueeze(0)), dim=1) + + predicted_indices = torch.cat((predicted_indices, predicted_digit.unsqueeze(0).unsqueeze(0)), dim=1) + + # Convert to text and yield immediately + predicted_digits = [reverse_vocab[idx.item()] for idx in predicted_indices.squeeze(0)] + cleaned_text = ' '.join(token.replace('', '') for token in predicted_digits + if token not in ['<|endoftext|>', '<<>>']) + logging.info(f"Generated caption: {cleaned_text}") + captions.append(cleaned_text) # Append instead of yield + + return captions # Return the list of captions + +def get_best_caption(image, captions, model, processor): + print(captions, "xxxxx", image) + + image_inputs = processor(images=image, return_tensors="pt") + text_inputs = processor.tokenizer( + captions, + padding=True, + truncation=True, + max_length=77, + return_tensors="pt", + add_special_tokens=True # Ensure special tokens are added + ) + + # Adjust input_ids to account for the extra token + input_ids = text_inputs["input_ids"] + # Shift indices greater than or equal to the PAD token id up by 1 + pad_token_id = processor.tokenizer.pad_token_id + input_ids[input_ids >= pad_token_id] += 1 + + # Combine inputs + inputs = { + "input_ids": input_ids, + "attention_mask": text_inputs["attention_mask"], + "pixel_values": image_inputs["pixel_values"] + } + + with torch.no_grad(): + outputs = model(**inputs) + + logits_per_caption = outputs.logits_per_text + best_caption_idx = torch.argmax(logits_per_caption).item() + + return captions[best_caption_idx] + +def main(): + st.title("🖼️ Amir's Image Caption Generator") + st.write("Upload an image and get AI-generated captions!") + + # Add GitHub link + st.markdown("[View on GitHub](https://github.com/Amirjab21/MLXW4)", unsafe_allow_html=True) + + + # Load models + with st.spinner("Loading models..."): + transformer, processor, clip_model, reverse_vocab, tokenizer = load_model() + + # File uploader + uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) + + if uploaded_file is not None: + # Create columns + col1, col2 = st.columns(2) + + with col1: + st.subheader("Uploaded Image") + image = Image.open(uploaded_file) + st.image(image, use_container_width=True) + + with col2: + st.subheader("Generated Captions") + # Create a placeholder for captions + captions_placeholder = st.empty() + + # Initialize an empty list to store captions + captions = [] + + # Process the image + processed_image = process_image(image) + data = processor(images=processed_image, return_tensors="pt") + + # Generate captions with streaming + transformer.eval() + MAX_SEQ_LENGTH = 35 + START_TOKEN_ID = 49406 + BEAM_WIDTH = 5 + + # Initialize beam with start token + current_sequences = torch.full((BEAM_WIDTH, 1), START_TOKEN_ID) + + # Generate first token with beam search + output = transformer.forward(data['pixel_values'].repeat(BEAM_WIDTH, 1, 1, 1), current_sequences) + output_probabilities = torch.softmax(output, dim=2) + top_k_values, top_k_indices = torch.topk(output_probabilities[0, 0], k=BEAM_WIDTH) + + # Create progress bar + progress_text = "Generating captions..." + caption_progress = st.progress(0) + + # Create BEAM_WIDTH different sequences + for beam_idx in range(BEAM_WIDTH): + current_sequence = torch.full((1, 1), START_TOKEN_ID) + predicted_indices = torch.zeros((1, 1)) + + # Use the beam_idx-th best first token + first_token = top_k_indices[beam_idx].item() + current_sequence = torch.cat((current_sequence, torch.tensor([first_token]).unsqueeze(0)), dim=1) + predicted_indices[0, 0] = first_token + + # Continue generating remaining tokens + for pos in range(1, MAX_SEQ_LENGTH): + output = transformer.forward(data['pixel_values'], current_sequence) + output_probabilities = torch.softmax(output, dim=2) + predicted_digit = torch.argmax(output_probabilities[0, pos]) + if predicted_digit.item() in [tokenizer.eos_token_id, tokenizer.pad_token_id]: + break + if pos < MAX_SEQ_LENGTH - 1: + current_sequence = torch.cat((current_sequence, torch.tensor([predicted_digit.item()]).unsqueeze(0)), dim=1) + + predicted_indices = torch.cat((predicted_indices, predicted_digit.unsqueeze(0).unsqueeze(0)), dim=1) + + # Convert to text and display immediately + predicted_digits = [reverse_vocab[idx.item()] for idx in predicted_indices.squeeze(0)] + cleaned_text = ' '.join(token.replace('', '') for token in predicted_digits + if token not in ['<|endoftext|>', '<<>>']) + + # Add the new caption to the list + captions.append(cleaned_text) + + # Update the display with all captions so far + caption_html = "" + for idx, caption in enumerate(captions, 1): + caption_html += f'
{idx}. {caption}
' + captions_placeholder.markdown(caption_html, unsafe_allow_html=True) + + # Update progress bar + caption_progress.progress((beam_idx + 1) / BEAM_WIDTH) + + # Clear the progress bar when done + caption_progress.empty() + + # Final update with all captions + st.success("Caption generation complete!") + +if __name__ == "__main__": + main() \ No newline at end of file From 90e0e952d54be7ee280ffd9d689274c4837c2738 Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 17:39:26 +0000 Subject: [PATCH 3/6] tokenizer requirement update --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 7603b8e..070ad9a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -82,7 +82,7 @@ sniffio==1.3.1 stack-data==0.6.3 starlette==0.45.3 sympy==1.13.1 -tokenizers==0.21.0 +tokenizers==0.18.0 torch==2.2.0 torchvision==0.17.0 tornado==6.4.2 From 0f4a5335b0aab5220288916a9faf589377844251 Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 17:41:28 +0000 Subject: [PATCH 4/6] new req --- requirements.txt | 114 ++++++----------------------------------------- 1 file changed, 14 insertions(+), 100 deletions(-) diff --git a/requirements.txt b/requirements.txt index 070ad9a..1638bf2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,100 +1,14 @@ -aiohappyeyeballs==2.4.6 -aiohttp==3.11.12 -aiosignal==1.3.2 -annotated-types==0.7.0 -anyio==4.8.0 -appnope==0.1.4 -asttokens==3.0.0 -attrs==25.1.0 -certifi==2025.1.31 -charset-normalizer==3.4.1 -click==8.1.8 -comm==0.2.2 -contourpy==1.3.1 -cycler==0.12.1 -datasets==3.3.0 -debugpy==1.8.12 -decorator==5.1.1 -dill==0.3.8 -docker-pycreds==0.4.0 -executing==2.2.0 -faiss-cpu==1.10.0 -fastapi==0.115.8 -filelock==3.17.0 -fonttools==4.56.0 -frozenlist==1.5.0 -fsspec==2024.12.0 -gitdb==4.0.12 -GitPython==3.1.44 -h11==0.14.0 -huggingface-hub==0.28.1 -idna==3.10 -ipykernel==6.29.5 -ipython==8.32.0 -jedi==0.19.2 -Jinja2==3.1.5 -joblib==1.4.2 -jupyter_client==8.6.3 -jupyter_core==5.7.2 -kiwisolver==1.4.8 -MarkupSafe==3.0.2 -matplotlib==3.8.2 -matplotlib-inline==0.1.7 -mpmath==1.3.0 -multidict==6.1.0 -multiprocess==0.70.16 -nest-asyncio==1.6.0 -networkx==3.4.2 -nltk==3.9.1 -numpy==1.26.3 -packaging==24.2 -pandas==2.2.3 -parso==0.8.4 -pexpect==4.9.0 -pillow==10.2.0 -platformdirs==4.3.6 -prompt_toolkit==3.0.50 -propcache==0.2.1 -protobuf==5.29.3 -psutil==7.0.0 -ptyprocess==0.7.0 -pure_eval==0.2.3 -pyarrow==19.0.0 -pydantic==2.10.6 -pydantic_core==2.27.2 -Pygments==2.19.1 -pyparsing==3.2.1 -python-dateutil==2.9.0.post0 -python-multipart==0.0.20 -pytz==2025.1 -PyYAML==6.0.2 -pyzmq==26.2.1 -regex==2024.11.6 -requests==2.32.3 -safetensors==0.5.2 -seaborn==0.13.2 -sentry-sdk==2.22.0 -setproctitle==1.3.4 -setuptools==75.8.0 -six==1.17.0 -smmap==5.0.2 -sniffio==1.3.1 -stack-data==0.6.3 -starlette==0.45.3 -sympy==1.13.1 -tokenizers==0.18.0 -torch==2.2.0 -torchvision==0.17.0 -tornado==6.4.2 -tqdm==4.67.1 -traitlets==5.14.3 -transformers==4.37.2 -typing_extensions==4.12.2 -tzdata==2025.1 -urllib3==2.3.0 -uvicorn==0.34.0 -wandb==0.19.6 -wcwidth==0.2.13 -xxhash==3.5.0 -yarl==1.18.3 -streamlit==1.31.1 +tqdm +transformers +torch +torchvision +numpy +pillow +tokenizers +streamlit +matplotlib +datasets +pandas +nltk +sentencepiece +wandb \ No newline at end of file From af33ca1d1f91b6d8653d2739330152459c18f8ab Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 18:15:35 +0000 Subject: [PATCH 5/6] add gitattributes for git-lfs --- .gitattributes | 1 + 1 file changed, 1 insertion(+) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..e92cf9b --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.pt filter=lfs diff=lfs merge=lfs -text From b47dec082274bfcdf19a3fa91cfe926ce6dc9cfa Mon Sep 17 00:00:00 2001 From: Amir Jabarivasal Date: Sun, 23 Feb 2025 18:22:58 +0000 Subject: [PATCH 6/6] add model.pt file --- checkpoints/best_model_sat.pt | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 checkpoints/best_model_sat.pt diff --git a/checkpoints/best_model_sat.pt b/checkpoints/best_model_sat.pt new file mode 100644 index 0000000..12b5222 --- /dev/null +++ b/checkpoints/best_model_sat.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49e445f48e8c35b7497338e1010222a0b0ad4a060ae9ab19d593d2a69615860b +size 1026684016