Update app.py

app.py

@@ -1,21 +1,18 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler, UNet2DConditionModel
+from diffusers import StableAudioPipeline
 from huggingface_hub import hf_hub_download
 import spaces
-from PIL import Image
-import requests
 from translatepy import Translator
 import numpy as np
 import random
-
+import soundfile as sf
 
 translator = Translator()
 
 # Constants
-model = "Corcelio/mobius"
-vae_model = "madebyollin/sdxl-vae-fp16-fix"
-MAX_SEED = np.iinfo(np.int32).max
+model = "stabilityai/stable-audio-open-1.0"
+# MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
 .gradio-container {
@@ -41,10 +38,9 @@ vae = AutoencoderKL.from_pretrained(
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
-    unet = UNet2DConditionModel.from_pretrained(model, subfolder="unet", torch_dtype=torch.float16).to("cuda")
-    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, unet=unet, torch_dtype=torch.float16).to("cuda")
-
-pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe = StableAudioPipeline.from_pretrained(
+        model,
+        torch_dtype=torch.float16).to("cuda")
 
 
 # Function 
@@ -52,131 +48,51 @@ pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.conf
 def generate_image(
     prompt,
     negative="low quality",
-    width=1024,
-    height=1024,
-    seed: int = -1,
-    nums=1,
-    scale=1.5,
-    steps=30,
-    clip=3):
+    second: float = 10.0):
     
-    if seed == -1:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-    generator = torch.Generator().manual_seed(seed)
+    # if seed == -1:
+    #     seed = random.randint(0, MAX_SEED)
+    # seed = int(seed)
+    # generator = torch.Generator().manual_seed(seed)
     
     prompt = str(translator.translate(prompt, 'English'))
 
     print(f'prompt:{prompt}')
 
-    
-    image = pipe(
+    audio = pipe(
         prompt, 
         negative_prompt=negative, 
-        width=width,
-        height=height,
-        guidance_scale=scale,
-        generator = generator,
-        num_inference_steps=steps,
-        num_images_per_prompt=nums,
-        clip_skip=clip,
-    ).images
-    return image, seed
-
-
-examples = [
-    "a cat eating a piece of cheese",
-    "a ROBOT riding a BLUE horse on Mars, photorealistic",
-    "Ironman VS Hulk, ultrarealistic",
-    "a CUTE robot artist painting on an easel",
-    "Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k",
-    "An alien holding sign board contain word 'Flash', futuristic, neonpunk",
-    "Kids going to school, Anime style"
-]
+        audio_end_in_s=second,
+    ).audios
 
+    os.makedirs("outputs", exist_ok=True)
+    base_count = len(glob(os.path.join("outputs", "*.mp4")))
+    audio_path = os.path.join("outputs", f"{base_count:06d}.wav")
+    
+    sf.write(audio_path, audio[0].T.float().cpu().numpy(), pipe.vae.samping_rate)
+    
+    return audio_path
 
 # Gradio Interface
 
-with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
-    gr.HTML("<h1><center>Mobius💠</center></h1>")
-    gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/mobius'>mobius</a> text-to-image generation</center><br><center>Adding default prompts to enhance.</center></p>")
-    with gr.Group():
-        with gr.Row():
-            prompt = gr.Textbox(label='Enter Your Prompt(Multi-Languages)', value="best quality, HD, aesthetic", scale=6)
-            submit = gr.Button(scale=1, variant='primary')
-    img = gr.Gallery(label='Mobius Generated Image',columns = 1, preview=True)
-    with gr.Accordion("Advanced Options", open=False):
-        with gr.Row():
-            negative = gr.Textbox(label="Negative prompt", value="low quality, ugly, blurry, poor face, bad anatomy")
-        with gr.Row():
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=1280,
-                step=8,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=1280,
-                step=8,
-                value=1024,
-            )
-        with gr.Row():
-            seed = gr.Slider(
-                label="Seed (-1 Get Random)",
-                minimum=-1,
-                maximum=MAX_SEED,
-                step=1,
-                value=-1,
-                scale=2,
-            )
-            nums = gr.Slider(
-                label="Image Numbers",
-                minimum=1,
-                maximum=4,
-                step=1,
-                value=1,
-                scale=1,
-            ) 
-        with gr.Row():
-            scale = gr.Slider(
-                label="Guidance",
-                minimum=3.5,
-                maximum=7,
-                step=0.1,
-                value=7,
-            )
-            steps = gr.Slider(
-                label="Steps",
-                minimum=1,
-                maximum=50,
-                step=1,
-                value=30,
-            )
-            clip = gr.Slider(
-                label="Clip Skip",
-                minimum=1,
-                maximum=10,
-                step=1,
-                value=3,
-            )
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=[img, seed],
-        fn=generate_image,
-        cache_examples="lazy",
-    )
-
-    prompt.submit(fn=generate_image,
-                 inputs=[prompt, negative, width, height, seed, nums, scale, steps, clip],
-                 outputs=[img, seed],
-                 )
-    submit.click(fn=generate_image,
-                 inputs=[prompt, negative, width, height, seed, nums, scale, steps, clip],
-                 outputs=[img, seed],
-                 )
-    
-demo.queue().launch()
+with gr.Blocks(theme='soft', css=css, title="Stable Audio Open") as iface:
+    with gr.Accordion(""):
+        gr.Markdown(DESCRIPTION)
+    with gr.Row():
+        output = gr.Audio(label="Podcast", type="filepath", interactive=False, autoplay=True, elem_classes="audio")  # Create an output textbox
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", placeholder="1000 BPM percussive sound of water drops")
+    with gr.Row():
+        negative = gr.Textbox(label="Negative prompt", placeholder="Low quality")
+        second =  gr.Slider(5.0, 60.0, value=10.0, label="Second", step=0.1),
+    with gr.Row():
+        submit_btn = gr.Button("🚀 Send")  # Create a submit button
+        clear_btn = gr.ClearButton(output_box, value="🗑️ Clear") # Create a clear button
+
+    # Set up the event listeners
+    submit_btn.click(main, inputs=[prompt, negative, second], outputs=output)
+
+
+#gr.close_all()
+
+iface.queue().launch(show_api=False)  # Launch the Gradio interface