app.py

"""
Step-Audio-EditX - Audio Editing Demo using StepFun API
"""
import logging
import gradio as gr

from stepfun_api import get_api_token, process_audio, transcribe_audio

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def get_supported_edit_types():
    """
    获取支持的编辑类型和选项

    Returns:
        Dict[str, list]: Dictionary of edit types and their options
    """
    return {
        "clone": [],
        "emotion": [
            'happy', 'angry', 'sad', 'humour', 'confusion', 'disgusted',
            'empathy', 'embarrass', 'fear', 'surprised', 'excited',
            'depressed', 'coldness', 'admiration', 'remove'
        ],
        "style": [
            'serious', 'arrogant', 'child', 'older', 'girl', 'pure',
            'sister', 'sweet', 'ethereal', 'whisper', 'gentle', 'recite',
            'generous', 'act_coy', 'warm', 'shy', 'comfort', 'authority',
            'chat', 'radio', 'soulful', 'story', 'vivid', 'program',
            'news', 'advertising', 'roar', 'murmur', 'shout', 'deeply', 'loudly',
            'remove', 'exaggerated'
        ],
        "vad": [],
        "denoise": [],
        "paralinguistic": [],
        "speed": ["faster", "slower", "more faster", "more slower"],
    }


class EditxTab:
    """Audio editing and voice cloning interface tab"""

    def __init__(self):
        self.edit_type_list = list(get_supported_edit_types().keys())
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

    def history_messages_to_show(self, messages):
        """Convert message history to gradio chatbot format"""
        show_msgs = []
        for message in messages:
            edit_type = message['edit_type']
            edit_info = message['edit_info']
            source_text = message['source_text']
            target_text = message['target_text']
            raw_audio_path = message['raw_audio_path']
            edit_audio_path = message['edit_audio_path']
            type_str = f"{edit_type}-{edit_info}" if edit_info is not None else f"{edit_type}"
            show_msgs.extend([
                {"role": "user", "content": f"任务类型：{type_str}\n文本：{source_text}"},
                {"role": "user", "content": gr.Audio(value=raw_audio_path, interactive=False)},
                {"role": "assistant", "content": f"输出音频：\n文本：{target_text}"},
                {"role": "assistant", "content": gr.Audio(value=edit_audio_path, interactive=False)}
            ])
        return show_msgs

    def generate_clone(self, prompt_text_input, prompt_audio_input, generated_text, edit_type, edit_info, state):
        """Generate cloned audio using API"""
        self.logger.info("Starting voice cloning via API")

        # Input validation
        if not prompt_text_input or prompt_text_input.strip() == "":
            error_msg = "[Error] Uploaded text cannot be empty."
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state
        if not prompt_audio_input:
            error_msg = "[Error] Uploaded audio cannot be empty."
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state
        if not generated_text or generated_text.strip() == "":
            error_msg = "[Error] Clone content cannot be empty."
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state
        if edit_type != "clone":
            error_msg = "[Error] CLONE button must use clone task."
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state

        try:
            # Call API
            output_path = process_audio(
                prompt_audio_input, prompt_text_input, generated_text, edit_type, edit_info
            )

            # Create message for history
            cur_msg = {
                "edit_type": edit_type,
                "edit_info": edit_info,
                "source_text": prompt_text_input,
                "target_text": generated_text,
                "raw_audio_path": prompt_audio_input,
                "edit_audio_path": output_path,
            }
            state["history_audio"].append((output_path, generated_text))
            state["history_messages"].append(cur_msg)

            show_msgs = self.history_messages_to_show(state["history_messages"])
            self.logger.info("Voice cloning completed successfully")
            return show_msgs, state

        except Exception as e:
            error_msg = f"[Error] Clone failed: {str(e)}"
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state
        
    def generate_edit(self, prompt_text_input, prompt_audio_input, generated_text, edit_type, edit_info, state):
        """Generate edited audio using API"""
        self.logger.info("Starting audio editing via API")

        # Input validation
        if not prompt_audio_input:
            error_msg = "[Error] Uploaded audio cannot be empty."
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state

        try:
            # Determine which audio to use
            if len(state["history_audio"]) == 0:
                audio_to_edit = prompt_audio_input
                text_to_use = prompt_text_input
            else:
                audio_to_edit, previous_text = state["history_audio"][-1]
                text_to_use = previous_text

            # For para-linguistic, use generated_text; otherwise use source text
            if edit_type not in {"paralinguistic"}:
                generated_text = text_to_use

            # Call API
            output_path = process_audio(
                audio_to_edit, text_to_use, generated_text, edit_type, edit_info
            )

            # Create message for history
            cur_msg = {
                "edit_type": edit_type,
                "edit_info": edit_info,
                "source_text": text_to_use,
                "target_text": generated_text,
                "raw_audio_path": audio_to_edit,
                "edit_audio_path": output_path,
            }
            state["history_audio"].append((output_path, generated_text))
            state["history_messages"].append(cur_msg)

            show_msgs = self.history_messages_to_show(state["history_messages"])
            self.logger.info("Audio editing completed successfully")
            return show_msgs, state

        except Exception as e:
            error_msg = f"[Error] Edit failed: {str(e)}"
            self.logger.error(error_msg)
            return self.history_messages_to_show(state["history_messages"]) + [{"role": "user", "content": error_msg}], state

    def clear_history(self, state):
        """Clear conversation history"""
        state["history_messages"] = []
        state["history_audio"] = []
        return [], state

    def auto_transcribe_audio(self, audio_path, state):
        """
        自动转录音频文件，一次性返回最终结果

        Args:
            audio_path: 音频文件路径
            state: 状态字典

        Returns:
            转录的文本内容和更新后的状态
        """
        if not audio_path:
            return "", state

        # 防止重复调用 - 简化逻辑
        if state.get("last_audio_path") == audio_path:
            self.logger.debug(f"⚠️ Skipping duplicate transcription request for: {audio_path}")
            return state.get("last_transcribed_text", ""), state

        try:
            # 更新音频路径
            state["last_audio_path"] = audio_path

            self.logger.info(f"🎙️ Starting auto transcription for: {audio_path}")

            # 使用stepfun_api中的transcribe_audio函数，不使用streaming模式
            transcribed_text = transcribe_audio(audio_path, streaming=False)

            # 转录完成，缓存结果
            state["last_transcribed_text"] = transcribed_text

            self.logger.info(f"✅ Auto transcription completed: {transcribed_text}")
            return transcribed_text, state

        except Exception as e:
            error_msg = f"[转录失败: {str(e)}]"
            self.logger.error(f"❌ Auto transcription failed: {str(e)}")
            state["last_transcribed_text"] = error_msg
            return error_msg, state

    def init_state(self):
        """Initialize conversation state"""
        return {
            "history_messages": [],
            "history_audio": [],
            "last_audio_path": None,  # 用于防重复调用
            "last_transcribed_text": ""  # 缓存最后的转录结果
        }

    def update_edit_info(self, category):
        """Update sub-task dropdown based on main task selection"""
        category_items = get_supported_edit_types()
        choices = category_items.get(category, [])
        value = None if len(choices) == 0 else choices[0]
        return gr.Dropdown(label="Sub-task", choices=choices, value=value)

    def register_components(self):
        """Register gradio components - maintaining exact layout from original"""
        with gr.Tab("Editx"):
            with gr.Row():
                with gr.Column():
                    self.model_input = gr.Textbox(label="Model Name", value="Step-Audio-EditX", scale=1)
                    self.prompt_text_input = gr.Textbox(label="Prompt Text", value="", scale=1)
                    self.prompt_audio_input = gr.Audio(
                        sources=["upload", "microphone"],
                        format="wav",
                        type="filepath",
                        label="Input Audio",
                    )
                    self.generated_text = gr.Textbox(label="Target Text", lines=1, max_lines=200, max_length=1000)
                with gr.Column():
                    with gr.Row():
                        self.edit_type = gr.Dropdown(label="Task", choices=self.edit_type_list, value="clone")
                        self.edit_info = gr.Dropdown(label="Sub-task", choices=[], value=None)
                    self.chat_box = gr.Chatbot(label="History", type="messages", height=480*1)
            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        self.button_tts = gr.Button("CLONE", variant="primary")
                        self.button_edit = gr.Button("EDIT", variant="primary")
                with gr.Column():
                    self.clean_history_submit = gr.Button("Clear History", variant="primary")

            gr.Markdown("---")
            gr.Markdown("""
                **Button Description:**
                - CLONE: Synthesizes audio based on uploaded audio and text, only used for clone mode, will clear history information when used.
                - EDIT: Edits based on uploaded audio, or continues to stack edit effects based on the previous round of generated audio.
                """)
            gr.Markdown("""
                **Operation Workflow:**
                - Upload the audio to be edited on the left side and fill in the corresponding text content of the audio;
                - If the task requires modifying text content (such as clone, para-linguistic), fill in the text to be synthesized in the "target text" field. For all other tasks, keep the uploaded audio text content unchanged;
                - Select tasks and subtasks on the right side (some tasks have no subtasks, such as vad, etc.);
                - Click the "CLONE" or "EDIT" button on the left side, and audio will be generated in the dialog box on the right side.
                """)
            gr.Markdown("""
                **Para-linguistic Description:**
                - Supported tags include: [Breathing] [Laughter] [Surprise-oh] [Confirmation-en] [Uhm] [Surprise-ah] [Surprise-wa] [Sigh] [Question-ei] [Dissatisfaction-hnn]
                - Example:
                    - Fill in "target text" field: "Great, the weather is so nice today." Click the "CLONE" button to get audio.
                    - Change "target text" field to: "Great[Laughter], the weather is so nice today[Surprise-ah]." Click the "EDIT" button to get para-linguistic audio.
                """)

    def register_events(self):
        """Register event handlers"""
        state = gr.State(self.init_state())

        self.button_tts.click(
            self.generate_clone,
            inputs=[self.prompt_text_input, self.prompt_audio_input, self.generated_text, self.edit_type, self.edit_info, state],
            outputs=[self.chat_box, state]
        )
        self.button_edit.click(
            self.generate_edit,
            inputs=[self.prompt_text_input, self.prompt_audio_input, self.generated_text, self.edit_type, self.edit_info, state],
            outputs=[self.chat_box, state]
        )
        self.clean_history_submit.click(self.clear_history, inputs=[state], outputs=[self.chat_box, state])
        self.edit_type.change(
            fn=self.update_edit_info,
            inputs=self.edit_type,
            outputs=self.edit_info,
        )

        # 音频上传时自动转录
        self.prompt_audio_input.change(
            fn=self.auto_transcribe_audio,
            inputs=[self.prompt_audio_input, state],
            outputs=[self.prompt_text_input, state]
        )


def create_demo():
    """Create and return the Gradio demo"""
    editx_tab = EditxTab()
    
    with gr.Blocks(
        theme=gr.themes.Soft(),
        title="🎙️ Step-Audio-EditX",
        css="""
:root {
    --font: "Helvetica Neue", Helvetica, Arial, sans-serif;
    --font-mono: "SFMono-Regular", Consolas, "Liberation Mono", Menlo, monospace;
}
"""
    ) as demo:
        gr.Markdown("## 🎙️ Step-Audio-EditX")
        gr.Markdown("Audio Editing and Zero-Shot Cloning using Step-Audio-EditX")
        
        editx_tab.register_components()
        editx_tab.register_events()
    
    return demo


# Main entry point
if __name__ == "__main__":
    logger.info("🚀 Starting Step-Audio-EditX Demo (API Mode)")
    logger.info(f"API Token configured: {'Yes' if get_api_token() else 'No'}")
    
    demo = create_demo()
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )