src/ui/tabs.py

"""
Gradio Tab Definitions

Defines all 6 tabs for the StackNetdemo application.
"""

import gradio as gr


def create_text_to_music_tab():
    """Create the Text to Music tab components."""
    with gr.Column():
        gr.Markdown("### Generate original music from a text description")

        prompt = gr.Textbox(
            label="Describe your music",
            placeholder="e.g., upbeat jazz with piano and saxophone, cheerful summer vibes",
            lines=3
        )

        with gr.Row():
            tags = gr.Textbox(
                label="Genre/Style Tags",
                placeholder="jazz, piano, instrumental",
                scale=2
            )
            instrumental = gr.Checkbox(
                label="Instrumental Only",
                value=False,
                scale=1
            )

        lyrics = gr.Textbox(
            label="Lyrics (optional)",
            placeholder="Write your lyrics here...",
            lines=4,
            visible=True
        )

        title = gr.Textbox(
            label="Song Title (optional)",
            placeholder="My Song"
        )

        generate_btn = gr.Button("Generate Music", variant="primary", size="lg")

        status = gr.Textbox(label="Status", interactive=False, visible=False)

        output_audio = gr.Audio(label="Generated Music", type="filepath")

    # Toggle lyrics visibility based on instrumental checkbox
    instrumental.change(
        fn=lambda x: gr.update(visible=not x),
        inputs=[instrumental],
        outputs=[lyrics],
        api_name=None
    )

    return {
        "prompt": prompt,
        "tags": tags,
        "instrumental": instrumental,
        "lyrics": lyrics,
        "title": title,
        "generate_btn": generate_btn,
        "status": status,
        "output_audio": output_audio
    }


def create_music_to_music_tab():
    """Create the Music to Music tab with sub-tabs for Cover and Stems."""
    with gr.Tabs() as sub_tabs:
        # Cover Song Sub-tab
        with gr.Tab("Create Cover"):
            with gr.Column():
                gr.Markdown("### Create music from reference audio (Diffusion)")

                cover_audio_input = gr.Audio(
                    label="Upload Audio",
                    type="filepath"
                )

                cover_style_prompt = gr.Textbox(
                    label="Style Direction",
                    placeholder="e.g., rock version with electric guitar, female vocalist",
                    lines=2
                )

                cover_tags = gr.Textbox(
                    label="Style Tags",
                    placeholder="rock, electric guitar"
                )

                cover_title = gr.Textbox(
                    label="Title (optional)",
                    placeholder="My Song"
                )

                cover_btn = gr.Button("Create", variant="primary", size="lg")

                cover_status = gr.Textbox(label="Status", interactive=False, visible=False)

                cover_output = gr.Audio(label="Song", type="filepath")

        # Extract Stems Sub-tab
        with gr.Tab("Extract Stems"):
            with gr.Column():
                gr.Markdown("### Separate audio into individual stems")

                stems_audio_input = gr.Audio(
                    label="Upload Audio",
                    type="filepath"
                )

                stems_btn = gr.Button("Extract Stems", variant="primary", size="lg")

                stems_status = gr.Textbox(label="Status", interactive=False, visible=False)

                gr.Markdown("**Extracted Stems:**")
                with gr.Row():
                    vocals_output = gr.Audio(label="Vocals", type="filepath")
                    drums_output = gr.Audio(label="Drums", type="filepath")
                with gr.Row():
                    bass_output = gr.Audio(label="Bass", type="filepath")
                    other_output = gr.Audio(label="Other", type="filepath")

    return {
        # Cover components
        "cover_audio_input": cover_audio_input,
        "cover_style_prompt": cover_style_prompt,
        "cover_tags": cover_tags,
        "cover_title": cover_title,
        "cover_btn": cover_btn,
        "cover_status": cover_status,
        "cover_output": cover_output,
        # Stems components
        "stems_audio_input": stems_audio_input,
        "stems_btn": stems_btn,
        "stems_status": stems_status,
        "vocals_output": vocals_output,
        "drums_output": drums_output,
        "bass_output": bass_output,
        "other_output": other_output
    }


def create_text_to_image_tab():
    """Create the Text to Image tab components."""
    with gr.Column():
        gr.Markdown("### Generate images from a text description")

        prompt = gr.Textbox(
            label="Describe your image",
            placeholder="e.g., a serene mountain landscape at sunset with snow-capped peaks",
            lines=3
        )

        # Example prompt buttons
        with gr.Row():
            example1_btn = gr.Button("Example 1", size="sm")
            example2_btn = gr.Button("Example 2", size="sm")
            example3_btn = gr.Button("Example 3", size="sm")
            example4_btn = gr.Button("Example 4", size="sm")

        format_type = gr.Dropdown(
            label="Format",
            choices=["image", "multi", "3d"],
            value="image"
        )

        generate_btn = gr.Button("Generate", variant="primary", size="lg")

        status = gr.Textbox(label="Status", interactive=False, visible=False)

        # Image output (for image and multi formats)
        output_image = gr.Image(label="Generated Image", type="filepath", visible=True)

        # 3D Model output (for 3d format)
        output_model = gr.Model3D(label="Generated 3D Model", visible=False)

    # Toggle output visibility based on format selection
    format_type.change(
        fn=lambda fmt: (
            gr.update(visible=(fmt != "3d")),
            gr.update(visible=(fmt == "3d"))
        ),
        inputs=[format_type],
        outputs=[output_image, output_model],
        api_name=None
    )

    # Example prompts
    example_prompts = [
        "An Adorable Alien Rambutan with large round eyes wearing acid-etched armor, in the style of a Pixar character design, soft studio lighting, highly detailed fur texture, whimsical expression, floating in a cosmic fruit bowl surrounded by stardust, 8k resolution, cinematic composition",
        "panning photo of a red 1989 two-door Mercedes-Benz 450 SLC driving at high speed along a coastal highway, golden hour lighting, motion blur on wheels and background, sharp focus on car body, Mediterranean Sea visible in background, professional automotive photography",
        "A high-contrast flash photograph captured on Fuji Superia X-Tra 400 film showing an elderly woman in a floral housedress washing a bright yellow Lamborghini Countach in her suburban driveway, garden hose in hand, curlers in hair, bemused expression, mundane meets extraordinary",
        '{"subject": "great horned owl", "setting": "ancient library at midnight", "style": "Dutch Golden Age painting", "lighting": "single candle illumination with dramatic chiaroscuro", "details": "owl perched on stack of leather-bound books, spectacles on beak, quill pen nearby", "mood": "scholarly and mysterious", "technical": "oil painting texture, visible brushstrokes, museum quality"}'
    ]

    # Wire up example buttons
    example1_btn.click(
        fn=lambda: example_prompts[0],
        inputs=[],
        outputs=[prompt],
        api_name=None
    )
    example2_btn.click(
        fn=lambda: example_prompts[1],
        inputs=[],
        outputs=[prompt],
        api_name=None
    )
    example3_btn.click(
        fn=lambda: example_prompts[2],
        inputs=[],
        outputs=[prompt],
        api_name=None
    )
    example4_btn.click(
        fn=lambda: example_prompts[3],
        inputs=[],
        outputs=[prompt],
        api_name=None
    )

    return {
        "prompt": prompt,
        "format_type": format_type,
        "generate_btn": generate_btn,
        "status": status,
        "output_image": output_image,
        "output_model": output_model
    }


def create_image_to_image_tab():
    """Create the Image to Image tab components."""
    with gr.Column():
        gr.Markdown("### Transform or edit an existing image")

        with gr.Row():
            input_image = gr.Image(
                label="Upload Image",
                type="filepath",
                scale=1
            )

            with gr.Column(scale=1):
                edit_prompt = gr.Textbox(
                    label="Edit Instructions",
                    placeholder="e.g., add dramatic sunset lighting, make it look like a painting",
                    lines=3
                )

                strength = gr.Slider(
                    label="Edit Strength",
                    minimum=0.1,
                    maximum=1.0,
                    value=0.5,
                    step=0.1
                )

        edit_btn = gr.Button("Transform Image", variant="primary", size="lg")

        status = gr.Textbox(label="Status", interactive=False, visible=False)

        output_image = gr.Image(label="Transformed Image", type="filepath")

    return {
        "input_image": input_image,
        "edit_prompt": edit_prompt,
        "strength": strength,
        "edit_btn": edit_btn,
        "status": status,
        "output_image": output_image
    }


def create_text_to_video_tab():
    """Create the Text to Video tab components."""
    with gr.Column():
        gr.Markdown("### Generate videos from a text description")

        prompt = gr.Textbox(
            label="Describe your video",
            placeholder="e.g., a drone shot flying over a tropical beach at golden hour",
            lines=3
        )

        with gr.Row():
            duration = gr.Slider(
                label="Duration (seconds)",
                minimum=3,
                maximum=30,
                value=10,
                step=1,
                scale=1
            )
            style = gr.Dropdown(
                label="Style",
                choices=["Cinematic", "Animation", "Documentary", "Abstract"],
                value="Cinematic",
                scale=1
            )

        generate_btn = gr.Button("Generate Video", variant="primary", size="lg")

        status = gr.Textbox(label="Status", interactive=False, visible=False)

        output_video = gr.Video(label="Generated Video")

    return {
        "prompt": prompt,
        "duration": duration,
        "style": style,
        "generate_btn": generate_btn,
        "status": status,
        "output_video": output_video
    }


def create_image_to_video_tab():
    """Create the Image to Video tab components."""
    with gr.Column():
        gr.Markdown("### Animate a static image into video")

        with gr.Row():
            input_image = gr.Image(
                label="Upload Image",
                type="filepath",
                scale=1
            )

            with gr.Column(scale=1):
                motion_prompt = gr.Textbox(
                    label="Motion Description",
                    placeholder="e.g., gentle zoom in, clouds moving slowly, water rippling",
                    lines=3
                )

                duration = gr.Slider(
                    label="Duration (seconds)",
                    minimum=3,
                    maximum=15,
                    value=5,
                    step=1
                )

        animate_btn = gr.Button("Animate Image", variant="primary", size="lg")

        status = gr.Textbox(label="Status", interactive=False, visible=False)

        output_video = gr.Video(label="Animated Video")

    return {
        "input_image": input_image,
        "motion_prompt": motion_prompt,
        "duration": duration,
        "animate_btn": animate_btn,
        "status": status,
        "output_video": output_video
    }


def create_all_tabs():
    """Create all tabs and return component references."""
    tabs = {}

    with gr.Tab("Text to Music", id="text-to-music"):
        tabs["text_to_music"] = create_text_to_music_tab()

    with gr.Tab("Music to Music", id="music-to-music"):
        tabs["music_to_music"] = create_music_to_music_tab()

    with gr.Tab("Text to Image", id="text-to-image"):
        tabs["text_to_image"] = create_text_to_image_tab()

    with gr.Tab("Image to Image", id="image-to-image"):
        tabs["image_to_image"] = create_image_to_image_tab()

    with gr.Tab("Text to Video", id="text-to-video"):
        tabs["text_to_video"] = create_text_to_video_tab()

    with gr.Tab("Image to Video", id="image-to-video"):
        tabs["image_to_video"] = create_image_to_video_tab()

    return tabs