tti.py

import gradio as gr
from langchain_groq import ChatGroq
import nltk
import requests
from PIL import Image
import io
import time
from gtts import gTTS
from dotenv import load_dotenv
import os

# **NLTK Setup**
nltk.download('punkt')
nltk.download('punkt_tab')
# Load .env file
load_dotenv()

# Retrieve API keys from environment variables
GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # e.g. export GROQ_API_KEY="your-groq-key"
HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN") 
# **Initialize ChatGroq LLM**
llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=GROQ_API_KEY  # Replace with your Groq API key
)

# **Stable Diffusion XL API Setup**
API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
headers_sdxl = {"Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}"}  # Replace with your Hugging Face API token
API_COOLDOWN = 3  # Seconds between API calls

def query_sdxl(prompt):
    """Query the Stable Diffusion XL API to generate an image from a prompt."""
    payload = {"inputs": prompt, "options": {"wait_for_model": True}}
    try:
        r = requests.post(API_URL, headers=headers_sdxl, json=payload)
        if r.status_code == 200:
            return r.content
        elif r.status_code == 429:
            wait = int(r.headers.get("Retry-After", API_COOLDOWN))
            print(f"Rate limited—waiting {wait}s…")
            time.sleep(wait)
            return query_sdxl(prompt)
        else:
            print(f"SDXL API error {r.status_code}: {r.text}")
    except Exception as e:
        print("Request failed:", e)
    return None

def generate_storyboard(story):
    """Generate a storyboard with images and voice-overs from a user's story."""
    # Split story into sentences (up to 10 frames)
    sentences = nltk.sent_tokenize(story)[:10]
    num_frames = len(sentences)

    # Build few-shot messages for ChatGroq
    messages = [
        ("system", f"You are a helpful assistant that rewrites a story into frame-by-frame prompts for image generation. "
                   f"Generate EXACTLY {num_frames} frames - no more, no less. "
                   "Ensure that each frame is a continuation of the previous one, maintaining consistency in characters, objects, colors, clothing, environment, lighting, and other visual elements. "
                   "Avoid vague references—refer back to previously introduced characters and objects explicitly. "
                   "Make sure the setting and visual context flow smoothly from one frame to the next."),
        # Example 1
        ("human", "A girl in a red dress walks into a forest.\n"
                  "She sees a white rabbit hopping past.\n"
                  "Curious, she follows the rabbit deeper into the woods.\n"
                  "She stumbles upon a glowing cave entrance.\n"
                  "The girl steps into the cave, her red dress glowing under the blue light."),
        ("assistant", "Frame 1: A young girl wearing a red dress walks into a dense forest surrounded by tall trees.\n"
                      "Frame 2: The young girl wearing a red dress walks, a small white rabbit hops past her feet on the dense forest path.\n"
                      "Frame 3: The young girl in the red dress follows the white rabbit, moving deeper into the darker parts of the dense forest.\n"
                      "Frame 4: The young girl in the red dress reaches a glowing blue cave entrance hidden between mossy rocks and trees in the dense forest.\n"
                      "Frame 5: Inside the glowing blue cave in the dense forest, the young girl's red dress softly illuminates the surroundings."),
        # Example 2
        ("human", "A boy flies a kite in a windy field.\n"
                  "The kite gets tangled in a tree.\n"
                  "He climbs the tree to untangle it.\n"
                  "Suddenly, it starts to rain.\n"
                  "The boy holds his kite and runs home drenched."),
        ("assistant", "Frame 1: A boy joyfully flies a colorful kite in a wide, windy green field.\n"
                      "Frame 2: A boy watches the kite gets tangled in the branches of a tall tree nearby in the windy green field.\n"
                      "Frame 3: Climbing up the tree carefully, the boy reaches out to untangle the kite in the windy green field.\n"
                      "Frame 4: Dark clouds roll in as rain begins to fall, soaking the boy on the tree in the windy green field.\n"
                      "Frame 5: Holding the damp kite, the boy runs across the field, drenched and smiling."),
        ("human", "\n".join(sentences))
    ]

    # Generate frame prompts with ChatGroq
    ai_msg = llm.invoke(messages)
    frames_txt = ai_msg.content.strip()

    # Extract prompts
    prompts = [
        line.partition(":")[2].strip()
        for line in frames_txt.splitlines()
        if line.lower().startswith("frame")
    ]
    prompts = prompts[:num_frames]

    # Generate images and voice-overs
    image_paths = []
    audio_paths = []
    for idx, prompt in enumerate(prompts, start=1):
        # Generate voice-over
        tts = gTTS(text=prompt, lang='en')
        audio_path = f"frame_{idx}.mp3"
        tts.save(audio_path)
        audio_paths.append(audio_path)

        # Generate image
        img_bytes = query_sdxl(prompt)
        if img_bytes:
            img = Image.open(io.BytesIO(img_bytes))
            img_path = f"frame_{idx}.png"
            img.save(img_path)
            image_paths.append(img_path)
        else:
            image_paths.append(None)  # Placeholder for failed images

        # Respect API cooldown
        if idx < len(prompts):
            time.sleep(API_COOLDOWN)

    # Prepare Gradio updates
    image_updates = gr.update(value=image_paths)
    audio_updates = [
        gr.update(value=audio_paths[i], visible=True) if i < num_frames else gr.update(value=None, visible=False)
        for i in range(10)
    ]
    return [image_updates] + audio_updates

# **Gradio Interface**
with gr.Blocks(title="Storyboard Generator") as demo:
    gr.Markdown("# Storyboard Generator\nEnter a story (up to 10 sentences) to generate a storyboard with images and voice-overs.")
    story_input = gr.Textbox(label="Enter your story", lines=5, placeholder="Type your story here...")
    generate_btn = gr.Button("Generate Storyboard")
    image_gallery = gr.Gallery(label="Storyboard Images", show_label=True)
    with gr.Column():
        audios = [gr.Audio(label=f"Frame {i+1} Voice-Over", visible=False) for i in range(10)]

    generate_btn.click(
        fn=generate_storyboard,
        inputs=story_input,
        outputs=[image_gallery] + audios
    )

# **Launch the App**
demo.launch()