Add reprompter

README.md

@@ -19,7 +19,7 @@ Stable Audio 3 is the next generation of Stable Audio: a focused, streamlined pl
 | [**Stable Audio 3 Small-Music**](https://huggingface.co/stabilityai/stable-audio-3-small-music) | `small-music` | SAME-Small | CPU | 433M | 120s | Lightweight music-only inference, no GPU required |
 | [**Stable Audio 3 Small-SFX**](https://huggingface.co/stabilityai/stable-audio-3-small-sfx) | `small-sfx` | SAME-Small | CPU | 433M | 120s | Lightweight sound effects-only inference, no GPU required |
 | [**Stable Audio 3 Medium**](https://huggingface.co/stabilityai/stable-audio-3-medium) | `medium` | SAME-Large | GPU (CUDA) | 1.4B | 380s | High Quality, Fast Inference |
-| **Stable Audio 3 Large** | — | SAME-Large | API only | 2.7B | 380s | Highest quality, API only. Not supported by this repo, see the [API docs](TODO) |
+| **Stable Audio 3 Large** | — | SAME-Large | API only | 2.7B | 380s | Highest quality, API only. Not supported by this repo, see the [API docs](https://stableaudio.com/user-guide) |
 
 Base (un-post-trained) checkpoints, the SAME autoencoders, and optimized variants are available in the [Extra Models collection](https://huggingface.co/collections/stabilityai/stable-audio-3-extra).
 
@@ -200,20 +200,9 @@ audio_out = ae.decode(latents)
 See [Autoencoder Workflows](docs/workflows/autoencoder.md) for encoding batches, chunked processing, and pre-encoding datasets for LoRA training.
 
 ## Hardware Support
-Stable Audio 3 scales from a laptop to a GPU server. Right now, we have scripts for generating audio with various backends.
+Stable Audio 3 scales from a laptop to a GPU server.
 
-TODO: FINISH THIS
-
-As an example,
-```
-bash optimized/run_trt.sh
-
-```
-
-
-Check out the `optimized/` directory for other options
-
-*More Hardware Support COMING SOON*
+*Hardware Support Scripts COMING SOON*
 
 
 ## Docs
@@ -232,7 +221,7 @@ Check out the `optimized/` directory for other options
 
 - [Harmonai Discord](https://discord.gg/cKpvjey8b): Check out our Harmonai Discord server run by the research team. Besides good discussions, we host weekly office hours talking all things AI audio and music and want to hear what you come up with!
 
-- [Underfit](https://github.com/dada-bots/underfit): A LoRA training poweruser dream from Dadabots. If LoRA training in this repo is not enough, check out some experimental tools there.
+- [Underfit](https://github.com/dada-bots/underfit): A LoRA training poweruser dream from Dadabots. If LoRA training in this repo is not enough, check out some experimental tools there like agentic LoRA orchestrations and monitoring.
 
 ---
 

docs/workflows/lora.md

@@ -44,7 +44,7 @@ uv run python scripts/train_lora.py \
 
 # Training Configuration
 
-> Looking for more advanced training tooling? Check out [Underfit](https://github.com/dada-bots/underfit) from Dadabots. There you can find TODO fill-this in
+> Looking for more advanced training tooling? Check out [Underfit](https://github.com/dada-bots/underfit) from Dadabots. There you can find scripts to orchestrate LoRA training runs, monitor runs from a dashboard, and much much more.
 
 ## CLI Arguments
 

stable_audio_3/interface/diffusion_cond.py

@@ -11,14 +11,15 @@
 from einops import rearrange
 
 from stable_audio_3.interface.aeiou import audio_spectrogram_image
-from stable_audio_3.verbose import vprint
 from stable_audio_3.inference.distribution_shift import LogSNRShift, FluxDistributionShift, DistributionShift, IdentityDistributionShift
 from stable_audio_3.models.lora import has_lora
+from stable_audio_3.interface.reprompt import reprompt as _reprompt_fn, get_model as _reprompt_get_model, is_model_cached as _reprompt_is_model_cached
 
 stable_audio_3_model = None
 sample_size = 5324800
 sample_rate = 44100
 n_loras = 0
+_LENGTH_EXTRACT_RE = re.compile(r' Length: (\d+) seconds\.?\s*$')
 
 
 # when using a prompt in a filename
@@ -287,21 +288,24 @@ def create_sampling_ui(stable_audio_3_model, default_prompt=None):
     if default_prompt is None:
         default_prompt = ""
 
-
+    _reprompt_model_id = "Qwen/Qwen3.5-2B"
+    _reprompt_cached = _reprompt_is_model_cached(_reprompt_model_id)
 
     with gr.Row():
         with gr.Column(scale=6):
             prompt = gr.Textbox(show_label=False, placeholder="Prompt", value=default_prompt)
             negative_prompt = gr.Textbox(show_label=False, placeholder="Negative prompt")
-        with gr.Column(scale=3):
-            with gr.Row():
-                generate_button = gr.Button("Generate", variant='primary', scale=1)
+        prompt_assistant_button = gr.Button(
+            "Prompt Assistant" if _reprompt_cached else "Download Prompt Assistant (~4.2 GB)",
+            scale=1
+        )
+        generate_button = gr.Button("Generate", variant='primary', scale=1)
 
     with gr.Row(equal_height=False):
         with gr.Column():
             with gr.Row(visible = True):
                 # Timing controls
-                seconds_total_slider = gr.Slider(minimum=0, maximum=512, step=1, value=sample_size//sample_rate, label="Seconds total", visible=has_seconds_total)
+                seconds_total_slider = gr.Slider(minimum=0, maximum=sample_size//sample_rate, step=1, value=sample_size//sample_rate, label="Seconds total", visible=has_seconds_total)
 
             with gr.Row():
                 # Steps slider
@@ -532,6 +536,26 @@ def update_inpaint_sliders(seconds_total):
         ],
         api_name="generate")
 
+    def _prompt_assistant_or_download(text, progress=gr.Progress(track_tqdm=True)):
+        if not _reprompt_is_model_cached(_reprompt_model_id):
+            _reprompt_get_model(_reprompt_model_id)
+            return text, gr.update(), gr.update(value="Prompt Assistant")
+        _, result, _ = _reprompt_fn(text, "Auto", "", _reprompt_model_id, 128, 1.11)
+        m = _LENGTH_EXTRACT_RE.search(result)
+        if m:
+            max_seconds = sample_size // sample_rate
+            seconds = min(int(m.group(1)), max_seconds)
+            result = result[:m.start()]
+        else:
+            seconds = gr.update()
+        return result, seconds, gr.update()
+
+    prompt_assistant_button.click(
+        fn=_prompt_assistant_or_download,
+        inputs=[prompt],
+        outputs=[prompt, seconds_total_slider, prompt_assistant_button],
+        concurrency_limit=1,
+    )
 
 def create_diffusion_cond_ui(model, gradio_title="", default_prompt=None):
     global sample_size, sample_rate, stable_audio_3_model