Update app.py

app.py

@@ -21,6 +21,9 @@ from funcs import (
     save_videos
 )
 from transformers import pipeline
+from diffusers import FluxPipeline
+from PIL import Image
+import numpy as np
 
 def download_model():
     REPO_ID = 'Doubiiu/DynamiCrafter_1024'
@@ -47,6 +50,23 @@ model = model.cuda()
 # 번역 모델 초기화
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
+# FLUX 파이프라인 초기화
+flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+flux_pipe.enable_model_cpu_offload()
+
+def generate_image_from_text(prompt, seed=0):
+    generator = torch.Generator("cpu").manual_seed(seed)
+    image = flux_pipe(
+        prompt,
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        num_inference_steps=50,
+        max_sequence_length=512,
+        generator=generator
+    ).images[0]
+    return image
+
 @spaces.GPU(duration=300)
 def infer(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, video_length=2):
     # 한글 입력 감지 및 번역
@@ -105,6 +125,17 @@ def infer(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, video
         save_videos(batch_samples, './', filenames=['output'], fps=save_fps)
     return video_path
 
+@spaces.GPU(duration=300)
+def infer_t2v(prompt, video_prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, video_length=2):
+    # 이미지 생성
+    image = generate_image_from_text(prompt, seed)
+    
+    # 이미지를 numpy 배열로 변환
+    image_np = np.array(image)
+    
+    # 비디오 생성을 위해 기존 infer 함수 호출
+    return infer(image_np, video_prompt, steps, cfg_scale, eta, fs, seed, video_length)
+
 i2v_examples = [
     ['prompts/1024/astronaut04.png', 'a man in an astronaut suit playing a guitar', 30, 7.5, 1.0, 6, 123, 2],
 ]
@@ -145,4 +176,31 @@ with gr.Blocks(analytics_enabled=False, css=css) as dynamicrafter_iface:
                         fn = infer
         )
 
+    with gr.Tab(label='T2V'):
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        t2v_input_text = gr.Text(label='Image Generation Prompt')
+                    with gr.Row():
+                        t2v_video_prompt = gr.Text(label='Video Generation Prompt')
+                    with gr.Row():
+                        t2v_seed = gr.Slider(label='Random Seed', minimum=0, maximum=10000, step=1, value=123)
+                        t2v_eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label='ETA', value=1.0)
+                        t2v_cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.5)
+                    with gr.Row():
+                        t2v_steps = gr.Slider(minimum=1, maximum=50, step=1, label="Sampling steps", value=30)
+                        t2v_motion = gr.Slider(minimum=5, maximum=20, step=1, label="FPS", value=8)
+                    with gr.Row():
+                        t2v_video_length = gr.Slider(minimum=2, maximum=8, step=1, label="Video Length (seconds)", value=2)
+                    t2v_end_btn = gr.Button("Generate")
+                with gr.Row():
+                    t2v_output_video = gr.Video(label="Generated Video", autoplay=True, show_share_button=True)
+
+        t2v_end_btn.click(
+            inputs=[t2v_input_text, t2v_video_prompt, t2v_steps, t2v_cfg_scale, t2v_eta, t2v_motion, t2v_seed, t2v_video_length],
+            outputs=[t2v_output_video],
+            fn=infer_t2v
+        )
+
 dynamicrafter_iface.queue(max_size=12).launch(show_api=True)