add pipe

app.py

@@ -7,7 +7,8 @@ from glob import glob
 from pathlib import Path
 from typing import Optional
 
-from diffusers import StableVideoDiffusionPipeline
+# from diffusers import StableVideoDiffusionPipeline
+from kandinsky import get_T2V_pipeline
 from diffusers.utils import load_image, export_to_video
 from PIL import Image
 
@@ -17,13 +18,20 @@ from huggingface_hub import hf_hub_download
 
 #gradio.helpers.CACHED_FOLDER = '/data/cache'
 
-pipe = StableVideoDiffusionPipeline.from_pretrained(
-    "multimodalart/stable-video-diffusion", torch_dtype=torch.float16, variant="fp16"
-)
-pipe.to("cuda")
+# pipe = StableVideoDiffusionPipeline.from_pretrained(
+#     "multimodalart/stable-video-diffusion", torch_dtype=torch.float16, variant="fp16"
+# )
+# pipe.to("cuda")
 #pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
 #pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
 
+device_map = {
+    "dit": torch.device('cuda'), 
+    "vae": torch.device('cuda'), 
+    "text_embedder": torch.device('cuda')
+}
+pipe = get_T2V_pipeline(device_map)
+
 max_64_bit_int = 2**63 - 1
 
 @spaces.GPU(duration=120)
@@ -40,21 +48,29 @@ def sample(
     output_folder: str = "outputs",
     progress=gr.Progress(track_tqdm=True)
 ):
-    if image.mode == "RGBA":
-        image = image.convert("RGB")
+    # if image.mode == "RGBA":
+    #     image = image.convert("RGB")
         
-    if(randomize_seed):
-        seed = random.randint(0, max_64_bit_int)
-    generator = torch.manual_seed(seed)
-    
+    # if(randomize_seed):
+    #     seed = random.randint(0, max_64_bit_int)
+    # generator = torch.manual_seed(seed)
+
     os.makedirs(output_folder, exist_ok=True)
     base_count = len(glob(os.path.join(output_folder, "*.mp4")))
     video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
 
-    frames = pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=0.1, num_frames=25).frames[0]
-    export_to_video(frames, video_path, fps=fps_id)
+    # frames = pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=0.1, num_frames=25).frames[0]
+    frames = pipe(
+        seed=seed,
+        time_length=12,
+        width = 672,
+        height = 384,
+        save_path=video_path,
+        text=prompt,
+    )
+    export_to_video(frames, video_path, fps=8)
     torch.manual_seed(seed)
-    
+
     return video_path, seed
 
 def resize_image(image, output_size=(1024, 576)):