Update app.py

app.py

@@ -2,13 +2,12 @@ import gradio as gr
 import torch
 from transformers import pipeline
 import os
-from huggingface_hub import login
 
 # --- App Configuration ---
 TITLE = "✍️ AI Story Outliner"
 DESCRIPTION = """
-Enter a prompt and get 10 unique story outlines from a powerful AI model.
-The app uses **Mistral-7B-v0.1**, a popular and capable open-source model, to generate creative outlines.
+Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
+The app uses **DistilGPT-2**, a reliable and lightweight model, to generate creative outlines.
 
 **How it works:**
 1.  Enter your story idea.
@@ -26,35 +25,22 @@ examples = [
 ]
 
 # --- Model Initialization ---
-# This section loads the Mistral-7B model, which requires authentication.
-# It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
+# This section loads a smaller, stable, and CPU-friendly model that requires no authentication.
 generator = None
 model_error = None
 
 try:
     print("Initializing model... This may take a moment.")
     
-    # Explicitly load the token from environment variables (for HF Spaces secrets).
-    hf_token = os.environ.get("HF_TOKEN")
-
-    if hf_token:
-        print("✅ HF_TOKEN secret found. Logging in...")
-        # Programmatically log in to Hugging Face. This is a more robust method.
-        login(token=hf_token)
-        print("✅ Login successful.")
-    else:
-        # If no token is found, raise an error to prevent the app from crashing later.
-        raise ValueError("Hugging Face token not found. Please set the HF_TOKEN secret in your Space settings.")
-
-    # Using 'mistralai/Mistral-7B-v0.1'.
-    # After login(), the token argument is no longer needed here as the session is authenticated.
+    # Using 'distilgpt2', a stable and widely supported model that does not require a token.
+    # This is much more suitable for a standard CPU environment.
     generator = pipeline(
         "text-generation",
-        model="mistralai/Mistral-7B-v0.1",
-        torch_dtype=torch.bfloat16, # More performant data type
+        model="distilgpt2",
+        torch_dtype=torch.float32, # Use float32 for wider CPU compatibility
         device_map="auto" # Will use GPU if available, otherwise CPU
     )
-    print("✅ mistralai/Mistral-7B-v0.1 model loaded successfully!")
+    print("✅ distilgpt2 model loaded successfully!")
 
 except Exception as e:
     model_error = e
@@ -80,19 +66,30 @@ def generate_stories(prompt: str) -> list[str]:
         return [""] * 10
 
     try:
-        # This prompt format is optimized for Mistral instruct models.
-        story_prompt = f"""[INST] Create a short story outline based on this idea: "{prompt}"
-The outline must have three parts: a dramatic hook, a concise ballad, and a satisfying finale. Use emojis for each section header. [/INST]
+        # A generic story prompt that works well with models like GPT-2.
+        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,
-            "top_p": 0.95,
+            "temperature": 0.9,
+            "top_k": 50,
             "pad_token_id": generator.tokenizer.eos_token_id
         }
         
@@ -104,11 +101,8 @@ The outline must have three parts: a dramatic hook, a concise ballad, and a sati
         # Extract the generated text.
         stories = []
         for out in outputs:
-            # The model will generate the prompt plus the continuation. We extract just the new part.
             full_text = out['generated_text']
-            # Split by the instruction closing tag to get only the model's response
-            generated_part = full_text.split("[/INST]")[-1].strip()
-            stories.append(generated_part)
+            stories.append(full_text)
 
         # Ensure we return exactly 10 stories, padding if necessary.
         while len(stories) < 10: