Update app.py

app.py

@@ -7,7 +7,7 @@ import os
 TITLE = "✍️ AI Story Outliner"
 DESCRIPTION = """
 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.
+The app uses **DistilGPT-2** to generate creative outlines formatted in Markdown.
 
 **How it works:**
 1.  Enter your story idea.
@@ -25,32 +25,26 @@ examples = [
 ]
 
 # --- Model Initialization ---
-# This section loads a smaller, CPU-friendly model.
-# It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
+# This section loads a smaller, CPU-friendly model that does not require a token.
+generator = None
+model_error = None
+
 try:
     print("Initializing model... This may take a moment.")
     
-    # Load the token from environment variables if it exists (for HF Spaces secrets)
-    hf_token = os.environ.get("HF_TOKEN", None)
-    
-    # Using a smaller model that is more suitable for running without a high-end GPU.
+    # Using a smaller, fully open model that does not require authentication.
     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
+        model="distilgpt2",
+        torch_dtype=torch.float32, # Use float32 for wider CPU compatibility
+        device_map="auto" # Will use GPU if available, otherwise CPU
     )
-    print("✅ TinyLlama model loaded successfully!")
+    print("✅ DistilGPT-2 model loaded successfully!")
 
 except Exception as e:
-    print(f"--- 🚨 Error loading models ---")
-    print(f"Error: {e}")
-    # Create a dummy function if models fail, so the app can still launch with an error message.
-    def failed_generator(prompt, **kwargs):
-        error_message = f"Model failed to load. Please check the console for errors. Error: {e}"
-        return [{'generated_text': error_message}]
-    generator = failed_generator
+    model_error = e
+    print(f"--- 🚨 Error loading model ---")
+    print(f"Error: {model_error}")
 
 
 # --- App Logic ---
@@ -58,47 +52,54 @@ def generate_stories(prompt: str) -> list[str]:
     """
     Generates 10 story outlines from the loaded model based on the user's prompt.
     """
+    # If the model failed to load, display the error in all output boxes.
+    if model_error:
+        error_message = f"**Model failed to load.**\n\nPlease check the console logs for details.\n\n**Error:**\n`{str(model_error)}`"
+        return [error_message] * 10
+        
     if not prompt:
         # 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|>
+    # Note: Simpler models like DistilGPT-2 may not follow complex instructions perfectly.
+    story_prompt = f"""
+Story Idea: "{prompt}"
+
+Create a short story outline based on this idea.
+
+### The Hook
+A dramatic opening.
+
+### The Ballad
+The main story, told concisely.
+
+### The Finale
+A clear and satisfying ending.
+---
 """
 
     # Parameters for the pipeline to generate 10 diverse results.
     params = {
-        "max_new_tokens": 250,
+        "max_new_tokens": 200,
         "num_return_sequences": 10,
         "do_sample": True,
-        "temperature": 0.8,
+        "temperature": 0.9,
         "top_k": 50,
         "top_p": 0.95,
     }
     
     # Generate 10 different story variations
-    outputs = generator(system_prompt, **params)
+    outputs = generator(story_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)
+        # The model will generate the prompt plus the continuation, so we format it.
+        full_text = out['generated_text']
+        stories.append(full_text)
 
-    # Ensure we return exactly 10 stories, padding with an error message if necessary.
+    # Ensure we return exactly 10 stories, padding if necessary.
     while len(stories) < 10:
         stories.append("Failed to generate a story for this slot.")