Create app.py

app.py

@@ -0,0 +1,187 @@
+# 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()