initial commit

README.md

@@ -1,8 +1,8 @@
 ---
-title: Stable Diffusion Music Vid Gen
+title: Stable Diffusion Video Gen
 emoji: 💻
-colorFrom: yellow
-colorTo: yellow
+colorFrom: purple
+colorTo: blue
 sdk: gradio
 sdk_version: 3.35.2
 app_file: app.py

app.py

@@ -0,0 +1,367 @@
+# Experimental app to help with the process of generating music videos
+# Requires youtube-dl to be installed
+# pip install youtube-dl
+
+import gradio as gr
+import librosa
+from pathlib import Path
+import numpy as np
+import random
+from io import BytesIO
+import soundfile as sf
+from matplotlib import pyplot as plt
+
+from stable_diffusion_videos import StableDiffusionWalkPipeline, generate_images, get_timesteps_arr
+
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import LMSDiscreteScheduler
+from diffusers.utils.import_utils import is_xformers_available
+import torch
+import youtube_dl
+import os
+
+pipe = StableDiffusionWalkPipeline.from_pretrained(
+    'runwayml/stable-diffusion-v1-5',
+    torch_dtype=torch.float16,
+    safety_checker=None,
+    vae=AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda"),
+    scheduler=LMSDiscreteScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+    )
+).to("cuda")
+
+if is_xformers_available():
+    pipe.enable_xformers_memory_efficient_attention()
+
+def download_example_clip(url, output_dir='./', output_filename='%(title)s.%(ext)s'):
+    if (Path(output_dir) / output_filename).exists():
+        return str(Path(output_dir) / output_filename)
+
+    files_before = os.listdir(output_dir) if os.path.exists(output_dir) else []
+    ydl_opts = {
+        'outtmpl': str(Path(output_dir) / output_filename),
+        'format': 'bestaudio',
+        'extract-audio': True,
+        'audio-format': 'mp3',
+        'audio-quality': 0,
+    }
+    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+        ydl.download([url])
+    
+    files_after = os.listdir(output_dir)
+    return str(Path(output_dir) / list(set(files_after) - set(files_before))[0])
+    
+def audio_data_to_buffer(y, sr):
+    audio_filepath = BytesIO()
+    audio_filepath.name = 'audio.wav'
+    sf.write(audio_filepath, y, samplerate=sr, format='WAV')
+    audio_filepath.seek(0)
+    return audio_filepath
+
+
+def plot_array(y):
+    fig = plt.figure()
+    x = np.arange(y.shape[0]) 
+    plt.title("Line graph") 
+    plt.xlabel("X axis") 
+    plt.ylabel("Y axis") 
+    plt.plot(x, y, color ="red") 
+    plt.savefig('timesteps_chart.png')
+    return fig
+
+def on_slice_btn_click(audio, audio_start_sec, duration, fps, smooth, margin):
+    if audio is None:
+        return [
+            gr.update(visible=False),
+            gr.update(visible=False),
+        ]
+
+    y, sr = librosa.load(audio, offset=audio_start_sec, duration=duration)
+    T = get_timesteps_arr(
+        audio_data_to_buffer(y, sr),
+        0,
+        duration,
+        fps=fps,
+        margin=margin,
+        smooth=smooth,
+    )
+    return [gr.update(value=(sr, y), visible=True), gr.update(value=plot_array(T), visible=True)]
+
+def on_audio_change_or_clear(audio):
+    if audio is None:
+        return [
+            gr.update(visible=False),
+            gr.update(visible=False)
+        ]
+    
+    duration = librosa.get_duration(filename=audio)
+    return [
+        gr.update(maximum=int(duration), visible=True),
+        gr.update(maximum=int(min(10, duration)), visible=True)
+    ]
+
+def on_update_weight_settings_btn_click(sliced_audio, duration, fps, smooth, margin):
+    if sliced_audio is None:
+        return gr.update(visible=False)
+
+    T = get_timesteps_arr(
+        sliced_audio,
+        0,
+        duration,
+        fps=fps,
+        margin=margin,
+        smooth=smooth,
+    )
+    return gr.update(value=plot_array(T), visible=True)
+
+
+def on_generate_images_btn_click(
+    prompt_a,
+    prompt_b,
+    seed_a,
+    seed_b,
+    output_dir,
+    num_inference_steps,
+    guidance_scale,
+    height,
+    width,
+    upsample,
+):
+    output_dir = Path(output_dir) / 'images'
+
+    if seed_a == -1:
+        seed_a = random.randint(0, 9999999)
+    if seed_b == -1:
+        seed_b = random.randint(0, 9999999)
+
+    image_a_fpath = generate_images(
+        pipe,
+        prompt_a,
+        seeds=[seed_a],
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        upsample=upsample,
+        output_dir=output_dir
+    )[0]
+    image_b_fpath = generate_images(
+        pipe,
+        prompt_b,
+        seeds=[seed_b],
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        upsample=upsample,
+        output_dir=output_dir
+    )[0]
+
+    return [
+        gr.update(value=image_a_fpath, visible=True),
+        gr.update(value=image_b_fpath, visible=True),
+        gr.update(value=seed_a),
+        gr.update(value=seed_b),
+    ]
+
+def on_generate_music_video_btn_click(
+    audio_filepath,
+    audio_start_sec,
+    duration,
+    fps,
+    smooth,
+    margin,
+    prompt_a,
+    prompt_b,
+    seed_a,
+    seed_b,
+    batch_size,
+    output_dir,
+    num_inference_steps,
+    guidance_scale,
+    height,
+    width,
+    upsample,
+):
+
+    if audio_filepath is None:
+        return gr.update(visible=False)
+
+    video_filepath = pipe.walk(
+        prompts=[prompt_a, prompt_b],
+        seeds=[seed_a, seed_b],
+        num_interpolation_steps=int(duration * fps),
+        output_dir=output_dir,
+        fps=fps,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        upsample=upsample,
+        batch_size=batch_size,
+        audio_filepath=audio_filepath,
+        audio_start_sec=audio_start_sec,
+        margin=margin,
+        smooth=smooth,
+    )
+    return gr.update(value=video_filepath, visible=True)
+
+
+audio_start_sec = gr.Slider(0, 10, 0, step=1, label="Start (sec)", interactive=True)
+duration = gr.Slider(0, 10, 1, step=1, label="Duration (sec)", interactive=True)
+slice_btn = gr.Button("Slice Audio")
+
+sliced_audio = gr.Audio(type='filepath')
+wav_plot = gr.Plot(label="Interpolation Weights Per Frame")
+
+fps = gr.Slider(1, 60, 12, step=1, label="FPS", interactive=True)
+smooth = gr.Slider(0, 1, 0.0, label="Smoothing", interactive=True)
+margin = gr.Slider(1.0, 20.0, 1.0, step=0.5, label="Margin Max", interactive=True)
+update_weight_settings_btn = gr.Button("Update Interpolation Weights")
+
+prompt_a = gr.Textbox(value='blueberry spaghetti', label="Prompt A")
+prompt_b = gr.Textbox(value='strawberry spaghetti', label="Prompt B")
+seed_a = gr.Number(-1, label="Seed A", precision=0, interactive=True)
+seed_b = gr.Number(-1, label="Seed B", precision=0, interactive=True)
+generate_images_btn = gr.Button("Generate Images")
+image_a = gr.Image(visible=False, label="Image A")
+image_b = gr.Image(visible=False, label="Image B")
+
+batch_size = gr.Slider(1, 32, 1, step=1, label="Batch Size", interactive=True)
+generate_music_video_btn = gr.Button("Generate Music Video")
+video = gr.Video(visible=False, label="Video")
+
+STEP_1_MARKDOWN = """
+## 1. Upload Some Audio
+
+Upload an audio file to use as the source for the music video.
+"""
+
+STEP_2_MARKDOWN = """
+## 2. Slice Portion of Audio for Generated Clip
+
+Here you can slice a portion of the audio to use for the generated music video. The longer the audio, the more frames will be generated (which will take longer).
+
+I suggest you use this app to make music videos in segments of 5-10 seconds at a time. Then, you can stitch the videos together using a video editor or ffmpeg later.
+
+**Warning**: If your audio file is short, I do no check that the duration you chose is not longer than the audio. It may cause some issues, so just be mindful of that.
+"""
+
+STEP_3_MARKDOWN = """
+## 3. Set Interpolation Weight Settings
+
+This section lets you play with the settings used to configure how we move through the latent space given the audio you sliced.
+
+If you look at the graph on the right, you'll see in the X-axis how many frames. The Y-axis is the weight of Image A as we move through the latent space.
+
+If you listen to the audio slice and look at the graph, you should see bumps at points where the audio energy is high (in our case, percussive energy).
+"""
+
+STEP_4_MARKDOWN = """
+## 4. Select Prompts, Seeds, Settings, and Generate Images
+
+Here you can select the settings for image generation.
+
+Then, you can select prompts and seeds for generating images.
+
+  - Image A will be first frame of the generated video.
+  - Image B will be last frame of the generated video.
+  - The video will be generated by interpolating between the two images using the audio you provided.
+
+If you set the seeds to -1, a random seed will be used and saved for you, so you can explore different images given the same prompt.
+"""
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(STEP_1_MARKDOWN)
+    audio = gr.Audio(type='filepath', interactive=True)
+    gr.Examples(
+        [
+            download_example_clip(
+                url='https://soundcloud.com/nateraw/thoughts',
+                output_dir='./music',
+                output_filename='thoughts.mp3'
+            )
+        ],
+        inputs=audio,
+        outputs=[audio_start_sec, duration],
+        fn=on_audio_change_or_clear,
+        cache_examples=False
+    )
+    audio.change(on_audio_change_or_clear, audio, [audio_start_sec, duration])
+    audio.clear(on_audio_change_or_clear, audio, [audio_start_sec, duration])
+
+    gr.Markdown(STEP_2_MARKDOWN)
+    audio_start_sec.render()
+    duration.render()
+    slice_btn.render()
+
+    slice_btn.click(on_slice_btn_click, [audio, audio_start_sec, duration, fps, smooth, margin], [sliced_audio, wav_plot])
+    sliced_audio.render()
+
+    gr.Markdown(STEP_3_MARKDOWN)
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            fps.render()
+            smooth.render()
+            margin.render()
+            update_weight_settings_btn.render()
+            update_weight_settings_btn.click(
+                on_update_weight_settings_btn_click,
+                [sliced_audio, duration, fps, smooth, margin],
+                wav_plot
+            )
+        with gr.Column(scale=3):
+            wav_plot.render()
+
+    gr.Markdown(STEP_4_MARKDOWN)
+    
+    with gr.Accordion("Additional Settings", open=False):
+        output_dir = gr.Textbox(value='./dreams', label="Output Directory")
+        num_inference_steps = gr.Slider(1, 200, 50, step=10, label="Diffusion Inference Steps", interactive=True)
+        guidance_scale = gr.Slider(1.0, 25.0, 7.5, step=0.5, label="Guidance Scale", interactive=True)
+        height = gr.Slider(512, 1024, 512, step=64, label="Height", interactive=True)
+        width = gr.Slider(512, 1024, 512, step=64, label="Width", interactive=True)
+        upsample = gr.Checkbox(value=False, label="Upsample with Real-ESRGAN")
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            prompt_a.render()
+        with gr.Column(scale=1):
+            seed_a.render()
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            prompt_b.render()
+        with gr.Column(scale=1):
+            seed_b.render()
+
+    generate_images_btn.render()
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_a.render()
+        with gr.Column(scale=1):
+            image_b.render()
+
+    generate_images_btn.click(
+        on_generate_images_btn_click,
+        [prompt_a, prompt_b, seed_a, seed_b, output_dir, num_inference_steps, guidance_scale, height, width, upsample],
+        [image_a, image_b, seed_a, seed_b]
+    )
+
+    gr.Markdown("## 5. Generate Music Video")
+    # TODO - add equivalent code snippet to generate music video
+    batch_size.render()
+    generate_music_video_btn.render()
+    generate_music_video_btn.click(
+        on_generate_music_video_btn_click,
+        [audio, audio_start_sec, duration, fps, smooth, margin, prompt_a, prompt_b, seed_a, seed_b, batch_size, output_dir, num_inference_steps, guidance_scale, height, width, upsample],
+        video
+    )
+    video.render()
+
+
+if __name__ == '__main__':
+    demo.launch(debug=True)

requirements.txt

@@ -0,0 +1,10 @@
+# A10G needs cuda 11.3
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch==1.12.0+cu113
+torchvision==0.13.0+cu113
+
+# Stable diffusion videos repo with the realesrgan extra for upsampling
+stable-diffusion-videos[realesrgan]
+
+# Youtube-dl for downloading examples
+youtube-dl