Update app.py

app.py

@@ -36,23 +36,27 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
         )
         llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+        # System prompt with clear structure instructions
         system_prompt = (
             f"You are an expert radio imaging producer specializing in sound design and music. "
-            f"Based on the user's concept and the selected duration of {duration} seconds, generate: "
-            f"1. A concise voice-over script. "
-            f"2. Suggestions for sound design. "
-            f"3. Music styles or track recommendations."
+            f"Based on the user's concept and the selected duration of {duration} seconds, produce the following: "
+            f"1. A concise voice-over script. Prefix this section with 'Voice-Over Script:'.\n"
+            f"2. Suggestions for sound design. Prefix this section with 'Sound Design Suggestions:'.\n"
+            f"3. Music styles or track recommendations. Prefix this section with 'Music Suggestions:'."
         )
 
         combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nOutput:"
         result = llama_pipeline(combined_prompt, max_new_tokens=300, do_sample=True, temperature=0.8)
 
+        # Parsing output
         generated_text = result[0]["generated_text"].split("Output:")[-1].strip()
-        if "Suggestions:" in generated_text and "Music Recommendations:" in generated_text:
-            script, rest = generated_text.split("Suggestions:", 1)
-            sound_design, music_suggestion = rest.split("Music Recommendations:", 1)
-            return script.strip(), sound_design.strip(), music_suggestion.strip()
-        return "Error: Could not parse output.", "", ""
+
+        # Extract sections based on prefixes
+        voice_script = generated_text.split("Voice-Over Script:")[1].split("Sound Design Suggestions:")[0].strip() if "Voice-Over Script:" in generated_text else "No voice-over script found."
+        sound_design = generated_text.split("Sound Design Suggestions:")[1].split("Music Suggestions:")[0].strip() if "Sound Design Suggestions:" in generated_text else "No sound design suggestions found."
+        music_suggestions = generated_text.split("Music Suggestions:")[1].strip() if "Music Suggestions:" in generated_text else "No music suggestions found."
+
+        return voice_script, sound_design, music_suggestions
     except Exception as e:
         return f"Error generating script: {e}", "", ""