Skip to content

Add Tesla T4 support, SDPA fallback, and Gradio demo#13

Open
akshatvishu wants to merge 10 commits into
inclusionAI:mainfrom
akshatvishu:feature/tesla-t4-support
Open

Add Tesla T4 support, SDPA fallback, and Gradio demo#13
akshatvishu wants to merge 10 commits into
inclusionAI:mainfrom
akshatvishu:feature/tesla-t4-support

Conversation

@akshatvishu

Copy link
Copy Markdown
  • add Tesla T4 fallback for environments without native bfloat16
  • upcast to float32 at generate time to avoid NaN audio on float16 hardware
  • add flash-attn fallback to native PyTorch scaled_dot_product_attention
  • include Gradio demo for easier local testing
  • preserve original bfloat16 path on supported GPUs

- add Tesla T4 fallback for environments without native bfloat16
- upcast to float32 at generate time to avoid NaN audio on float16 hardware
- add flash-attn fallback to native PyTorch scaled_dot_product_attention
- include Gradio demo for easier local testing
- preserve original bfloat16 path on supported GPUs

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
@gemini-code-assist

Copy link
Copy Markdown

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the model's compatibility and user experience by adding explicit support for Tesla T4 GPUs and providing a user-friendly Gradio demo. The changes ensure stable audio generation across different hardware configurations through intelligent data type handling and attention mechanism fallbacks, making the system more accessible and robust for various audio synthesis tasks.

Highlights

  • Tesla T4 GPU Support: Implemented robust support for Tesla T4 GPUs by introducing a dynamic data type strategy. The model now automatically upcasts to float32 during generation on float16 hardware (like T4) to prevent NaN audio outputs, while preserving bfloat16 on supported GPUs (A100/H100).
  • SDPA Fallback for Attention: Added a fallback mechanism for Flash Attention. If Flash Attention is not available or cannot be imported, the system now gracefully defaults to native PyTorch's scaled_dot_product_attention (SDPA), ensuring broader compatibility and T4 friendliness.
  • Gradio Demo Integration: Introduced a new Gradio-based web interface for easier local testing and interaction with the audio generation model. This demo supports various tasks including TTS, TTA, Music/BGM generation, zero-shot cloning, multi-speaker podcasts, and text generation, with comprehensive controls for voice attributes and decode settings.
  • Code Refinements and Readability: Performed several code cleanups and refactorings across modeling_bailingmm.py and audio_tokenizer/audio_encoder.py to improve clarity, remove redundant code, and enhance maintainability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • audio_tokenizer/audio_encoder.py
    • Modified Flash Attention import to be optional, introducing a HAS_FLASH_ATTN flag.
    • Implemented a fallback to native PyTorch scaled_dot_product_attention (SDPA) if Flash Attention is unavailable, explicitly noting it as T4 friendly.
  • gradio_demo.py
    • Added a new Gradio application for demonstrating audio generation capabilities.
    • Included logic for dynamic LLM data type resolution (bfloat16 or float16) based on CUDA support.
    • Implemented a function to pick the optimal attention backend (Flash Attention, SDPA, or eager).
    • Provided a comprehensive UI with controls for task type, text input, podcast mode, reference audio, voice attributes, BGM settings, and decode parameters.
    • Incorporated a mechanism to upcast waveforms to float32 to prevent NaN/Inf issues on float16 hardware.
  • modeling_bailingmm.py
    • Removed explicit self.llm_dytpe initialization in __init__ to rely on from_pretrained behavior.
    • Added a _infer_dtype property to accurately determine the model's actual data type post-initialization.
    • Introduced a detailed comment explaining the data type strategy for bfloat16, float16 (T4), and the necessity of full float32 upcasting on T4 to avoid numerical instability.
    • Refactored get_rope_index and prepare_input_embed for improved readability and conciseness, including consolidating spk_token_id logic.
    • Updated sample and generate_text methods to use the inferred run_dtype for torch.autocast.
    • Implemented an explicit self.float() call in the generate method to upcast the entire model to float32 when running on float16 hardware (e.g., T4), ensuring stable audio output.
  • requirements.txt
    • Added gradio>=5.0.0 as a new dependency for the Gradio demo.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces significant improvements by adding fallback support for Tesla T4 GPUs and a comprehensive Gradio demo for easier testing. However, it introduces critical security vulnerabilities: the Gradio UI allows arbitrary model paths leading to Remote Code Execution (RCE), the model's prompt construction is vulnerable to prompt injection, and the application is publicly shared by default, increasing remote exploitation risk. Additionally, a missing dependency in requirements.txt will prevent the Gradio demo from functioning correctly.

Comment thread gradio_demo.py
Comment thread gradio_demo.py
Comment thread modeling_bailingmm.py
Comment thread gradio_demo.py Outdated
Comment thread modeling_bailingmm.py Outdated
Comment thread modeling_bailingmm.py Outdated
…usionAI#12)

- Restore the outer autocast scope in generate() to fix the test regression.
- Implement dynamic _autocast_context() to handle hardware-specific dtypes.
- Update T4 policy to load float32 directly in demo and test scripts.
- Retain defensive float16 -> float32 upcast in generate() for manual callers.

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Normalize attention backends to 'eager' across all nested Qwen
modules to prevent FlashAttention2 import errors on Colab T4
hardware.

- Implement backend policy in cookbooks/test.py and gradio_demo.py.
- Force top-level BailingMM config to 'eager' implementation.
- Recursively override nested Qwen2 configs in AudioVAE (encoder/decoder
  backbones) to bypass automatic FlashAttention detection.
- Maintain float32 for T4 while keeping bf16 paths for supported GPUs.

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
…rate

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
… of the repo.

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
@akshatvishu akshatvishu marked this pull request as ready for review March 11, 2026 15:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant