Update app.py

app.py

@@ -43,12 +43,18 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
         combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script and music suggestion:"
         result = llama_pipeline(combined_prompt, max_new_tokens=200, do_sample=True, temperature=0.9)
 
-        generated_text = result[0]["generated_text"].split("Refined script and music suggestion:")[-1].strip()
-        script, music_suggestion = generated_text.split("Music Suggestion:")
-        return script.strip(), music_suggestion.strip()
+        # Ensure the output contains the expected structure
+        generated_text = result[0]["generated_text"]
+        if "Music Suggestion:" in generated_text:
+            script, music_suggestion = generated_text.split("Music Suggestion:", 1)
+            return script.strip(), music_suggestion.strip()
+        else:
+            # Fallback if the output format isn't as expected
+            return generated_text.strip(), "No specific music suggestion found."
     except Exception as e:
         return f"Error generating script: {e}", None
 
+
 # ---------------------------------------------------------------------
 # Voice-Over Generation Function
 # ---------------------------------------------------------------------