Update app.py

app.py

@@ -69,7 +69,7 @@ def get_modelscope_pipeline(
     lora.merge_and_unload()
     pipe.unet = lora
 
-    pipe = pipe.to(device)
+    # pipe = pipe.to(device)
 
     return pipe
 
@@ -136,22 +136,48 @@ def get_animatediff_pipeline(
     lora.merge_and_unload()
     pipe.unet = lora
 
-    pipe = pipe.to(device)
+    # pipe = pipe.to(device)
     return pipe
 
 
-# pipe_dict = {
-#     "ModelScope T2V": {"WebVid": None, "LAION-aes": None, "Anime": None, "Realistic": None, "3D Cartoon": None},
-#     "AnimateDiff (SD1.5)": {"WebVid": None, "LAION-aes": None},
-#     "AnimateDiff (RealisticVision)": {"WebVid": None, "LAION-aes": None},
-#     "AnimateDiff (epiCRealism)": {"WebVid": None, "LAION-aes": None},
-# }
-cache_pipeline = {
-    "base_model": None,
-    "variant": None,
-    "pipeline": None,
+pipe_dict = {
+    "ModelScope T2V": {"WebVid": None, "LAION-aes": None, "Anime": None, "Realistic": None, "3D Cartoon": None},
+    "AnimateDiff (SD1.5)": {"WebVid": None, "LAION-aes": None},
+    "AnimateDiff (RealisticVision)": {"WebVid": None, "LAION-aes": None},
+    "AnimateDiff (epiCRealism)": {"WebVid": None, "LAION-aes": None},
 }
+# cache_pipeline = {
+#     "base_model": None,
+#     "variant": None,
+#     "pipeline": None,
+# }
 
+def init_pipelines():
+    for base_model in variants.keys():
+        for variant in variants[base_model]:
+            if pipe_dict[base_model][variant] is None:
+                if base_model == "ModelScope T2V":
+                    pipe_dict[base_model][variant] = get_modelscope_pipeline(mcm_variant=variant)
+                elif base_model == "AnimateDiff (SD1.5)":
+                    pipe_dict[base_model][variant] = get_animatediff_pipeline(
+                        real_variant=None,
+                        motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+                        mcm_variant=variant,
+                    )
+                elif base_model == "AnimateDiff (RealisticVision)":
+                    pipe_dict[base_model][variant] = get_animatediff_pipeline(
+                        real_variant="realvision",
+                        motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+                        mcm_variant=variant,
+                    )
+                elif base_model == "AnimateDiff (epiCRealism)":
+                    pipe_dict[base_model][variant] = get_animatediff_pipeline(
+                        real_variant="epicrealism",
+                        motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+                        mcm_variant=variant,
+                    )
+                else:
+                    raise ValueError(f"Unknown base_model {base_model}")
 
 @spaces.GPU(duration=120)
 def infer(
@@ -180,45 +206,46 @@ def infer(
     #         )
     #     else:
     #         raise ValueError(f"Unknown base_model {base_model}")
-    if (
-        cache_pipeline["base_model"] == base_model
-        and cache_pipeline["variant"] == variant
-    ):
-        pass
-    else:
-        if base_model == "ModelScope T2V":
-            pipeline = get_modelscope_pipeline(mcm_variant=variant)
-        elif base_model == "AnimateDiff (SD1.5)":
-            pipeline = get_animatediff_pipeline(
-                real_variant=None,
-                motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
-                mcm_variant=variant,
-            )
-        elif base_model == "AnimateDiff (RealisticVision)":
-            pipeline = get_animatediff_pipeline(
-                real_variant="realvision",
-                motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
-                mcm_variant=variant,
-            )
-        elif base_model == "AnimateDiff (epiCRealism)":
-            pipeline = get_animatediff_pipeline(
-                real_variant="epicrealism",
-                motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
-                mcm_variant=variant,
-            )
-        else:
-            raise ValueError(f"Unknown base_model {base_model}")
+    # if (
+    #     cache_pipeline["base_model"] == base_model
+    #     and cache_pipeline["variant"] == variant
+    # ):
+    #     pass
+    # else:
+    #     if base_model == "ModelScope T2V":
+    #         pipeline = get_modelscope_pipeline(mcm_variant=variant)
+    #     elif base_model == "AnimateDiff (SD1.5)":
+    #         pipeline = get_animatediff_pipeline(
+    #             real_variant=None,
+    #             motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+    #             mcm_variant=variant,
+    #         )
+    #     elif base_model == "AnimateDiff (RealisticVision)":
+    #         pipeline = get_animatediff_pipeline(
+    #             real_variant="realvision",
+    #             motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+    #             mcm_variant=variant,
+    #         )
+    #     elif base_model == "AnimateDiff (epiCRealism)":
+    #         pipeline = get_animatediff_pipeline(
+    #             real_variant="epicrealism",
+    #             motion_module_path="guoyww/animatediff-motion-adapter-v1-5-2",
+    #             mcm_variant=variant,
+    #         )
+    #     else:
+    #         raise ValueError(f"Unknown base_model {base_model}")
 
-        cache_pipeline["base_model"] = base_model
-        cache_pipeline["variant"] = variant
-        cache_pipeline["pipeline"] = pipeline
+    #     cache_pipeline["base_model"] = base_model
+    #     cache_pipeline["variant"] = variant
+    #     cache_pipeline["pipeline"] = pipeline
 
+    pipe_dict[base_model][variant] = pipe_dict[base_model][variant].to(device)
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator("cpu").manual_seed(seed)
 
-    output = cache_pipeline["pipeline"](
+    output = pipe_dict[base_model][variant](
         prompt=prompt,
         num_frames=16,
         guidance_scale=1.0,
@@ -238,7 +265,8 @@ def infer(
         fps=7,
     )
     print(f"Saved to {save_path}")
-    return save_path
+    pipe_dict[base_model][variant] = pipe_dict[base_model][variant].to("cpu")
+    return save_path, seed
 
 
 examples = [
@@ -402,7 +430,7 @@ with gr.Blocks(css=css) as demo:
             inputs=[base_model, variant_dropdown, prompt, num_inference_steps],
             cache_examples=True,
             fn=infer,
-            outputs=[result],
+            outputs=[result, seed],
         )
 
     run_button.click(
@@ -415,7 +443,7 @@ with gr.Blocks(css=css) as demo:
             seed,
             randomize_seed,
         ],
-        outputs=[result],
+        outputs=[result, seed],
     )
 
 demo.queue().launch()