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()