ByteDance/AnimateDiff-Lightning text2vid support

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+max-line-length = 88
+extend-ignore = E203,E701

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,73 @@
+name: Bug report
+description: Create a report to help us improve.
+labels:
+  - "bug"
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## Bug report
+        Please fill out the following information to help us understand your issue.
+
+        > [!IMPORTANT]
+        > This repository is only related to the AI pipelines and AI models used on the AI subnet. For core bugs with the AI branch of the go-livepeer subnet software, please report them at the [go-livepeer issues page](https://github.com/livepeer/go-livepeer/issues).
+  - type: textarea
+    attributes:
+      label: Describe the bug
+      description: A clear and concise description of what the bug is.
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Reproduction steps
+      description: "How do you trigger this bug? Please walk us through it step by step."
+      value: |
+        1. Go to '...'
+        2. Click on '....'
+        3. Scroll down to '....'
+        4. See error
+  - type: textarea
+    attributes:
+      label: Expected behaviour
+      description:
+        A clear and concise description of what you expected to happen.
+  - type: dropdown
+    id: severity
+    attributes:
+      label: Severity
+      description: "How severe is this bug?"
+      options:
+        - Minor
+        - Major
+        - Critical
+  - type: textarea
+    attributes:
+      label: Screenshots / Live demo link
+      description: If applicable, add screenshots to help explain your problem.
+      placeholder: Paste the github-readme-stats link as markdown image
+  - type: dropdown
+    id: os
+    attributes:
+      label: OS
+      description: "What operating system are you using?"
+      options:
+        - Windows
+        - Mac
+        - Linux
+  - type: dropdown
+    id: running_on
+    attributes:
+      label: Running on
+      description: "Where are you running the application?"
+      options:
+        - Local
+        - Docker
+  - type: input
+    id: ai_worker_version
+    attributes:
+      label: AI-worker version
+      description: "What version of the AI-worker are you using?"
+  - type: textarea
+    attributes:
+      label: Additional context
+      description: Add any other context about the problem here.

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: true
+contact_links:
+  - name: Question
+    url: https://discord.gg/livepeer
+    about: "Have a question? Join us in the 'ai-video' channel on the Livepeer Discord server. We're here to help!"

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,44 @@
+name: Feature request
+description: Suggest an idea for this project.
+labels:
+  - "enhancement"
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## Feature Request
+        Please fill out the following information to help us understand your request.
+
+        > [!IMPORTANT]
+        > This repository is only related to the AI pipelines and AI models used on the AI subnet. For core features related to the AI branch of the go-livepeer subnet software, please request them at the [go-livepeer issues page](https://github.com/livepeer/go-livepeer/issues).
+  - type: textarea
+    attributes:
+      label: Is your feature request related to a problem? Please describe.
+      description:
+        A clear and concise description of what the problem is. Ex. I'm always
+        frustrated when [...]
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Describe the solution you'd like
+      description: A clear and concise description of what you want to happen.
+  - type: textarea
+    attributes:
+      label: Describe alternatives you've considered
+      description:
+        A clear and concise description of any alternative solutions or features
+        you've considered.
+  - type: textarea
+    attributes:
+      label: Use Case
+      description: "Please describe why you want this feature to be added. This will help us prioritize your request."
+  - type: textarea
+    attributes:
+      label: Expected Outcome
+      description: "What do you expect to happen once this feature is implemented?"
+  - type: textarea
+    attributes:
+      label: Additional context
+      description:
+        Add any other context or screenshots about the feature request here.

.github/ISSUE_TEMPLATE/model_request.yml

@@ -0,0 +1,42 @@
+name: AI model request
+description: Request a model to be added to a given AI pipeline.
+labels:
+  - "enhancement"
+  - "model"
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## AI Model Request
+        Please fill out the following information to help us understand your needs.
+  - type: textarea
+    attributes:
+      label: Which model do you want us to add support for?
+      description: "I would like to request support for the following model: [model name]"
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Hugging Face model hub link
+      description: "Please provide a link to the model on the Hugging Face model hub"
+    validations:
+      required: true
+  - type: dropdown
+    id: pipeline
+    attributes:
+      label: Pipeline
+      description: To which pipeline should we add this model?
+      options:
+        - text-to-image
+        - image-to-image
+        - image-to-video
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Use Case
+      description: "Please describe why you want this model to be added. This will help us prioritize your request"
+  - type: textarea
+    attributes:
+      label: Additional context
+      description: Add any other context or screenshots about the feature request here.

.github/ISSUE_TEMPLATE/pipeline_request.yml

@@ -0,0 +1,37 @@
+name: AI Pipeline request
+description: Request a new AI pipeline to be added to the project.
+labels:
+  - "enhancement"
+  - "pipeline"
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## AI Pipeline Request
+        Please fill out the following information to help us understand your needs.
+  - type: textarea
+    attributes:
+      label: What problem would this pipeline solve?
+      description: "Please describe the problem that would be solved by adding this pipeline. Ex. I'm unable to process [type of data] because [...]"
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Describe the pipeline you'd like
+      description: "Please provide a detailed description of the pipeline you'd like to be added."
+  - type: textarea
+    attributes:
+      label: Describe alternatives you've considered
+      description: "Please describe any alternative solutions or pipelines you've considered."
+  - type: input
+    attributes:
+      label: Link to a supported model
+      description: "If applicable, please provide a link to a model that is supported by this pipeline."
+  - type: textarea
+    attributes:
+      label: Use Case
+      description: "Please describe why you want this pipeline to be added. This will help us prioritize your request."
+  - type: textarea
+    attributes:
+      label: Additional context
+      description: "Add any other context or screenshots about the feature request here."

runner/.devcontainer/devcontainer.json

@@ -32,8 +32,8 @@
 
 	// Use 'mounts' to make a list of local folders available inside the container.
 	"mounts": [
-		"source=${localWorkspaceFolder}/models,target=/models,type=bind"
-		// "source=${localEnv:HOME}/.lpData/models,target=/models,type=bind"
+		// "source=${localWorkspaceFolder}/models,target=/models,type=bind"
+		"source=${localEnv:HOME}/.lpData/models,target=/models,type=bind"
     ]
 
 	// Use 'postCreateCommand' to run commands after the container is created.

runner/.flake8

@@ -0,0 +1,3 @@
+[flake8]
+max-line-length = 88
+extend-ignore = E203,E701

runner/app/main.py

@@ -42,6 +42,10 @@ def load_pipeline(pipeline: str, model_id: str) -> any:
             from app.pipelines.image_to_video import ImageToVideoPipeline
 
             return ImageToVideoPipeline(model_id)
+        case "text-to-video":
+            from app.pipelines.text_to_video import TextToVideoPipeline
+
+            return TextToVideoPipeline(model_id)
         case "frame-interpolation":
             raise NotImplementedError("frame-interpolation pipeline not implemented")
         case "upscale":
@@ -66,6 +70,10 @@ def load_route(pipeline: str) -> any:
             from app.routes import image_to_video
 
             return image_to_video.router
+        case "text-to-video":
+            from app.routes import text_to_video
+
+            return text_to_video.router
         case "frame-interpolation":
             raise NotImplementedError("frame-interpolation pipeline not implemented")
         case "upscale":

runner/app/pipelines/image_to_image.py

@@ -92,15 +92,25 @@ def __init__(self, model_id: str):
                 model_id, **kwargs
             ).to(torch_device)
 
-        if os.environ.get("SFAST"):
+        if os.getenv("SFAST", "").strip().lower() == "true":
             logger.info(
-                "ImageToImagePipeline will be dynamicallly compiled with stable-fast for %s",
+                "ImageToImagePipeline will be dynamically compiled with stable-fast "
+                "for %s",
                 model_id,
             )
             from app.pipelines.sfast import compile_model
 
             self.ldm = compile_model(self.ldm)
 
+            # Warm-up the pipeline.
+            # TODO: Not yet supported for ImageToImagePipeline.
+            if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+                logger.warning(
+                    "The 'SFAST_WARMUP' flag is not yet supported for the "
+                    "ImageToImagePipeline and will be ignored. As a result the first "
+                    "call may be slow if 'SFAST' is enabled."
+                )
+
     def __call__(self, prompt: str, image: PIL.Image, **kwargs) -> List[PIL.Image]:
         seed = kwargs.pop("seed", None)
         if seed is not None:

runner/app/pipelines/image_to_video.py

@@ -8,13 +8,16 @@
 from typing import List
 import logging
 import os
+import time
 
 from PIL import ImageFile
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 logger = logging.getLogger(__name__)
 
+SFAST_WARMUP_ITERATIONS = 2  # Model warm-up iterations when SFAST is enabled.
+
 
 class ImageToVideoPipeline(Pipeline):
     def __init__(self, model_id: str):
@@ -40,17 +43,52 @@ def __init__(self, model_id: str):
         self.ldm = StableVideoDiffusionPipeline.from_pretrained(model_id, **kwargs)
         self.ldm.to(get_torch_device())
 
-        if os.environ.get("SFAST"):
+        if os.getenv("SFAST", "").strip().lower() == "true":
             logger.info(
-                "ImageToVideoPipeline will be dynamicallly compiled with stable-fast for %s",
+                "ImageToVideoPipeline will be dynamically compiled with stable-fast "
+                "for %s",
                 model_id,
             )
             from app.pipelines.sfast import compile_model
 
             self.ldm = compile_model(self.ldm)
 
+            # Warm-up the pipeline.
+            # NOTE: Initial calls may be slow due to compilation. Subsequent calls will
+            # be faster.
+            if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+                # Retrieve default model params.
+                # TODO: Retrieve defaults from Pydantic class in route.
+                warmup_kwargs = {
+                    "image": PIL.Image.new("RGB", (576, 1024)),
+                    "height": 576,
+                    "width": 1024,
+                    "fps": 6,
+                    "motion_bucket_id": 127,
+                    "noise_aug_strength": 0.02,
+                    "decode_chunk_size": 25,
+                }
+
+                logger.info("Warming up ImageToVideoPipeline pipeline...")
+                total_time = 0
+                for ii in range(SFAST_WARMUP_ITERATIONS):
+                    t = time.time()
+                    try:
+                        self.ldm(**warmup_kwargs).frames
+                    except Exception as e:
+                        # FIXME: When out of memory, pipeline is corrupted.
+                        logger.error(f"ImageToVideoPipeline warmup error: {e}")
+                        raise e
+                    iteration_time = time.time() - t
+                    total_time += iteration_time
+                    logger.info(
+                        "Warmup iteration %s took %s seconds", ii + 1, iteration_time
+                    )
+                logger.info("Total warmup time: %s seconds", total_time)
+
     def __call__(self, image: PIL.Image, **kwargs) -> List[List[PIL.Image]]:
         if "decode_chunk_size" not in kwargs:
+            # Decrease decode_chunk_size to reduce memory usage.
             kwargs["decode_chunk_size"] = 4
 
         seed = kwargs.pop("seed", None)

runner/app/pipelines/text_to_image.py

@@ -111,15 +111,25 @@ def __init__(self, model_id: str):
                 self.ldm.vae.decode, mode="max-autotune", fullgraph=True
             )
 
-        if os.environ.get("SFAST"):
+        if os.getenv("SFAST", "").strip().lower() == "true":
             logger.info(
-                "TextToImagePipeline will be dynamicallly compiled with stable-fast for %s",
+                "TextToImagePipeline will be dynamically compiled with stable-fast for "
+                "%s",
                 model_id,
             )
             from app.pipelines.sfast import compile_model
 
             self.ldm = compile_model(self.ldm)
 
+            # Warm-up the pipeline.
+            # TODO: Not yet supported for ImageToImagePipeline.
+            if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+                logger.warning(
+                    "The 'SFAST_WARMUP' flag is not yet supported for the "
+                    "TextToImagePipeline and will be ignored. As a result the first "
+                    "call may be slow if 'SFAST' is enabled."
+                )
+
     def __call__(self, prompt: str, **kwargs) -> List[PIL.Image]:
         seed = kwargs.pop("seed", None)
         if seed is not None: