Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| # ------------------------------------------------------------ | |
| # Whisper Small LS-100 → Text → Image (CPU optimized) | |
| # Adds ⚡ Fast / 🎨 Quality mode for Hugging Face Spaces | |
| # Author: Murali Chandran (codeninja3d) | |
| # ------------------------------------------------------------ | |
| import os, gc, torch, gradio as gr | |
| try: | |
| import torchaudio | |
| except Exception: | |
| torchaudio = None | |
| import soundfile as sf | |
| from transformers import WhisperProcessor, WhisperForConditionalGeneration | |
| from diffusers import AutoPipelineForText2Image | |
| # ------------------------------------------------------------ | |
| # Config | |
| # ------------------------------------------------------------ | |
| DEVICE = "cpu" | |
| ASR_MODEL_ID = "codeninja3d/whisper-small-ls100" | |
| # ------------------------------------------------------------ | |
| # 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: | |
| if torchaudio: | |
| try: | |
| wav, sr = torchaudio.load(audio) | |
| except Exception: | |
| data, sr = sf.read(audio) | |
| wav = torch.tensor(data, dtype=torch.float32) | |
| else: | |
| 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) | |
| # ------------------------------------------------------------ | |
| # Load Whisper | |
| # ------------------------------------------------------------ | |
| processor = WhisperProcessor.from_pretrained(ASR_MODEL_ID) | |
| asr_model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_ID).to(DEVICE).eval() | |
| # ------------------------------------------------------------ | |
| # 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] | |
| # ------------------------------------------------------------ | |
| # Main logic | |
| # ------------------------------------------------------------ | |
| def transcribe_and_generate(audio, mode): | |
| if audio is None: | |
| return "⚠️ Please record or upload an audio clip.", None | |
| # --- 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) | |
| text = processor.tokenizer.batch_decode(ids, skip_special_tokens=True)[0].strip() | |
| # --- 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" | |
| gc.collect() | |
| pipe = get_pipeline(model_id) | |
| try: | |
| with torch.inference_mode(): | |
| image = pipe(prompt=prompt, num_inference_steps=steps, | |
| guidance_scale=cfg, width=width, height=height).images[0] | |
| return text, image | |
| except Exception as e: | |
| return f"{text}\\n\\n❗ Image generation failed: {e}", None | |
| # ------------------------------------------------------------ | |
| # Gradio UI | |
| # ------------------------------------------------------------ | |
| with gr.Blocks(theme="soft") as demo: | |
| 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") | |
| 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, mode], | |
| outputs=[text_out, img_out], | |
| ) | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |