import torch import torchaudio from einops import rearrange import gradio as gr import spaces import os import uuid from pydub import AudioSegment import numpy as np import random import torch from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler '''AUDIO''' # Importing the model-related functions from stable_audio_tools import get_pretrained_model from stable_audio_tools.inference.generation import generate_diffusion_cond # Load the model outside of the GPU-decorated function def load_model(): print("Loading model...") model, model_config = get_pretrained_model("stabilityai/stable-audio-open-1.0") print("Model loaded successfully.") return model, model_config # Function to set up, generate, and process the audio @spaces.GPU(duration=120) # Allocate GPU only when this function is called def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7): print(f"Prompt received: {prompt}") print(f"Settings: Duration={seconds_total}s, Steps={steps}, CFG Scale={cfg_scale}") device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # Fetch the Hugging Face token from the environment variable hf_token = os.getenv('HF_TOKEN') print(f"Hugging Face token: {hf_token}") # Use pre-loaded model and configuration model, model_config = load_model() sample_rate = model_config["sample_rate"] sample_size = model_config["sample_size"] print(f"Sample rate: {sample_rate}, Sample size: {sample_size}") model = model.to(device) print("Model moved to device.") # Set up text and timing conditioning conditioning = [{ "prompt": prompt, "seconds_start": 0, "seconds_total": seconds_total }] print(f"Conditioning: {conditioning}") # Generate stereo audio print("Generating audio...") output = generate_diffusion_cond( model, steps=steps, cfg_scale=cfg_scale, conditioning=conditioning, sample_size=sample_size, sigma_min=0.3, sigma_max=500, sampler_type="dpmpp-3m-sde", device=device ) print("Audio generated.") # Rearrange audio batch to a single sequence output = rearrange(output, "b d n -> d (b n)") print("Audio rearranged.") # Peak normalize, clip, convert to int16 output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu() print("Audio normalized and converted.") # Generate a unique filename for the output unique_filename = f"output_{uuid.uuid4().hex}.wav" print(f"Saving audio to file: {unique_filename}") # Save to file torchaudio.save(unique_filename, output, sample_rate) print(f"Audio saved: {unique_filename}") # Convert WAV to MP3 using pydub without ffmpeg audio = AudioSegment.from_wav(unique_filename) full_path_mp3 = unique_filename.replace('wav', 'mp3') audio.export(full_path_mp3, format="mp3") print(f"Audio converted and saved to MP3: {full_path_mp3}") # Return the path to the generated audio file return full_path_mp3 '''DIFFUSION''' device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float16 repo = "stabilityai/stable-diffusion-3-medium-diffusers" pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1344 @spaces.GPU def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)): if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator().manual_seed(seed) image = pipe( prompt = prompt, negative_prompt = negative_prompt, guidance_scale = guidance_scale, num_inference_steps = num_inference_steps, width = width, height = height, generator = generator ).images[0] return image, seed ''' # Setting up the Gradio Interface interface = gr.Interface( fn=generate_audio, inputs=[ gr.Textbox(label="Prompt", placeholder="Enter your text prompt here"), gr.Slider(0, 47, value=30, label="Duration in Seconds"), gr.Slider(10, 150, value=100, step=10, label="Number of Diffusion Steps"), gr.Slider(1, 15, value=7, step=0.1, label="CFG Scale") ], outputs=gr.Audio(type="filepath", label="Generated Audio"), title="Stable Audio Generator", description="Generate variable-length stereo audio at 44.1kHz from text prompts using Stable Audio Open 1.0.", )''' with gr.Blocks() as demo: with gr.Tab("SD3"): with gr.Column: gr.Markdown(f""" # Demo [Stable Diffusion 3 Medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) Learn more about the [Stable Diffusion 3 series](https://stability.ai/news/stable-diffusion-3). Try on [Stability AI API](https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post), [Stable Assistant](https://stability.ai/stable-assistant), or on Discord via [Stable Artisan](https://stability.ai/stable-artisan). Run locally with [ComfyUI](https://github.com/comfyanonymous/ComfyUI) or [diffusers](https://github.com/huggingface/diffusers) """) with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0) result = gr.Image(label="Result", show_label=False) with gr.Accordion("Advanced Settings", open=False): negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", ) seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024, ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024, ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=5.0, ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=28, ) with gr.Tab("Audio"): audio_prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here") audio_duration = gr.Slider(0, 47, value=30, label="Duration in Seconds") audio_steps = gr.Slider(10, 150, value=100, step=10, label="Number of Diffusion Steps") audio_cfg = gr.Slider(1, 15, value=7, step=0.1, label="CFG Scale") audio_process_button = gr.Button("Process Audio") audio_output = gr.Audio(type="filepath", label="Generated Audio") audio_process_button.click(generate_audio, [audio_prompt, audio_duration, audio_steps, audio_cfg], [audio_output]) gr.on( triggers=[run_button.click, prompt.submit, negative_prompt.submit], fn = infer, inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], outputs = [result, seed] ) # Pre-load the model to avoid multiprocessing issues model, model_config = load_model() demo.launch()