Update gradio_app.py

Added (Float 16) Checkbox for Low VRAM inferance
This commit is contained in:
drbaph
2024-06-06 21:48:52 +01:00
committed by GitHub
parent e2b224d04f
commit 6fbfe602bf

View File

@@ -10,7 +10,7 @@ from datetime import datetime
import gradio as gr import gradio as gr
# Define the function to generate audio based on a prompt # Define the function to generate audio based on a prompt
def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type): def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type, model_half):
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
# Download model # Download model
@@ -20,6 +20,16 @@ def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_ti
model = model.to(device) model = model.to(device)
# Print model data type before conversion
print("Model data type before conversion:", next(model.parameters()).dtype)
# Convert model to float16 if model_half is True
if model_half:
model = model.to(torch.float16)
# Print model data type after conversion
print("Model data type after conversion:", next(model.parameters()).dtype)
# Set up text and timing conditioning # Set up text and timing conditioning
conditioning = [{ conditioning = [{
"prompt": prompt, "prompt": prompt,
@@ -41,11 +51,19 @@ def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_ti
seed=seed seed=seed
) )
# Print output data type
print("Output data type:", output.dtype)
# Rearrange audio batch to a single sequence # Rearrange audio batch to a single sequence
output = rearrange(output, "b d n -> d (b n)") output = rearrange(output, "b d n -> d (b n)")
# Peak normalize, clip, convert to int16, and save to temporary file # Peak normalize, clip, and convert to int16 directly if model_half is used
output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu() output = output.div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767)
if model_half:
output = output.to(torch.int16).cpu()
else:
output = output.to(torch.float32).to(torch.int16).cpu()
torchaudio.save("temp_output.wav", output, sample_rate) torchaudio.save("temp_output.wav", output, sample_rate)
# Convert to MP3 format using pydub # Convert to MP3 format using pydub
@@ -74,7 +92,7 @@ def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_ti
return full_path return full_path
def audio_generator(prompt, sampler_type, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed): def audio_generator(prompt, sampler_type, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, model_half):
try: try:
print("Generating audio with parameters:") print("Generating audio with parameters:")
print("Prompt:", prompt) print("Prompt:", prompt)
@@ -85,8 +103,9 @@ def audio_generator(prompt, sampler_type, steps, cfg_scale, sigma_min, sigma_max
print("Sigma Max:", sigma_max) print("Sigma Max:", sigma_max)
print("Generation Time:", generation_time) print("Generation Time:", generation_time)
print("Seed:", seed) print("Seed:", seed)
print("Model Half Precision:", model_half)
filename = generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type) filename = generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type, model_half)
return gr.Audio(filename), f"Generated: {filename}" return gr.Audio(filename), f"Generated: {filename}"
except Exception as e: except Exception as e:
return str(e) return str(e)
@@ -106,16 +125,13 @@ sampler_dropdown = gr.Dropdown(
], ],
value="dpmpp-3m-sde" value="dpmpp-3m-sde"
) )
steps_slider = gr.Slider(minimum=0, maximum=200, label="Steps", step=1) steps_slider = gr.Slider(minimum=0, maximum=200, label="Steps", step=1, value=100)
steps_slider.value = 100 # Set the default value here cfg_scale_slider = gr.Slider(minimum=0, maximum=15, label="CFG Scale", step=0.1, value=7)
cfg_scale_slider = gr.Slider(minimum=0, maximum=15, label="CFG Scale", step=0.1)
cfg_scale_slider.value = 7 # Set the default value here
sigma_min_slider = gr.Slider(minimum=0, maximum=50, label="Sigma Min", step=0.1, value=0.3) sigma_min_slider = gr.Slider(minimum=0, maximum=50, label="Sigma Min", step=0.1, value=0.3)
sigma_max_slider = gr.Slider(minimum=0, maximum=1000, label="Sigma Max", step=1, value=500) sigma_max_slider = gr.Slider(minimum=0, maximum=1000, label="Sigma Max", step=0.1, value=500)
generation_time_slider = gr.Slider(minimum=0, maximum=47, label="Generation Time (seconds)", step=1) generation_time_slider = gr.Slider(minimum=0, maximum=47, label="Generation Time (seconds)", step=1, value=47)
generation_time_slider.value = 47 # Set the default value here seed_slider = gr.Slider(minimum=-1, maximum=999999, label="Seed", step=1, value=123456)
seed_slider = gr.Slider(minimum=-1, maximum=999999, label="Seed", step=1) model_half_checkbox = gr.Checkbox(label="Low VRAM (float16)", value=False)
seed_slider.value = 77212 # Set the default value here
output_textbox = gr.Textbox(label="Output") output_textbox = gr.Textbox(label="Output")
@@ -124,7 +140,7 @@ description = "[Github Repository](https://github.com/Saganaki22/StableAudioWebU
gr.Interface( gr.Interface(
audio_generator, audio_generator,
[prompt_textbox, sampler_dropdown, steps_slider, cfg_scale_slider, sigma_min_slider, sigma_max_slider, generation_time_slider, seed_slider], [prompt_textbox, sampler_dropdown, steps_slider, cfg_scale_slider, sigma_min_slider, sigma_max_slider, generation_time_slider, seed_slider, model_half_checkbox],
[gr.Audio(), output_textbox], [gr.Audio(), output_textbox],
title=title, title=title,
description=description description=description