Update app.py

app.py

@@ -4,16 +4,15 @@ from transformers import pipeline
 import os
 
 # --- App Configuration ---
-TITLE = "✍️ AI Story Weaver"
+TITLE = "✍️ AI Story Outliner"
 DESCRIPTION = """
-Enter a prompt, a topic, or the beginning of a story, and get three different continuations from powerful open-source AI models.
-This app uses:
-- **Mistral-7B-Instruct-v0.2**
-- **Google's Gemma-7B-IT**
-- **Meta's Llama-3-8B-Instruct**
-
-**⚠️ Hardware Warning:** These are very large models. Loading them requires a powerful GPU with significant VRAM (ideally > 24GB).
-The initial loading process may take several minutes. You will also need to install the `accelerate` library: `pip install accelerate`
+Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
+The app uses **TinyLlama-1.1B** to generate creative outlines formatted in Markdown.
+
+**How it works:**
+1.  Enter your story idea.
+2.  The AI will generate 10 different story outlines.
+3.  Each outline has a dramatic beginning and is concise, like a song.
 """
 
 # --- Example Prompts for Storytelling ---
@@ -26,52 +25,84 @@ examples = [
 ]
 
 # --- Model Initialization ---
-# This section loads the models. It requires significant hardware resources.
-# `device_map="auto"` and `torch_dtype="auto"` help manage resources by using available GPUs and half-precision.
+# This section loads a smaller, CPU-friendly model.
+# It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
 try:
-    print("Initializing models... This may take several minutes.")
+    print("Initializing model... This may take a moment.")
     
-    # NOTE: For Llama-3, you may need to log in to Hugging Face and accept the license agreement.
-    # from huggingface_hub import login
-    # login("YOUR_HF_TOKEN")
-
-    generator1 = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", torch_dtype="auto", device_map="auto")
-    print("✅ Mistral-7B loaded.")
+    # Load the token from environment variables if it exists (for HF Spaces secrets)
+    hf_token = os.environ.get("HF_TOKEN", None)
     
-    generator2 = pipeline("text-generation", model="google/gemma-7b-it", torch_dtype="auto", device_map="auto")
-    print("✅ Gemma-7B loaded.")
-
-    generator3 = pipeline("text-generation", model="meta-llama/Llama-3-8B-Instruct", torch_dtype="auto", device_map="auto")
-    print("✅ Llama-3-8B loaded.")
-
-    print("All models loaded successfully! 🎉")
+    # Using a smaller model that is more suitable for running without a high-end GPU.
+    generator = pipeline(
+        "text-generation",
+        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        torch_dtype=torch.bfloat16, # More efficient dtype
+        device_map="auto", # Will use GPU if available, otherwise CPU
+        token=hf_token
+    )
+    print("✅ TinyLlama model loaded successfully!")
 
 except Exception as e:
     print(f"--- 🚨 Error loading models ---")
     print(f"Error: {e}")
-    print("Please ensure you have 'torch' and 'accelerate' installed, have sufficient VRAM, and are logged into Hugging Face if required.")
     # Create a dummy function if models fail, so the app can still launch with an error message.
     def failed_generator(prompt, **kwargs):
-        return [{'generated_text': "A model failed to load. Please check the console for errors. You may need more VRAM or need to accept model license terms on Hugging Face."}]
-    generator1 = generator2 = generator3 = failed_generator
+        error_message = f"Model failed to load. Please check the console for errors. Error: {e}"
+        return [{'generated_text': error_message}]
+    generator = failed_generator
 
 
 # --- App Logic ---
-def generate_stories(prompt: str) -> tuple[str, str, str]:
-    """Generates text from the three loaded models based on the user's prompt."""
+def generate_stories(prompt: str) -> list[str]:
+    """
+    Generates 10 story outlines from the loaded model based on the user's prompt.
+    """
     if not prompt:
-        return "Please enter a prompt to start.", "", ""
-        
-    # We use 'max_new_tokens' to control the length of the generated story.
-    # Increased to 200 for more substantial story continuations.
-    params = {"max_new_tokens": 200, "do_sample": True, "temperature": 0.7, "top_p": 0.95}
-    
-    # Generate from all three models
-    out1 = generator1(prompt, **params)[0]['generated_text']
-    out2 = generator2(prompt, **params)[0]['generated_text']
-    out3 = generator3(prompt, **params)[0]['generated_text']
+        # Return a list of 10 empty strings to clear the outputs
+        return [""] * 10
+
+    # A detailed system prompt to guide the model's output format and structure.
+    system_prompt = f"""
+<|system|>
+You are an expert storyteller. Your task is to take a user's prompt and write 
+a short story as a Markdown outline. The story must have a dramatic arc and be 
+the length of a song. Use emojis to highlight the story sections.
+
+**Your Story Outline Structure:**
+- 🎬 **The Hook:** A dramatic opening.
+- 🎼 **The Ballad:** The main story, told concisely.
+- 🔚 **The Finale:** A clear and satisfying ending.</s>
+<|user|>
+{prompt}</s>
+<|assistant|>
+"""
+
+    # Parameters for the pipeline to generate 10 diverse results.
+    params = {
+        "max_new_tokens": 250,
+        "num_return_sequences": 10,
+        "do_sample": True,
+        "temperature": 0.8,
+        "top_k": 50,
+        "top_p": 0.95,
+    }
     
-    return out1, out2, out3
+    # Generate 10 different story variations
+    outputs = generator(system_prompt, **params)
+
+    # Extract the generated text and clean it up.
+    stories = []
+    for out in outputs:
+        # Remove the system prompt from the beginning of the output
+        cleaned_text = out['generated_text'].replace(system_prompt, "").strip()
+        stories.append(cleaned_text)
+
+    # Ensure we return exactly 10 stories, padding with an error message if necessary.
+    while len(stories) < 10:
+        stories.append("Failed to generate a story for this slot.")
+
+    return stories
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !important;}") as demo:
@@ -85,16 +116,22 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !i
                 label="Your Story Prompt 👇",
                 placeholder="e.g., 'The last dragon on Earth lived not in a cave, but in a library...'"
             )
-            generate_button = gr.Button("Weave a Story ✨", variant="primary")
+            generate_button = gr.Button("Generate 10 Outlines ✨", variant="primary")
+    
+    gr.Markdown("---")
+    gr.Markdown("## 📖 Your 10 Story Outlines")
     
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.TabItem("Mistral-7B"):
-                    gen1_output = gr.TextArea(label="Mistral's Tale", interactive=False, lines=12)
-                with gr.TabItem("Gemma-7B"):
-                    gen2_output = gr.TextArea(label="Gemma's Chronicle", interactive=False, lines=12)
-                with gr.TabItem("Llama-3-8B"):
-                    gen3_output = gr.TextArea(label="Llama's Legend", interactive=False, lines=12)
+    # Create 10 markdown components to display the stories in two columns
+    story_outputs = []
+    with gr.Row():
+        with gr.Column():
+            for i in range(5):
+                md = gr.Markdown(label=f"Story Outline {i + 1}")
+                story_outputs.append(md)
+        with gr.Column():
+            for i in range(5, 10):
+                md = gr.Markdown(label=f"Story Outline {i + 1}")
+                story_outputs.append(md)
 
     gr.Examples(
         examples=examples,
@@ -105,7 +142,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !i
     generate_button.click(
         fn=generate_stories,
         inputs=input_area,
-        outputs=[gen1_output, gen2_output, gen3_output],
+        outputs=story_outputs,
         api_name="generate"
     )