hskwon7/ISOM5240-Individual-Assignment

1# modules.py2 3"""4modules.py5 6Helper functions for the Image-to-Story Streamlit application.7Provides:8- Cached loaders for Hugging Face pipelines (captioning, story generation)9- Inference functions: generate_caption, generate_story_simple, generate_audio10"""11import streamlit as st12import re13from transformers import pipeline14from gtts import gTTS15import io16 17@st.cache_resource18def load_captioner():19    """Load and cache BLIP image captioning pipeline."""20    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")21 22@st.cache_resource23def load_story_gen():24    """Load and cache the genre-story-generator-v2 text-generation pipeline."""25    return pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")26 27def trim_to_sentence(text: str, max_words: int) -> str:28    """29    Trim the story to the last complete sentence under max_words.30    If no sentence fits, fallback to the first max_words words.31    """32    sentences = re.split(r'(?<=[.!?])\s+', text)33    trimmed = []34    count = 035    for s in sentences:36        wc = len(s.split())37        if count + wc <= max_words:38            trimmed.append(s)39            count += wc40        else:41            break42    if trimmed:43        return " ".join(trimmed)44    # fallback to naive word trim45    return " ".join(text.split()[:max_words])46 47def generate_caption(captioner, image) -> str:48    """Run the captioner pipeline on the PIL image."""49    raw = captioner(image)50    first = raw[0]51    return first.get("generated_text", "") if isinstance(first, dict) else str(first)52 53def generate_story_simple(storyteller, prompt_text: str,54                          min_words: int = 50, max_words: int = 100) -> str:55    """56    Generate a 50–100 word story:57      1. Sample ~120 tokens with nucleus sampling.58      2. If under min_words, re-sample ~200 tokens with higher top_p.59      3. Trim to last sentence under max_words.60    """61    out = storyteller(prompt_text, max_new_tokens=120,62                      do_sample=True, top_p=0.9, num_return_sequences=1)63    story = out[0]["generated_text"]64    if len(story.split()) < min_words:65        out = storyteller(prompt_text, max_new_tokens=200,66                          do_sample=True, top_p=0.95, num_return_sequences=1)67        story = out[0]["generated_text"]68    return trim_to_sentence(story, max_words)69 70def generate_audio(text: str) -> (bytes, str):71    """72    Convert text to MP3 bytes using gTTS.73    Returns (audio_bytes, mime_type) for use in st.audio(...).74    """75    tts = gTTS(text=text, lang="en")76    buf = io.BytesIO()77    tts.write_to_fp(buf)78    return buf.getvalue(), "audio/mp3"79