Update app.py

app.py

@@ -14,6 +14,7 @@ import torch
 
 from spectro import wav_bytes_from_spectrogram_image
 from diffusers import StableDiffusionPipeline
+from diffusers import EulerAncestralDiscreteScheduler
  
 import io
 from os import path
@@ -38,8 +39,10 @@ tips = {"en": "Tips: The input text will be translated into English for generati
 
 count = 0
 
+
 model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+eulera = EulerAncestralDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, scheduler=eulera)
 pipe = pipe.to("cuda")
 
 model_id2 = "riffusion/riffusion-model-v1"
@@ -60,23 +63,23 @@ def translate_language(text_prompts):
             text_prompts = language_translation_model.translate(text_prompts, language_code, 'en')
     except Exception as e:
         error_text = str(e)
-        return {status_text:error_text, language_tips_text:gr.update(visible=False)}
+        return {status_text:error_text, language_tips_text:gr.update(visible=False), translated_language:text_prompts, trigger_component: gr.update(value=count, visible=False)}
     if language_code in tips:
         tips_text = tips[language_code]
     else:
         tips_text = tips['en']
-    if language_code == 'zh':
+    if language_code == 'en':
         return {language_tips_text:gr.update(visible=False), translated_language:text_prompts, trigger_component: gr.update(value=count, visible=False)}
     else:
         return {language_tips_text:gr.update(visible=True, value=tips_text), translated_language:text_prompts, trigger_component:  gr.update(value=count, visible=False)}
 
 
 
-def get_result(text_prompts, style_indx, musicAI_indx):
+def get_result(text_prompts, style_indx, musicAI_indx, duration):
     style = style_list_EN[style_indx]
     prompt = style + "," + text_prompts
         
-    sdresult = pipe(prompt)
+    sdresult = pipe(prompt, negative_prompt = "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", num_inference_steps=50, guidance_scale=7, width=576, height=576)
     image_output = sdresult.images[0] if not sdresult.nsfw_content_detected[0] else Image.open("nsfw_placeholder.jpg")
     
     print("Generated image with prompt " + prompt)
@@ -91,15 +94,18 @@ def get_result(text_prompts, style_indx, musicAI_indx):
     
     interrogate_prompt = img_to_text(imagefile, "ViT-L (best for Stable Diffusion 1.*)", "fast", fn_index=1)[0]
     print(interrogate_prompt)
-    spec_image, music_output = get_music(interrogate_prompt + ", " + style_list_EN[style_indx], musicAI_indx)
+    spec_image, music_output = get_music(interrogate_prompt + ", " + style_list_EN[style_indx], musicAI_indx, duration)
     
     video_merged = merge_video(music_output, image_output)
-    return {spec_result:spec_image, video_result:video_merged, status_text:'Success'}
-
+    return {spec_result:spec_image, video_result:video_merged, status_text:'Success', share_button:gr.update(visible=True), community_icon:gr.update(visible=True), loading_icon:gr.update(visible=True)}
 
-def get_music(prompt, musicAI_indx):
+def get_music(prompt, musicAI_indx, duration):
     if musicAI_indx == 0: 
-        spec = pipe2(prompt).images[0]
+        if duration == 5:
+            width_duration=512
+        else :
+            width_duration = 512 + ((int(duration)-5) * 128)
+        spec = pipe2(prompt, height=512, width=width_duration).images[0]
         print(spec)
         wav = wav_bytes_from_spectrogram_image(spec)
         with open("output.wav", "wb") as f:
@@ -148,7 +154,9 @@ def merge_video(mp3file_name, image):
     fps = 12
     slide_time = audio_length
     fourcc = cv2.VideoWriter.fourcc(*'MJPG')
-    out = cv2.VideoWriter(file_name, fourcc, fps, (512, 512))
+
+    #W, H should be the same as input image
+    out = cv2.VideoWriter(file_name, fourcc, fps, (576, 576))
     
     # for image in img_list:
     #     cv_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
@@ -188,6 +196,11 @@ def merge_video(mp3file_name, image):
     mergedclip.to_videofile('mergedvideo.mp4')
     return 'mergedvideo.mp4'
 
+def change_music_generator(current_choice):
+    if current_choice == 0:
+        return gr.update(visible=True)
+    return gr.update(visible=False)
+
 title="文生图生音乐视频 Text to Image to Music to Video with Riffusion"
 
 description="An AI art generation pipeline, which supports text-to-image-to-music task."
@@ -263,6 +276,22 @@ css = """
             font-weight: bold;
             font-size: 115%;
         }
+        #share-btn-container {
+            display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
+        }
+        #share-btn {
+            all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;
+        }
+        #share-btn * {
+            all: unset;
+        }
+        #share-btn-container div:nth-child(-n+2){
+            width: auto !important;
+            min-height: 0px !important;
+        }
+        #share-btn-container .wrap {
+            display: none !important;
+        }
 """
 
 block = gr.Blocks(css=css)
@@ -429,6 +458,7 @@ with block:
             </div>
         """
     )
+    
     with gr.Group():
         with gr.Box():
             with gr.Row().style(mobile_collapse=False, equal_height=True):
@@ -437,6 +467,7 @@ with block:
                     show_label=False,
                     max_lines=1,
                     placeholder="Enter your prompt, multiple languages are supported now.",
+                    elem_id="input-prompt",
                 ).style(
                     border=(True, False, True, True),
                     rounded=(True, False, False, True),
@@ -453,6 +484,7 @@ with block:
         '像素风格(Pixel Style)', '概念艺术(Conceptual Art)', '未来主义(Futurism)', '赛博朋克(Cyberpunk)', '写实风格(Realistic style)', 
         '洛丽塔风格(Lolita style)', '巴洛克风格(Baroque style)', '超现实主义(Surrealism)', '默认(Default)'], value='默认(Default)', type="index")
         musicAI = gr.Dropdown(label="音乐生成技术(AI Music Generator)", choices=['Riffusion', 'Mubert AI'], value='Riffusion', type="index")
+        duration_input = gr.Slider(label="Duration in seconds", minimum=5, maximum=10, step=1, value=5, elem_id="duration-slider", visible=True)
         status_text = gr.Textbox(
             label="处理状态(Process status)",
             show_label=True,
@@ -460,35 +492,45 @@ with block:
             interactive=False
         )
         
-        video_result = gr.Video(type=None, label='Final Merged video')
-        spec_result = gr.Image()
-        
-        trigger_component = gr.Textbox(vaule="", visible=False) # This component is used for triggering inference funtion.
-        translated_language = gr.Textbox(vaule="", visible=False)
-        
-        
-        ex = gr.Examples(examples=examples, fn=translate_language_example, inputs=[text, styles], outputs=[language_tips_text, status_text, trigger_component, translated_language], cache_examples=False)
-        ex.dataset.headers = [""]
 
-        
-        text.submit(translate_language, inputs=[text], outputs=[language_tips_text, status_text, trigger_component, translated_language])
-        btn.click(translate_language, inputs=[text], outputs=[language_tips_text, status_text, trigger_component, translated_language])
-        trigger_component.change(fn=get_result, inputs=[translated_language, styles, musicAI], outputs=[spec_result, video_result, status_text])
-        
-        
-        gr.Markdown(
-            """  
-      Space by [@DGSpitzer](https://www.youtube.com/channel/UCzzsYBF4qwtMwJaPJZ5SuPg)❤️ [@大谷的游戏创作小屋](https://space.bilibili.com/176003)
-      [![Twitter Follow](https://img.shields.io/twitter/follow/DGSpitzer?label=%40DGSpitzer&style=social)](https://twitter.com/DGSpitzer)
-      ![visitors](https://visitor-badge.glitch.me/badge?page_id=dgspitzer_txt2img2video)
-            """
-        )
-        gr.HTML('''
-        <div class="footer">
-                    <p>Model：<a href="https://huggingface.co/riffusion/riffusion-model-v1" style="text-decoration: underline;" target="_blank">Riffusion</a>
-                    </p>
-        </div>
-        ''')
+    with gr.Column(elem_id="col-container"):
+        with gr.Group(elem_id="share-btn-container"):
+            community_icon = gr.HTML(community_icon_html, visible=False)
+            loading_icon = gr.HTML(loading_icon_html, visible=False)
+            share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
+
+    share_button.click(None, [], [], _js=share_js)
+
+    video_result = gr.Video(type=None, label='Final Merged video', elem_id="output-video")
+    spec_result = gr.Image()
+    
+    trigger_component = gr.Textbox(vaule="", visible=False) # This component is used for triggering inference funtion.
+    translated_language = gr.Textbox(vaule="", visible=False)
+    
+    
+    ex = gr.Examples(examples=examples, fn=translate_language_example, inputs=[text, styles], outputs=[language_tips_text, status_text, trigger_component, translated_language], cache_examples=False)
+    ex.dataset.headers = [""]
+
+    
+    musicAI.change(fn=change_music_generator, inputs=[musicAI], outputs=[duration_input])
+    text.submit(translate_language, inputs=[text], outputs=[language_tips_text, status_text, trigger_component, translated_language])
+    btn.click(translate_language, inputs=[text], outputs=[language_tips_text, status_text, trigger_component, translated_language])
+    trigger_component.change(fn=get_result, inputs=[translated_language, styles, musicAI, duration_input], outputs=[spec_result, video_result, status_text, share_button, community_icon, loading_icon])
+
 
+    gr.Markdown(
+        """  
+  Space by [@DGSpitzer](https://www.youtube.com/channel/UCzzsYBF4qwtMwJaPJZ5SuPg)❤️ [@大谷的游戏创作小屋](https://space.bilibili.com/176003)
+  [![Twitter Follow](https://img.shields.io/twitter/follow/DGSpitzer?label=%40DGSpitzer&style=social)](https://twitter.com/DGSpitzer)
+  ![visitors](https://visitor-badge.glitch.me/badge?page_id=dgspitzer_txt2img2video)
+        """
+    )
+    gr.HTML('''
+    <div class="footer">
+                <p>Model：<a href="https://huggingface.co/riffusion/riffusion-model-v1" style="text-decoration: underline;" target="_blank">Riffusion</a>
+                </p>
+    </div>
+    ''')
+    
 
-block.queue(concurrency_count=128).launch()
+block.queue().launch()