Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import sys | |
| st.write("Python executable being used:") | |
| st.write(sys.executable) | |
| from PIL import Image | |
| import base64 | |
| import requests | |
| import json | |
| import os | |
| import re | |
| import torch | |
| from peft import PeftModel, PeftConfig | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import argparse | |
| import io | |
| from utils.model_utils import get_model_caption | |
| from utils.image_utils import overlay_caption | |
| def load_models(): | |
| base_model = AutoModelForCausalLM.from_pretrained("google/gemma-2b") | |
| tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b") | |
| model_angry = PeftModel.from_pretrained(base_model, "NursNurs/outputs_gemma2b_angry") | |
| model_happy = PeftModel.from_pretrained(base_model, "NursNurs/outputs_gemma2b_happy") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| base_model.to(device) | |
| model_happy.to(device) | |
| model_angry.to(device) | |
| # Load the adapters for specific moods | |
| base_model.load_adapter("NursNurs/outputs_gemma2b_happy", "happy") | |
| base_model.load_adapter("NursNurs/outputs_gemma2b_angry", "angry") | |
| return base_model, tokenizer, model_happy, model_angry, device | |
| # x = st.slider('Select a value') | |
| # st.write(x, 'squared is', x * x) | |
| def generate_meme_from_image(img_path, base_model, tokenizer, hf_token, device='cuda'): | |
| caption = get_model_caption(img_path, base_model, tokenizer, hf_token) | |
| print(caption) | |
| image = overlay_caption(caption, img_path) | |
| return image, caption | |
| st.title("Image Upload and Processing App") | |
| def main(): | |
| st.title("Meme Generator with Mood") | |
| base_model, tokenizer, model_happy, model_angry, device = load_models() | |
| # Input widget to upload an image | |
| uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"]) | |
| # Input widget to add Hugging Face token | |
| hf_token = st.text_input("Enter your Hugging Face Token", type='default') | |
| # Dropdown to select mood | |
| # mood = st.selectbox("Select Mood", options=["happy", "angry"]) | |
| # Directory for saving the meme (optional, but you can let users set this if needed) | |
| output_dir = "results" | |
| if uploaded_image is not None and hf_token: | |
| # Convert uploaded image to a PIL image | |
| img = Image.open(uploaded_image) | |
| # Generate meme when button is pressed | |
| if st.button("Generate Meme"): | |
| with st.spinner('Generating meme...'): | |
| image, caption = generate_meme_from_image(img, base_model, tokenizer, hf_token, device) | |
| # Display the output | |
| st.image(image, caption=f"Generated Meme: {caption}") | |
| # # Optionally allow downloading the meme | |
| # buf = io.BytesIO() | |
| # image.save(buf, format="PNG") | |
| # byte_im = buf.getvalue() | |
| st.download_button( | |
| label="Download Image with Caption", | |
| data=image, | |
| file_name="captioned_image.jpg", | |
| mime="image/jpeg" | |
| ) | |
| if __name__ == '__main__': | |
| main() | |
| # # Upload the image | |
| # uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"]) | |
| # # Process and display if image is uploaded | |
| # if uploaded_image is not None: | |
| # image = Image.open(uploaded_image) | |
| # st.image(image, caption="Uploaded Image", use_column_width=True) |