File size: 4,841 Bytes
7cd5e3d
 
0c39aa5
 
7cd5e3d
 
 
0c39aa5
7cd5e3d
0c39aa5
7cd5e3d
 
0c39aa5
7cd5e3d
 
 
 
0c39aa5
 
 
 
7cd5e3d
0c39aa5
 
 
 
 
 
 
 
 
7cd5e3d
 
 
0c39aa5
7cd5e3d
0c39aa5
7cd5e3d
0c39aa5
7cd5e3d
 
0c39aa5
7cd5e3d
 
 
0c39aa5
 
 
 
7cd5e3d
0c39aa5
 
 
7cd5e3d
 
 
0c39aa5
 
 
 
 
 
 
 
 
 
 
7cd5e3d
0c39aa5
 
 
 
7cd5e3d
 
 
0c39aa5
 
 
7cd5e3d
 
 
0c39aa5
7cd5e3d
 
0c39aa5
 
 
 
 
 
 
 
 
 
7cd5e3d
0c39aa5
 
7cd5e3d
 
 
0c39aa5
 
7cd5e3d
 
0c39aa5
7cd5e3d
0c39aa5
7cd5e3d
0c39aa5
7cd5e3d
0c39aa5
 
 
 
 
7cd5e3d
 
0c39aa5
7cd5e3d
0c39aa5
 
7cd5e3d
0c39aa5
7cd5e3d
 
0c39aa5
 
7cd5e3d
 
0c39aa5
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#!/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)