Drop gpu_lock: Gradio's default concurrency_limit=1 already serializes process_fn

app.py

@@ -62,20 +62,11 @@ VARIANTS = {
     },
 }
 
-# Only one variant's full stack (text encoder + DiT + VAE) is kept on the GPU at a
-# time: the two variants need different Qwen2.5-Omni-7B layer-sampling depths, so
-# the encoder can't be shared as-is (QwenOmniThinkerExtractor bakes dit_depth into
-# which LLM layers it extracts at construction time). Switching variants evicts the
-# previous one and rebuilds from the local HF cache (no re-download).
+# Only one variant is kept on the GPU at a time; switching evicts the other.
 _active_name = None
 _active_entry = None    # built on CPU at first; moved to GPU on first real use
 _loading = True
 _load_error = None
-
-# Only one request runs at a time. The scheduler object is shared and stateful,
-# so two requests running together would corrupt each other's results. This lock
-# also covers _active_name/_active_entry below, so no separate lock is needed.
-_gpu_lock = threading.Lock()
 _active_ready = False  # has _active_entry been moved onto the GPU yet?
 
 
@@ -162,9 +153,8 @@ def get_variant(name: str):
     """Return the active variant's (model, extractor, vae, ...) tuple, ready to
     use on the GPU — building and/or activating it first if needed.
 
-    Caller must hold _gpu_lock and be inside an @spaces.GPU call — that's both
-    what makes the unlocked _active_* globals safe, and what makes it safe to
-    touch CUDA in _activate_variant()."""
+    Must be called from inside process_fn (i.e. inside @spaces.GPU) — that's
+    what makes it safe to touch CUDA in _activate_variant()."""
     global _active_name, _active_entry, _active_ready
     if name != _active_name:
         print(f"Switching variant: {_active_name!r} -> {name!r}")
@@ -214,108 +204,107 @@ def process_fn(
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
         out_path = f.name
 
-    with _gpu_lock:
-        model, omni_extractor, audio_vae, vae_scale_factor, vae_sample_rate, gen_target_frames, scheduler = get_variant(model_variant)
-
-        if mode == "Generate":
-            # Text-to-audio/speech, no input audio involved. sound_type picks which of
-            # the model's trained prompt templates to build (see README task table) —
-            # voice/background only apply to the Speech templates, unused otherwise.
-            if sound_type == "Sound Effect":
-                tagged_prompt = f"[Audio] {prompt}"
-            elif sound_type == "Speech":
-                tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}"'
-            else:  # "Speech + Background"
-                tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}" [Audio] {background}'
-
-            latents = sample_latents(
-                model, scheduler, omni_extractor, [tagged_prompt],
-                num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
-                target_frames=gen_target_frames,
-            )
-            latents = latents * (1.0 / vae_scale_factor)
-            decode_and_save(audio_vae, latents, [duration], [out_path], sample_rate=vae_sample_rate)
-
-        elif mode == "Edit":
-            # Source audio + instruction. Mask is all-ones: the whole clip is editable,
-            # source_latents just gives the model something to condition on.
-            # sound_type again picks [Audio] vs [Speech] as the edit's sub-tag; "Speech
-            # + Background" isn't a real edit template, so it also falls back to [Audio].
-            if not input_audio_path:
-                raise gr.Error("Edit mode requires an input audio track.")
-            edit_target = "Speech" if sound_type == "Speech" else "Audio"
-            tgt_wav_len = int(MAX_AUDIO_DURATION * vae_sample_rate)
-            src_wav = load_source_audio(input_audio_path, target_sr=vae_sample_rate,
-                                         target_length=tgt_wav_len, device=DEVICE)
-            src_latent = audio_vae.wrapped_encode(src_wav) * vae_scale_factor
-            mask_wav = make_edit_mask(src_wav.shape[-1], DEVICE)
-            mask_lat = downsample_mask(mask_wav, src_latent.shape[-1])
-
-            latents = sample_latents(
-                model, scheduler, omni_extractor, [f"[Edit] [{edit_target}] {prompt}"],
-                source_latents=src_latent, masks=mask_lat,
-                num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
-            )
-            latents = latents * (1.0 / vae_scale_factor)
-            edit_duration = src_wav.shape[-1] / vae_sample_rate
-            decode_and_save(audio_vae, latents, [edit_duration], [out_path], sample_rate=vae_sample_rate)
-
-        elif mode == "Clone Voice":
-            # Reference clip + text to speak in that voice.
-            if not input_audio_path:
-                raise gr.Error("Clone Voice mode requires a reference audio track.")
-
-            # 3.0s matches REF_DURATION, the reference-clip length the model was trained with.
-            ref_wav = load_ref_audio(input_audio_path, target_sr=vae_sample_rate,
-                                      max_ref_duration=3.0, device=DEVICE)
-            ref_audio_duration = ref_wav.shape[-1] / vae_sample_rate
-            if ref_audio_duration + 1.0 > duration:  # 1.0s floor so some cloned speech always fits
-                raise gr.Error(f"Reference clip ({ref_audio_duration:.1f}s) leaves under 1s for "
-                                f"the cloned speech at Duration={duration:.1f}s. Increase Duration.")
-
-            # Get the ref transcript (typed one wins over auto-transcription), then
-            # combine it with what the user wants said next into one prompt.
-            # Exclude load_ref_audio's trailing silence pad (tail_pad_s=0.1) so Whisper
-            # doesn't hallucinate tokens over silence.
-            speech_samples = max(ref_wav.shape[-1] - int(0.1 * vae_sample_rate), 1)
-            resolved_ref_text = ref_text.strip() or transcribe_ref_audio(
-                ref_wav[..., :speech_samples], sr=vae_sample_rate,
-            )
-            combined_text = join_ref_target_text(resolved_ref_text, prompt)
-            full_prompt = f"[Speech with voice] {combined_text}"
-
-            # Build one waveform [ref audio | silence] and encode it as a single clip —
-            # the model generates the target portion conditioned on the ref portion.
-            total_wav_len = int(duration * vae_sample_rate)
-            ref_wav_1d = ref_wav.squeeze(0) if ref_wav.dim() == 2 else ref_wav
-            ref_wav_len = min(ref_wav_1d.shape[-1], total_wav_len)
-            source_wav = torch.zeros(1, total_wav_len, device=DEVICE)
-            source_wav[:, :ref_wav_len] = ref_wav_1d[:ref_wav_len]
-            source_latent = audio_vae.wrapped_encode(source_wav) * vae_scale_factor
-
-            mask_wav = torch.zeros(1, 1, total_wav_len, device=DEVICE)
-            mask_wav[:, :, :ref_wav_len] = 2.0  # 2 = reference region, 0 = target (see sample_latents docstring)
-            mask_latent = torch.nn.functional.interpolate(
-                mask_wav, size=source_latent.shape[-1]
-            ).to(torch.bfloat16)
-
-            latents = sample_latents(
-                model, scheduler, omni_extractor, [full_prompt],
-                source_latents=source_latent, masks=mask_latent,
-                num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
-            )
-            latents_full = latents * (1.0 / vae_scale_factor)
-
-            # Decode the full ref+target latent, then crop out just the target —
-            # the decoder needs the ref portion as context to decode cleanly.
-            ref_samples = int(ref_audio_duration * vae_sample_rate)
-            decode_and_save_full(
-                audio_vae, latents_full, ref_samples,
-                [duration - ref_audio_duration], [out_path], sample_rate=vae_sample_rate,
-            )
-
-        else:
-            raise gr.Error(f"Unknown mode: {mode}")
+    model, omni_extractor, audio_vae, vae_scale_factor, vae_sample_rate, gen_target_frames, scheduler = get_variant(model_variant)
+
+    if mode == "Generate":
+        # Text-to-audio/speech, no input audio involved. sound_type picks which of
+        # the model's trained prompt templates to build (see README task table) —
+        # voice/background only apply to the Speech templates, unused otherwise.
+        if sound_type == "Sound Effect":
+            tagged_prompt = f"[Audio] {prompt}"
+        elif sound_type == "Speech":
+            tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}"'
+        else:  # "Speech + Background"
+            tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}" [Audio] {background}'
+
+        latents = sample_latents(
+            model, scheduler, omni_extractor, [tagged_prompt],
+            num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
+            target_frames=gen_target_frames,
+        )
+        latents = latents * (1.0 / vae_scale_factor)
+        decode_and_save(audio_vae, latents, [duration], [out_path], sample_rate=vae_sample_rate)
+
+    elif mode == "Edit":
+        # Source audio + instruction. Mask is all-ones: the whole clip is editable,
+        # source_latents just gives the model something to condition on.
+        # sound_type again picks [Audio] vs [Speech] as the edit's sub-tag; "Speech
+        # + Background" isn't a real edit template, so it also falls back to [Audio].
+        if not input_audio_path:
+            raise gr.Error("Edit mode requires an input audio track.")
+        edit_target = "Speech" if sound_type == "Speech" else "Audio"
+        tgt_wav_len = int(MAX_AUDIO_DURATION * vae_sample_rate)
+        src_wav = load_source_audio(input_audio_path, target_sr=vae_sample_rate,
+                                     target_length=tgt_wav_len, device=DEVICE)
+        src_latent = audio_vae.wrapped_encode(src_wav) * vae_scale_factor
+        mask_wav = make_edit_mask(src_wav.shape[-1], DEVICE)
+        mask_lat = downsample_mask(mask_wav, src_latent.shape[-1])
+
+        latents = sample_latents(
+            model, scheduler, omni_extractor, [f"[Edit] [{edit_target}] {prompt}"],
+            source_latents=src_latent, masks=mask_lat,
+            num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
+        )
+        latents = latents * (1.0 / vae_scale_factor)
+        edit_duration = src_wav.shape[-1] / vae_sample_rate
+        decode_and_save(audio_vae, latents, [edit_duration], [out_path], sample_rate=vae_sample_rate)
+
+    elif mode == "Clone Voice":
+        # Reference clip + text to speak in that voice.
+        if not input_audio_path:
+            raise gr.Error("Clone Voice mode requires a reference audio track.")
+
+        # 3.0s matches REF_DURATION, the reference-clip length the model was trained with.
+        ref_wav = load_ref_audio(input_audio_path, target_sr=vae_sample_rate,
+                                  max_ref_duration=3.0, device=DEVICE)
+        ref_audio_duration = ref_wav.shape[-1] / vae_sample_rate
+        if ref_audio_duration + 1.0 > duration:  # 1.0s floor so some cloned speech always fits
+            raise gr.Error(f"Reference clip ({ref_audio_duration:.1f}s) leaves under 1s for "
+                            f"the cloned speech at Duration={duration:.1f}s. Increase Duration.")
+
+        # Get the ref transcript (typed one wins over auto-transcription), then
+        # combine it with what the user wants said next into one prompt.
+        # Exclude load_ref_audio's trailing silence pad (tail_pad_s=0.1) so Whisper
+        # doesn't hallucinate tokens over silence.
+        speech_samples = max(ref_wav.shape[-1] - int(0.1 * vae_sample_rate), 1)
+        resolved_ref_text = ref_text.strip() or transcribe_ref_audio(
+            ref_wav[..., :speech_samples], sr=vae_sample_rate,
+        )
+        combined_text = join_ref_target_text(resolved_ref_text, prompt)
+        full_prompt = f"[Speech with voice] {combined_text}"
+
+        # Build one waveform [ref audio | silence] and encode it as a single clip —
+        # the model generates the target portion conditioned on the ref portion.
+        total_wav_len = int(duration * vae_sample_rate)
+        ref_wav_1d = ref_wav.squeeze(0) if ref_wav.dim() == 2 else ref_wav
+        ref_wav_len = min(ref_wav_1d.shape[-1], total_wav_len)
+        source_wav = torch.zeros(1, total_wav_len, device=DEVICE)
+        source_wav[:, :ref_wav_len] = ref_wav_1d[:ref_wav_len]
+        source_latent = audio_vae.wrapped_encode(source_wav) * vae_scale_factor
+
+        mask_wav = torch.zeros(1, 1, total_wav_len, device=DEVICE)
+        mask_wav[:, :, :ref_wav_len] = 2.0  # 2 = reference region, 0 = target (see sample_latents docstring)
+        mask_latent = torch.nn.functional.interpolate(
+            mask_wav, size=source_latent.shape[-1]
+        ).to(torch.bfloat16)
+
+        latents = sample_latents(
+            model, scheduler, omni_extractor, [full_prompt],
+            source_latents=source_latent, masks=mask_latent,
+            num_inference_steps=steps, guidance_scale=guidance, device=DEVICE,
+        )
+        latents_full = latents * (1.0 / vae_scale_factor)
+
+        # Decode the full ref+target latent, then crop out just the target —
+        # the decoder needs the ref portion as context to decode cleanly.
+        ref_samples = int(ref_audio_duration * vae_sample_rate)
+        decode_and_save_full(
+            audio_vae, latents_full, ref_samples,
+            [duration - ref_audio_duration], [out_path], sample_rate=vae_sample_rate,
+        )
+
+    else:
+        raise gr.Error(f"Unknown mode: {mode}")
 
     return out_path