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 **DistilGPT-2** to generate creative outlines formatted in Markdown.
+The app uses **Tencent's Hunyuan-1.8B** to generate creative outlines formatted in Markdown.
 
 **How it works:**
 1.  Enter your story idea.
@@ -34,19 +34,19 @@ try:
     print("Initializing model... This may take a moment.")
     
     # Explicitly load the token from environment variables (for HF Spaces secrets)
-    # This makes the authentication more robust.
+    # This makes the authentication more robust, overriding any bad default credentials.
     hf_token = os.environ.get("HF_TOKEN", None)
     
-    # Using a smaller, fully open model. Passing the token explicitly handles
-    # environments that might otherwise send invalid credentials.
+    # Using a smaller model from the user's list.
+    # Passing the token explicitly to ensure correct authentication.
     generator = pipeline(
         "text-generation",
-        model="distilgpt2",
-        torch_dtype=torch.float32, # Use float32 for wider CPU compatibility
+        model="tencent/Hunyuan-1.8B-Instruct",
+        torch_dtype=torch.bfloat16, # Use bfloat16 for better performance if available
         device_map="auto", # Will use GPU if available, otherwise CPU
         token=hf_token
     )
-    print("✅ DistilGPT-2 model loaded successfully!")
+    print("✅ Tencent/Hunyuan-1.8B-Instruct model loaded successfully!")
 
 except Exception as e:
     model_error = e
@@ -68,32 +68,19 @@ def generate_stories(prompt: str) -> list[str]:
         # 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.
-    # 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.
----
-"""
+    # This prompt format is specific to the Hunyuan model.
+    system_instruction = "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."
+    
+    story_prompt = f"<|im_start|>system\n{system_instruction}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
 
     # Parameters for the pipeline to generate 10 diverse results.
     params = {
-        "max_new_tokens": 200,
+        "max_new_tokens": 250,
         "num_return_sequences": 10,
         "do_sample": True,
-        "temperature": 0.9,
-        "top_k": 50,
+        "temperature": 0.8,
         "top_p": 0.95,
+        "pad_token_id": generator.tokenizer.eos_token_id # Suppress warning
     }
     
     # Generate 10 different story variations
@@ -102,9 +89,10 @@ A clear and satisfying ending.
     # Extract the generated text and clean it up.
     stories = []
     for out in outputs:
-        # The model will generate the prompt plus the continuation, so we format it.
+        # Remove the prompt part from the full generated text
         full_text = out['generated_text']
-        stories.append(full_text)
+        assistant_response = full_text.split("<|im_start|>assistant\n")[-1]
+        stories.append(assistant_response)
 
     # Ensure we return exactly 10 stories, padding if necessary.
     while len(stories) < 10: