add picture generation for background

app.py

@@ -3,30 +3,46 @@ import spaces
 import gradio as gr
 
 from transformers import MusicgenForConditionalGeneration
-model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-sampling_rate = model.config.audio_encoder.sampling_rate
+music_gen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+sampling_rate = music_gen_model.config.audio_encoder.sampling_rate
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+music_gen_model.to(device)
 
 from transformers import AutoProcessor
 processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
 
+from diffusers import DiffusionPipeline
+
+# sd_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+sd_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+sd_pipe.to(device)
+
+
 @spaces.GPU
 def generate_music(desc):
+    
     inputs = processor(text=[desc], padding=True, return_tensors="pt")
-    audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)
+    audio_values = music_gen_model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)
     return sampling_rate, audio_values[0][0].cpu().numpy()
     
-# gr.Interface(fn=generate_music, inputs="text", outputs="audio").launch()
+@spaces.GPU
+def generate_pic(desc):
+    return sd_pipe(prompt=desc).images[0]
 
 with gr.Blocks() as app:
     with gr.Row():
         music_desc = gr.TextArea(label="Music Description")
+        music_pic = gr.Image(label="Music Image(StableDiffusion)")
         music_player = gr.Audio(label="Play My Tune")
 
-    gen_btn = gr.Button("Get Some Tune!!")
-    gen_btn.click(fn=generate_music, inputs=[music_desc], outputs=[music_player])
+
+    gen_pic_btn = gr.Button("Gen Picture")
+    gen_music_btn = gr.Button("Get Some Tune!!")
+
+    gen_pic_btn.click(fn=generate_pic, inputs=[music_desc], outputs=[music_pic])
+    gen_music_btn.click(fn=generate_music, inputs=[music_desc], outputs=[music_player])
     
 if __name__ == '__main__':
     app.launch()

requirements.txt

@@ -2,4 +2,7 @@ transformers
 torch
 torchvision
 torchaudio
-spaces
+spaces
+accelerate
+safetensors
+diffusers