Spaces:
Sleeping
Sleeping
| # app.py — Story Generator with Elegant UI | |
| import os, json, re, pathlib, base64, time, uuid | |
| from huggingface_hub import InferenceClient | |
| from PIL import Image | |
| import gradio as gr | |
| # ---------- Config ---------- | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| if not HF_TOKEN: | |
| raise RuntimeError("⚠️ Set HF_TOKEN environment variable (use Spaces Secrets).") | |
| CHAT_MODEL = "meta-llama/Llama-3.1-8b-instruct" | |
| IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell" | |
| SAFETY_MODEL = "meta-llama/Meta-Llama-Guard-2-8B" # Safety model for content moderation | |
| OUT_DIR = pathlib.Path("/tmp/generated") | |
| OUT_DIR.mkdir(parents=True, exist_ok=True) | |
| client = InferenceClient(api_key=HF_TOKEN, provider="auto") | |
| # ---------- Utility ---------- | |
| def try_chat_completion(model_id, messages, max_tokens=3000): | |
| r = client.chat.completions.create(model=model_id, messages=messages, max_tokens=max_tokens) | |
| return r.choices[0].message.get("content", "") | |
| def extract_json_from_text(text): | |
| if not text: | |
| return None | |
| match = re.search(r'```json\s*(\{[\s\S]*\})\s*```', text) | |
| if match: | |
| json_str = match.group(1) | |
| else: | |
| match = re.search(r'\{[\s\S]*\}', text) | |
| if not match: | |
| return None | |
| json_str = match.group(0) | |
| try: | |
| return json.loads(json_str) | |
| except json.JSONDecodeError as e: | |
| print(f"--- JSON PARSING FAILED ---\nError: {e}\nContent: {json_str}\n--------------------------") | |
| return None | |
| def hf_text_to_image(model_id, prompt, out_path): | |
| img = client.text_to_image(prompt=prompt, model=model_id) | |
| if isinstance(img, Image.Image): | |
| img.save(out_path) | |
| return out_path | |
| raise RuntimeError("Invalid image response") | |
| def is_content_inappropriate(text_to_check): | |
| """Uses Llama Guard to check for inappropriate content.""" | |
| try: | |
| r = client.chat.completions.create( | |
| model=SAFETY_MODEL, | |
| messages=[{"role": "user", "content": f"Please evaluate if the following content is safe or unsafe based on typical safety guidelines for an AI assistant. Output 'safe' or 'unsafe'.\n\nContent: \"{text_to_check}\""}], | |
| max_tokens=20, | |
| temperature=0.1 | |
| ) | |
| response = r.choices[0].message.get("content", "").lower() | |
| return "unsafe" in response | |
| except Exception as e: | |
| print(f"Error in safety check: {e}") | |
| return False | |
| # ---------- Story logic ---------- | |
| def make_prompt(user_prompt, nscenes=6, nsent=5): | |
| # This prompt is excellent and requires no changes. | |
| return f""" | |
| You are a creative story writer. Your task is to write a compelling story based on a user's prompt. | |
| You MUST return the story in a single, valid JSON object. Do not write any text or explanations outside of the JSON structure. | |
| Here is an example of the required JSON format: | |
| {{ | |
| "title": "A descriptive title for the entire story", | |
| "scenes": [ | |
| {{ | |
| "id": 1, | |
| "text": "The full story text for this scene. This should be a complete paragraph with around {nsent} sentences, describing the events and setting.", | |
| "visual_prompt": "A detailed, vivid description for an image generation model, capturing the key visual elements of this scene." | |
| }} | |
| ] | |
| }} | |
| Please use the following details for the story: | |
| - Story Prompt: "{user_prompt}" | |
| - Total number of scenes: {nscenes} | |
| Now, generate the story in the specified JSON format. | |
| """ | |
| def generate_story_and_images(prompt, nscenes, nsent, img_model): | |
| start = time.time() | |
| logs = [] | |
| logs.append("🎬 Generating story...") | |
| raw = try_chat_completion(CHAT_MODEL, [{"role": "user", "content": make_prompt(prompt, nscenes, nsent)}]) | |
| story = extract_json_from_text(raw) | |
| if not story: | |
| story = {"title": "Untitled", "scenes": [{"id": i + 1, "text": f"Scene {i+1}. The AI failed to generate a proper story, or the JSON was malformed.", "visual_prompt": prompt} for i in range(nscenes)]} | |
| logs.append("⚠️ Failed to parse story JSON, using fallback.") | |
| else: | |
| logs.append("✅ Story JSON parsed.") | |
| image_paths = [] | |
| for s in story["scenes"]: | |
| visual_prompt = s.get("visual_prompt", s.get("text", prompt)) | |
| if is_content_inappropriate(visual_prompt): | |
| gr.Warning(f"Visual prompt for Scene {s['id']} was moderated for safety. Generating a default image.") | |
| visual_prompt = "A serene landscape with gentle colors." | |
| name = OUT_DIR / f"{uuid.uuid4().hex[:6]}_scene_{s['id']}.png" | |
| logs.append(f"🎨 Generating image for Scene {s['id']}...") | |
| hf_text_to_image(img_model, visual_prompt, str(name)) | |
| image_paths.append(str(name)) | |
| total = time.time() - start | |
| logs.append(f"✨ Done in {total:.2f}s") | |
| return story, image_paths, "\n".join(logs) | |
| # ---------- UI ---------- | |
| def build_ui(): | |
| css = """ | |
| .main-container { max-width: 1400px; margin: 0 auto; } | |
| .prompt-section { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 32px; border-radius: 16px; margin-bottom: 24px; } | |
| .prompt-box textarea { font-size: 16px !important; } | |
| .story-panel { background: rgba(255,255,255,0.05); padding: 24px; border-radius: 12px; backdrop-filter: blur(10px); border: 1px solid rgba(255,255,255,0.1); max-height: 800px; overflow-y: auto; } | |
| .story-title { font-size: 32px; font-weight: 700; margin-bottom: 24px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent; } | |
| .story-content { line-height: 1.8; font-size: 16px; color: rgba(255,255,255,0.9); } | |
| .story-content p { margin-bottom: 16px; } | |
| """ | |
| with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Story Generator") as demo: | |
| gr.Markdown("# 📚 AI Story Generator", elem_classes="main-title") | |
| with gr.Column(elem_classes="main-container"): | |
| with gr.Column(elem_classes="prompt-section"): | |
| prompt_box = gr.Textbox(label="✨ Enter your story idea", placeholder="e.g. A pirate discovering a hidden island...", lines=3, elem_classes="prompt-box") | |
| with gr.Row(): | |
| generate_btn = gr.Button("🚀 Generate Story", variant="primary", size="lg", scale=3) | |
| with gr.Column(scale=1): | |
| with gr.Accordion("⚙️ Settings", open=False): | |
| nscenes = gr.Slider(2, 12, value=6, step=1, label="📖 Number of scenes") | |
| nsent = gr.Slider(2, 8, value=5, step=1, label="📝 Sentences per scene") | |
| img_model = gr.Dropdown(choices=[IMAGE_MODEL], value=IMAGE_MODEL, label="🎨 Image model") | |
| log_box = gr.Textbox(label="📋 Generation Logs", lines=6, interactive=False) | |
| with gr.Row(): | |
| with gr.Column(scale=5, elem_classes="story-panel"): | |
| story_html = gr.HTML("<div style='text-align:center;padding:40px;color:#888;'>Your story will appear here...<br><br>Click 'Generate Story' to begin! ✨</div>") | |
| with gr.Column(scale=7): | |
| image_gallery = gr.Gallery( | |
| label="📷 Scene Visuals", show_label=False, elem_id="gallery", | |
| columns=2, object_fit="cover", height="auto", preview=True | |
| ) | |
| def on_generate(prompt, nscenes, nsent, img_model): | |
| if is_content_inappropriate(prompt): | |
| gr.Warning("Your prompt seems to violate the safety policy. Please try again.") | |
| return "<div style='text-align:center;padding:40px;color:#888;'>Prompt rejected due to safety policy.</div>", [], "Prompt rejected by safety filter." | |
| story, imgs, logs = generate_story_and_images(prompt, int(nscenes), int(nsent), img_model) | |
| story_output_html = f"<div class='story-title'>{story.get('title', 'Untitled')}</div>\n<div class='story-content'>\n" | |
| for s in story.get('scenes', []): | |
| scene_text = s.get('text', '') | |
| if is_content_inappropriate(scene_text): | |
| scene_text = f"**[Scene {s['id']} was moderated for safety and replaced with a placeholder.]**" | |
| gr.Info(f"Scene {s['id']} content was flagged and replaced.") | |
| story_output_html += f"<p>{scene_text}</p>\n\n" | |
| story_output_html += "</div>" | |
| return story_output_html, imgs, logs | |
| generate_btn.click( | |
| on_generate, | |
| inputs=[prompt_box, nscenes, nsent, img_model], | |
| outputs=[story_html, image_gallery, log_box] | |
| ) | |
| return demo | |
| app = build_ui() | |
| if __name__ == "__main__": | |
| app.launch() |