Switch to inference api

.gitignore

@@ -165,3 +165,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+*.DS_STORE

app.py

@@ -4,6 +4,13 @@ import os
 from torch import autocast
 from diffusers import StableDiffusionPipeline
 from moviepy.editor import AudioFileClip, ImageClip
+from pathlib import Path
+
+
+output_dir = Path("temp/")
+output_dir.mkdir(exist_ok=True, parents=True)
+os.chdir(output_dir)
+
 
 def process_inputs(prompt, audio):
     image = get_stable_diffusion_image(prompt)
@@ -12,8 +19,8 @@ def process_inputs(prompt, audio):
 
 
 def add_static_image_to_audio(image, audio):
-    """Create and save a video file to `output_path` after 
-    combining a static image that is located in `image_path` 
+    """Create and save a video file to `output_path` after
+    combining a static image that is located in `image_path`
     with an audio file in `audio_path`"""
     # create the audio clip object
     audio_clip = AudioFileClip(audio)
@@ -26,21 +33,20 @@ def add_static_image_to_audio(image, audio):
     # set the FPS to 1
     video_clip.fps = 1
     # write the resuling video clip
-    path = "temp/video_out.mp4"
+    path = "out.mp4"
     video_clip.write_videofile(path)
     return path
 
+
 def get_stable_diffusion_image(prompt):
-    model_id = "CompVis/stable-diffusion-v1-4"
-    device = "cuda"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=os.getenv("HF_TOKEN") if os.getenv("HF_TOKEN")  else True)
-    pipe = pipe.to(device)
-    with autocast("cuda"):
-        image = pipe(prompt, guidance_scale=7.5)["sample"][0]
-        path = "temp/out.jpg"
-        image.save(path)
-        return path
-    
-
-iface = gr.Interface(fn=process_inputs, inputs=["text", gr.Audio(type="filepath")], outputs="video")
+    path = "temp/image_out.png"
+    stable_diffusion = gr.Blocks.load(name="spaces/stabilityai/stable-diffusion")
+    gallery_dir = stable_diffusion(prompt, fn_index=2)
+    # Rename gallery dir to sdout
+    return [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)][0]
+
+
+iface = gr.Interface(
+    fn=process_inputs, inputs=["text", gr.Audio(type="filepath")], outputs="video"
+)
 iface.launch()