Update app.py

app.py

@@ -7,6 +7,7 @@ import torch
 from einops import rearrange
 from PIL import Image
 from transformers import pipeline
+from moviepy.editor import ImageClip, concatenate_videoclips, AudioFileClip
 
 from flux.cli import SamplingOptions
 from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
@@ -15,7 +16,7 @@ from pulid.pipeline_flux import PuLIDPipeline
 from pulid.utils import resize_numpy_image_long
 from prompt_template import prompt_dict
 from style_template import styles
-
+from images import image_dict
 NSFW_THRESHOLD = 0.85
 PROMPTS_NAMES = list(prompt_dict.keys())
 DEFAULT_PROMPT_NAME = "Paris"
@@ -29,11 +30,33 @@ def get_prompt (prompt_name : str):
     prompt = prompt_dict.get(prompt_name , prompt_dict[DEFAULT_PROMPT_NAME])
     return prompt[0]
 
+def get_image(prompt_name):
+    return image_dict[prompt_name]
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + ' ' + negative
 
+def generate_video(prompt_name , output_image):
+    images = get_image(prompt_name)
+    image_clips = [ImageClip(img_path).set_duration(2).resize((1024,1024)) 
+                   for img_path in images]
+    
+    video_clip =[]
+    try :
+        video_clip.append(output_image.crossfadein(0.5))
+        for i in range(len(image_clips)):
+            video_clip.append(image_clips[i].crossfadeout(0.5))
+        final_video = concatenate_videoclips(video_clip, method="compose")
+        final_video.write_videofile(filename=f"{prompt_name}.mp4", fps=24)
+        return final_video
+    except : 
+        for i in range(len(image_clips)):
+            video_clip.append(image_clips[i].crossfadeout(0.5))
+        final_video = concatenate_videoclips(video_clip, method="compose")
+        final_video.write_videofile(filename=f"{prompt_name}.mp4", fps=24)
+        return final_video
+
 
 def get_models(name: str, device: torch.device, offload: bool):
     t5 = load_t5(device, max_length=128)
@@ -200,7 +223,7 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                 id_image = gr.Image(label="ID Image")
                 
                 generate_btn = gr.Button("Generate")
-                
+                video_btn = gr.Button("Generate Video")
                 id_weight = gr.Slider(0.0, 3.0, 1, step=0.05, label="id weight")
 
                 width = gr.Slider(256, 1536, 896, step=16, label="Width")
@@ -225,6 +248,7 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                 output_image = gr.Image(label="Generated Image", format='png')
                 seed_output = gr.Textbox(label="Used Seed")
                 intermediate_output = gr.Gallery(label='Output', elem_id="gallery", visible=args.dev)
+                output_video = gr.Video(label= "Generated Video" , format='mp4')
 
         
         generate_btn.click(
@@ -233,7 +257,11 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                     neg_prompt, timestep_to_start_cfg, max_sequence_length],
             outputs=[output_image, seed_output, intermediate_output],
         )
-
+        video_btn.click(
+            fn=generate_video,
+            inputs=[prompt_name,output_image],
+            outputs=[output_video]
+        )
     return demo