ImageGenTest / app.py
jonloporto's picture
Upload app.py
4eb567d verified
import gradio as gr
from transformers import pipeline
import torch
from diffusers import DiffusionPipeline
# Load speech-to-text model (Whisper)
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base")
# Load image generation model (Stable Diffusion)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe = pipe.to(device)
# Speech-to-text function
def transcribe_audio(audio):
"""Convert audio to text using Whisper"""
if audio is None:
return ""
try:
# Gradio Audio with type="numpy" returns tuple of (sample_rate, audio_data)
if isinstance(audio, tuple):
sample_rate, audio_data = audio
# Create a dictionary with the audio data for the pipeline
result = transcriber({"array": audio_data, "sampling_rate": sample_rate})
else:
result = transcriber(audio)
text = result.get("text", "").strip()
return text if text else "No speech detected"
except Exception as e:
return f"Error transcribing audio: {str(e)}"
# Image generation function
def generate_image_from_text(prompt):
"""Generate an image from a text prompt using Stable Diffusion"""
if not prompt or prompt.strip() == "":
return None, "Please provide a text prompt"
try:
with torch.no_grad():
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
return image, f"✓ Generated image from prompt: '{prompt}'"
except Exception as e:
return None, f"Error generating image: {str(e)}"
# Combined function: speech -> text -> image
def speech_to_image(audio):
"""Convert speech to text, then generate image from the text"""
# Step 1: Convert speech to text
text_prompt = transcribe_audio(audio)
if text_prompt.startswith("Error"):
return None, text_prompt
# Step 2: Generate image from text
image, status = generate_image_from_text(text_prompt)
return image, f"Transcript: '{text_prompt}'\n\n{status}"
# Gradio interface with tabs
with gr.Blocks(title="AI Image Generation from Speech") as demo:
gr.Markdown("# 🎨 AI Image Generation from Speech")
gr.Markdown("Speak your image description, and the AI will generate an image based on your words!")
with gr.Tab("🎤 Speech to Image"):
gr.Markdown("Record or upload audio with your image description")
audio_input = gr.Audio(label="Record Audio", type="numpy")
generate_btn = gr.Button("Generate Image from Speech", variant="primary")
output_image = gr.Image(label="Generated Image")
output_text = gr.Textbox(label="Status", interactive=False)
generate_btn.click(
fn=speech_to_image,
inputs=audio_input,
outputs=[output_image, output_text]
)
with gr.Tab("⌨️ Text to Image"):
gr.Markdown("Or type a description directly")
text_input = gr.Textbox(
label="Enter Image Description",
placeholder="e.g., a beautiful sunset over mountains",
lines=3
)
text_generate_btn = gr.Button("Generate Image", variant="primary")
text_output_image = gr.Image(label="Generated Image")
text_output_status = gr.Textbox(label="Status", interactive=False)
text_generate_btn.click(
fn=generate_image_from_text,
inputs=text_input,
outputs=[text_output_image, text_output_status]
)
# Launch the interface
if __name__ == "__main__":
demo.launch()