Fix warmup text length and ve attribute for torch.compile
All checks were successful
Build ROCm Image / build (push) Successful in 3m35s
All checks were successful
Build ROCm Image / build (push) Successful in 3m35s
- Warmup now uses a ~170-char representative sentence so torch.compile JIT-compiles for typical token sequence lengths. Previously "Warmup." compiled for very short shapes, causing a full re-compile (17s) on the first real HA request and pushing total synthesis past 30s. - Compile model.ve (voice encoder) in addition to s3gen — both are convolutional and hit the MIOpen workspace=0 bug. - Fix _patch_timing: attribute is model.ve not model.voice_encoder, so the timing wrap was silently skipping the speaker embedding. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
21
engine.py
21
engine.py
@@ -49,15 +49,18 @@ def load_model() -> bool:
|
||||
|
||||
if torch.cuda.is_available():
|
||||
# torch.compile replaces MIOpen's convolution path with Triton-generated
|
||||
# kernels, bypassing the workspace=0 fallback entirely. We compile only
|
||||
# s3gen (HiFiGAN vocoder + flow matching) since that's the bottleneck.
|
||||
# kernels, bypassing the workspace=0 fallback entirely. We compile s3gen
|
||||
# (HiFiGAN vocoder + flow matching) and the voice encoder (ve) since both
|
||||
# are convolutional and hit the workspace=0 bug.
|
||||
# suppress_errors=True falls back to eager for any op compile can't handle.
|
||||
try:
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
chatterbox_model.s3gen = torch.compile(chatterbox_model.s3gen, dynamic=True)
|
||||
logger.info("s3gen compiled with torch.compile")
|
||||
except Exception:
|
||||
logger.warning("torch.compile unavailable, running s3gen in eager mode", exc_info=True)
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
for attr, label in [("s3gen", "s3gen"), ("ve", "ve")]:
|
||||
try:
|
||||
obj = getattr(chatterbox_model, attr)
|
||||
setattr(chatterbox_model, attr, torch.compile(obj, dynamic=True))
|
||||
logger.info(f"{label} compiled with torch.compile")
|
||||
except Exception:
|
||||
logger.warning(f"torch.compile unavailable for {label}, running in eager mode", exc_info=True)
|
||||
|
||||
_patch_timing(chatterbox_model)
|
||||
logger.info("Model loaded successfully")
|
||||
@@ -87,7 +90,7 @@ def _patch_timing(model) -> None:
|
||||
pass
|
||||
try:
|
||||
# Speaker/voice encoder — xvector embedding from reference audio
|
||||
_wrap(model.voice_encoder, "forward", "voice_encoder (speaker embedding)")
|
||||
_wrap(model.ve, "forward", "ve (speaker embedding)")
|
||||
except AttributeError:
|
||||
pass
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user