Update app.py

app.py

@@ -64,51 +64,65 @@ 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.
+    print("--- Button clicked. Attempting to generate stories... ---")
+    
+    # If the model failed to load during startup, display that error.
     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)}`"
+        error_message = f"**Model failed to load during startup.**\n\nPlease check the console logs for details.\n\n**Error:**\n`{str(model_error)}`"
+        print(f"Returning startup error: {error_message}")
         return [error_message] * 10
         
     if not prompt:
         # Return a list of 10 empty strings to clear the outputs
         return [""] * 10
 
-    # A clear, instructive prompt format that works well with models like Phi-2.
-    story_prompt = f"""Instruct: Create a short story outline based on this idea: "{prompt}"
+    try:
+        # A clear, instructive prompt format that works well with models like Phi-2.
+        story_prompt = f"""Instruct: Create a short story outline based on this idea: "{prompt}"
 The outline should have three parts: a dramatic hook, a concise ballad, and a satisfying finale. Use emojis.
 Output:
 ### 🎬 The Hook
 """
 
-    # Parameters for the pipeline to generate 10 diverse results.
-    params = {
-        "max_new_tokens": 250,
-        "num_return_sequences": 10,
-        "do_sample": True,
-        "temperature": 0.9,
-        "top_k": 50,
-        "top_p": 0.95,
-        "pad_token_id": generator.tokenizer.eos_token_id
-    }
-    
-    # Generate 10 different story variations
-    outputs = generator(story_prompt, **params)
-
-    # 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']
-        # Add back the part of the prompt we want to see in the output
-        story_start = "### 🎬 The Hook\n"
-        generated_part = full_text.split(story_start)[-1]
-        stories.append(story_start + generated_part)
-
-    # Ensure we return exactly 10 stories, padding if necessary.
-    while len(stories) < 10:
-        stories.append("Failed to generate a story for this slot.")
-
-    return stories
+        # Parameters for the pipeline to generate 10 diverse results.
+        params = {
+            "max_new_tokens": 250,
+            "num_return_sequences": 10,
+            "do_sample": True,
+            "temperature": 0.9,
+            "top_k": 50,
+            "top_p": 0.95,
+            "pad_token_id": generator.tokenizer.eos_token_id
+        }
+        
+        print("Generating text with the model...")
+        # Generate 10 different story variations
+        outputs = generator(story_prompt, **params)
+        print("✅ Text generation complete.")
+
+        # 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']
+            # Add back the part of the prompt we want to see in the output
+            story_start = "### 🎬 The Hook\n"
+            generated_part = full_text.split(story_start)[-1]
+            stories.append(story_start + generated_part)
+
+        # Ensure we return exactly 10 stories, padding if necessary.
+        while len(stories) < 10:
+            stories.append("Failed to generate a story for this slot.")
+
+        return stories
+
+    except Exception as e:
+        # Catch any errors that happen DURING generation and display them in the UI.
+        print(f"--- 🚨 Error during story generation ---")
+        print(f"Error: {e}")
+        runtime_error_message = f"**An error occurred during story generation.**\n\nPlease check the console logs for details.\n\n**Error:**\n`{str(e)}`"
+        return [runtime_error_message] * 10
+
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !important;}") as demo: