Update app.py

app.py

@@ -7,6 +7,7 @@ from diffusers import (
     StableDiffusionImg2ImgPipeline,
     FluxPipeline,
     DiffusionPipeline,
+    DPMSolverMultistepScheduler,
 )
 from transformers import (
     pipeline as transformers_pipeline,
@@ -383,6 +384,48 @@ def retrain_models():
     pass
 
 
+def generate_text_to_video_ms_1_7b(prompt, num_frames=200):
+    blob_name = f"diffusers/text_to_video_ms_1_7b:{prompt}:{num_frames}"
+    video_bytes = load_object_from_gcs(blob_name)
+    if not video_bytes:
+        try:
+            with tqdm(total=1, desc="Generating video") as pbar:
+                video_frames = text_to_video_ms_1_7b_pipeline(
+                    prompt, num_inference_steps=25, num_frames=num_frames
+                ).frames
+                pbar.update(1)
+            video_path = export_to_video(video_frames)
+            with open(video_path, "rb") as f:
+                video_bytes = f.read()
+            save_object_to_gcs(blob_name, video_bytes)
+            os.remove(video_path)
+        except Exception as e:
+            print(f"Failed to generate video: {e}")
+            return None
+    return video_bytes
+
+
+def generate_text_to_video_ms_1_7b_short(prompt):
+    blob_name = f"diffusers/text_to_video_ms_1_7b_short:{prompt}"
+    video_bytes = load_object_from_gcs(blob_name)
+    if not video_bytes:
+        try:
+            with tqdm(total=1, desc="Generating short video") as pbar:
+                video_frames = text_to_video_ms_1_7b_short_pipeline(
+                    prompt, num_inference_steps=25
+                ).frames
+                pbar.update(1)
+            video_path = export_to_video(video_frames)
+            with open(video_path, "rb") as f:
+                video_bytes = f.read()
+            save_object_to_gcs(blob_name, video_bytes)
+            os.remove(video_path)
+        except Exception as e:
+            print(f"Failed to generate short video: {e}")
+            return None
+    return video_bytes
+
+
 text_to_image_pipeline = get_model_or_download(
     "stabilityai/stable-diffusion-2",
     "diffusers/text_to_image_model",
@@ -517,6 +560,23 @@ gemma_2_27b_it_pipeline = transformers_pipeline(
     model="google/gemma-2-27b-it",
     model_kwargs={"torch_dtype": torch.bfloat16},
 )
+text_to_video_ms_1_7b_pipeline = DiffusionPipeline.from_pretrained(
+    "damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16"
+)
+text_to_video_ms_1_7b_pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+    text_to_video_ms_1_7b_pipeline.scheduler.config
+)
+text_to_video_ms_1_7b_pipeline.enable_model_cpu_offload()
+text_to_video_ms_1_7b_pipeline.enable_vae_slicing()
+text_to_video_ms_1_7b_short_pipeline = DiffusionPipeline.from_pretrained(
+    "damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16"
+)
+text_to_video_ms_1_7b_short_pipeline.scheduler = (
+    DPMSolverMultistepScheduler.from_config(
+        text_to_video_ms_1_7b_short_pipeline.scheduler.config
+    )
+)
+text_to_video_ms_1_7b_short_pipeline.enable_model_cpu_offload()
 
 tools = []
 
@@ -686,11 +746,14 @@ gemma_2_2b_it_tab = gr.Interface(
     outputs=gr.Textbox(label="Gemma 2 2B IT Response:"),
     title="Gemma 2 2B IT",
 )
+
+
 def generate_gemma_2_27b(prompt):
     input_ids = gemma_2_27b_tokenizer(prompt, return_tensors="pt")
     outputs = gemma_2_27b_model.generate(**input_ids, max_new_tokens=32)
     return gemma_2_27b_tokenizer.decode(outputs[0])
 
+
 gemma_2_27b_tab = gr.Interface(
     fn=generate_gemma_2_27b,
     inputs=[gr.Textbox(label="Prompt:")],
@@ -703,6 +766,21 @@ gemma_2_27b_it_tab = gr.Interface(
     outputs=gr.Textbox(label="Gemma 2 27B IT Response:"),
     title="Gemma 2 27B IT",
 )
+text_to_video_ms_1_7b_tab = gr.Interface(
+    fn=generate_text_to_video_ms_1_7b,
+    inputs=[
+        gr.Textbox(label="Prompt:"),
+        gr.Slider(50, 200, 200, step=1, label="Number of Frames:"),
+    ],
+    outputs=gr.Video(),
+    title="Text to Video MS 1.7B",
+)
+text_to_video_ms_1_7b_short_tab = gr.Interface(
+    fn=generate_text_to_video_ms_1_7b_short,
+    inputs=[gr.Textbox(label="Prompt:")],
+    outputs=gr.Video(),
+    title="Text to Video MS 1.7B Short",
+)
 
 app = gr.TabbedInterface(
     [
@@ -734,6 +812,8 @@ app = gr.TabbedInterface(
         gemma_2_2b_it_tab,
         gemma_2_27b_tab,
         gemma_2_27b_it_tab,
+        text_to_video_ms_1_7b_tab,
+        text_to_video_ms_1_7b_short_tab,
     ],
     [
         "Generate Image",
@@ -764,6 +844,8 @@ app = gr.TabbedInterface(
         "Gemma 2 2B IT",
         "Gemma 2 27B",
         "Gemma 2 27B IT",
+        "Text to Video MS 1.7B",
+        "Text to Video MS 1.7B Short",
     ],
 )