optimised for faster generation

app.py

@@ -1,181 +1,130 @@
 #!/usr/bin/env python3
 # ------------------------------------------------------------
-# Whisper Small LS-100 + Stable Diffusion (CPU-only)
-# Works on Hugging Face Spaces free tier (no GPU)
+# Whisper Small LS-100  →  Text  →  Image (CPU optimized)
+# Adds ⚡ Fast / 🎨 Quality mode for Hugging Face Spaces
 # Author: Murali Chandran (codeninja3d)
 # ------------------------------------------------------------
 
-import os
-import gc
-import torch
-import gradio as gr
-
-# Prefer torchaudio but safely fall back to soundfile if TorchCodec backend bites
+import os, gc, torch, gradio as gr
 try:
-    import torchaudio  # may try TorchCodec in new versions
-    _HAS_TORCHAUDIO = True
+    import torchaudio
 except Exception:
-    _HAS_TORCHAUDIO = False
     torchaudio = None
-
-try:
-    import soundfile as sf  # CPU-safe decoder (libsndfile)
-except Exception as e:
-    raise RuntimeError("soundfile (pysoundfile) is required on CPU-only Spaces: pip install soundfile") from e
+import soundfile as sf
 
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 from diffusers import AutoPipelineForText2Image
 
-# -------------------------------
-# Global config & environment
-# -------------------------------
-DEVICE = "cpu"  # free Spaces are CPU-only
+# ------------------------------------------------------------
+# Config
+# ------------------------------------------------------------
+DEVICE = "cpu"
 ASR_MODEL_ID = "codeninja3d/whisper-small-ls100"
-IMG_MODEL_ID = "runwayml/stable-diffusion-v1-5"  # SD 1.5 is lighter for CPU
-DEFAULT_SIZE = 512  # safe default for CPU Spaces
-DEFAULT_STEPS = 15
-DEFAULT_CFG = 7.5
-
-# -------------------------------
-# Safe audio loader (CPU, no TorchCodec)
-# -------------------------------
-def load_audio_cpu(path_or_tuple, target_sr=16000):
-    """Load mono waveform at target_sr as torch.float32 tensor on CPU.
-    Accepts (sr, numpy) tuple from Gradio mic OR a filepath for uploads.
-    """
-    if isinstance(path_or_tuple, tuple):
-        sr, wav_np = path_or_tuple
-        wav = torch.tensor(wav_np, dtype=torch.float32)
-        # Gradio numpy shape: [samples, channels] → make [1, samples] mono
+
+# ------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------
+def load_audio(audio, target_sr=16000):
+    """Load and normalize mono waveform (CPU-safe)."""
+    if isinstance(audio, tuple):
+        sr, arr = audio
+        wav = torch.tensor(arr, dtype=torch.float32)
         if wav.ndim > 1:
             wav = wav.mean(dim=-1, keepdim=True).T
     else:
-        path = path_or_tuple
-        # Try torchaudio first (if present), otherwise soundfile
-        if _HAS_TORCHAUDIO:
+        if torchaudio:
             try:
-                wav, sr = torchaudio.load(path)
+                wav, sr = torchaudio.load(audio)
             except Exception:
-                data, sr = sf.read(path)
+                data, sr = sf.read(audio)
                 wav = torch.tensor(data, dtype=torch.float32)
         else:
-            data, sr = sf.read(path)
+            data, sr = sf.read(audio)
             wav = torch.tensor(data, dtype=torch.float32)
-
     if wav.ndim > 1:
         wav = wav.mean(dim=0, keepdim=True)
+    if sr != target_sr and torchaudio:
+        wav = torchaudio.functional.resample(wav, sr, target_sr)
+    wav = wav / wav.abs().max()
+    return wav.squeeze(0)
 
-    if sr != target_sr:
-        if _HAS_TORCHAUDIO:
-            wav = torchaudio.functional.resample(wav, sr, target_sr)
-        else:
-            # Basic resample with torch if torchaudio not available
-            import torch.nn.functional as F
-            wav = wav.unsqueeze(0)  # [1,1,T]
-            T = wav.shape[-1]
-            new_T = int(T * target_sr / sr)
-            wav = F.interpolate(wav, size=new_T, mode="linear", align_corners=False).squeeze(0)
-        sr = target_sr
-
-    # normalize to [-1, 1]
-    maxv = wav.abs().max().clamp(min=1e-8)
-    wav = wav / maxv
-    return wav.squeeze(0), sr
-
-# -------------------------------
-# Load models
-# -------------------------------
+# ------------------------------------------------------------
+# Load Whisper
+# ------------------------------------------------------------
 processor = WhisperProcessor.from_pretrained(ASR_MODEL_ID)
 asr_model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_ID).to(DEVICE).eval()
 
-pipe = AutoPipelineForText2Image.from_pretrained(
-    IMG_MODEL_ID,
-    torch_dtype=torch.float32,  # CPU
-).to(DEVICE)
-pipe.set_progress_bar_config(disable=True)
+# ------------------------------------------------------------
+# Load both diffusion models lazily (to save RAM)
+# ------------------------------------------------------------
+pipe_cache = {}
+
+def get_pipeline(model_id):
+    if model_id not in pipe_cache:
+        pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float32).to(DEVICE)
+        pipe.set_progress_bar_config(disable=True)
+        pipe_cache[model_id] = pipe
+    return pipe_cache[model_id]
 
-# -------------------------------
-# Core function
-# -------------------------------
-def transcribe_and_generate(audio, size_choice, steps, cfg):
+# ------------------------------------------------------------
+# Main logic
+# ------------------------------------------------------------
+def transcribe_and_generate(audio, mode):
     if audio is None:
         return "⚠️ Please record or upload an audio clip.", None
 
-    # --- Audio preprocessing ---
-    wav, _ = load_audio_cpu(audio, target_sr=16000)
-    inputs = processor.feature_extractor(
-        wav.numpy(), sampling_rate=16000, return_tensors="pt"
-    ).to(DEVICE)
-
+    # --- Transcribe ---
+    wav = load_audio(audio)
+    inputs = processor.feature_extractor(wav.numpy(), sampling_rate=16000, return_tensors="pt").to(DEVICE)
     forced_ids = processor.get_decoder_prompt_ids(language="en", task="transcribe")
 
     with torch.no_grad():
-        ids = asr_model.generate(
-            **inputs,
-            forced_decoder_ids=forced_ids,
-            max_new_tokens=192,
-            num_beams=3,
-            early_stopping=True,
-        )
+        ids = asr_model.generate(**inputs, forced_decoder_ids=forced_ids, max_new_tokens=192, num_beams=3)
     text = processor.tokenizer.batch_decode(ids, skip_special_tokens=True)[0].strip()
 
-    # --- Diffusion generation (CPU) ---
+    # --- Image generation ---
+    if mode == "Fast":
+        model_id = "stabilityai/sd-turbo"
+        width = height = 384
+        steps, cfg = 4, 1.0
+    else:
+        model_id = "runwayml/stable-diffusion-v1-5"
+        width = height = 512
+        steps, cfg = 15, 7.5
+
     prompt = f"photorealistic, soft light, {text}, detailed"
-    width = height = int(size_choice)
+    gc.collect()
+    pipe = get_pipeline(model_id)
 
     try:
         with torch.inference_mode():
-            image = pipe(
-                prompt=prompt,
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=width,
-                height=height,
-            ).images[0]
-        gc.collect()
+            image = pipe(prompt=prompt, num_inference_steps=steps,
+                         guidance_scale=cfg, width=width, height=height).images[0]
         return text, image
     except Exception as e:
-        # Graceful fallback if OOM or other runtime error: retry smaller size
-        if width > 512:
-            try:
-                with torch.inference_mode():
-                    image = pipe(
-                        prompt=prompt,
-                        num_inference_steps=int(steps),
-                        guidance_scale=float(cfg),
-                        width=512,
-                        height=512,
-                    ).images[0]
-                gc.collect()
-                return text + "\n\n(ℹ️ Auto-fallback to 512×512 due to CPU limits)", image
-            except Exception:
-                pass
-        return text + f"\n\n❗ Image generation failed: {type(e).__name__}: {e}", None
+        return f"{text}\\n\\n❗ Image generation failed: {e}", None
 
-# -------------------------------
+# ------------------------------------------------------------
 # Gradio UI
-# -------------------------------
+# ------------------------------------------------------------
 with gr.Blocks(theme="soft") as demo:
-    gr.Markdown("## 🎙️ Whisper → 🖼️ Image (CPU-only demo)\nUpload/record audio; get transcription and a generated image.")
+    gr.Markdown("""
+    ## 🎙️ Whisper → 🖼️ Image Generator  
+    Choose **Fast** (⚡ quick demo) or **Quality** (🎨 detailed, slower).
+    """)
+
     with gr.Row():
         audio_in = gr.Audio(type="numpy", label="🎧 Record or Upload Audio")
-        with gr.Column():
-            size_choice = gr.Dropdown([
-                "384", "512", "640", "768"
-            ], value=str(DEFAULT_SIZE), label="Image Size (square)")
-            steps = gr.Slider(5, 30, value=DEFAULT_STEPS, step=1, label="Diffusion Steps")
-            cfg = gr.Slider(1.0, 12.0, value=DEFAULT_CFG, step=0.5, label="CFG (guidance)")
-    with gr.Row():
-        text_out = gr.Textbox(label="📄 Transcription")
-    img_out = gr.Image(label="🖼️ Generated Image", show_label=True)
+        mode = gr.Radio(["Fast", "Quality"], value="Fast", label="Mode")
 
+    text_out = gr.Textbox(label="📄 Transcription")
+    img_out = gr.Image(label="🖼️ Generated Image")
     run_btn = gr.Button("Transcribe + Generate", variant="primary")
+
     run_btn.click(
         fn=transcribe_and_generate,
-        inputs=[audio_in, size_choice, steps, cfg],
-        outputs=[text_out, img_out]
+        inputs=[audio_in, mode],
+        outputs=[text_out, img_out],
     )
 
-# Limit concurrency for CPU stability
-demo.queue().launch(server_name="0.0.0.0", server_port=7860)
-
+demo.launch(server_name="0.0.0.0", server_port=7860)