| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import torch |
| from diffusers import DiffusionPipeline |
|
|
| from .base import ImageToVideoGenerator |
|
|
|
|
| class DiffusersImageToVideoAdapter(ImageToVideoGenerator): |
| def __init__(self, model_id: str, device: str = "cuda", torch_dtype: Any | None = None, pipeline: Any | None = None, **kwargs) -> None: |
| if pipeline is None: |
| dtype = torch_dtype |
| if dtype is None and str(device).startswith("cuda"): |
| dtype = torch.bfloat16 |
| pipeline_cls = kwargs.pop("pipeline_cls", None) or _default_i2v_pipeline_cls(model_id) |
| pipeline = pipeline_cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs) |
| self.pipe = pipeline |
| self.device = device |
| if hasattr(self.pipe, "to"): |
| self.pipe.to(device) |
| if hasattr(self.pipe, "set_progress_bar_config"): |
| self.pipe.set_progress_bar_config(disable=True) |
|
|
| def generate(self, image: Any, prompt: str, *, generator: Any | None = None, **kwargs) -> Any: |
| kwargs.setdefault("output_type", "pt") |
| try: |
| out = self.pipe(image=image, prompt=prompt, generator=generator, **kwargs) |
| except TypeError as exc: |
| raise TypeError( |
| "The configured image-to-video model does not support the default " |
| "`image=..., prompt=..., generator=..., **kwargs` signature. " |
| "Pass a custom ImageToVideoGenerator adapter." |
| ) from exc |
| return self._normalize_output(out) |
|
|
| def generate_batch(self, images, prompts, *, generators=None, **kwargs): |
| images = list(images) |
| prompts = list(prompts) |
| if not prompts: |
| return [] |
| kwargs.setdefault("output_type", "pt") |
| try: |
| out = self.pipe(image=images, prompt=prompts, generator=generators, **kwargs) |
| frames = getattr(out, "frames", None) |
| if frames is None: |
| frames = getattr(out, "videos", None) |
| if torch.is_tensor(frames) and frames.ndim == 5 and frames.shape[0] == len(prompts): |
| return [frames[i] for i in range(len(prompts))] |
| if isinstance(frames, (list, tuple)) and len(frames) == len(prompts): |
| return list(frames) |
| except (RuntimeError, TypeError, ValueError): |
| pass |
| |
| gens = generators if isinstance(generators, (list, tuple)) else [generators] * len(prompts) |
| return [self.generate(img, p, generator=g, **kwargs) for img, p, g in zip(images, prompts, gens)] |
|
|
| @staticmethod |
| def _normalize_output(out: Any) -> Any: |
| frames = getattr(out, "frames", None) |
| if frames is None: |
| frames = getattr(out, "videos", None) |
| if frames is None: |
| return out |
| if torch.is_tensor(frames): |
| return frames[0] if frames.ndim == 5 else frames |
| if isinstance(frames, (list, tuple)) and frames: |
| return frames[0] |
| return frames |
|
|
|
|
| def _default_i2v_pipeline_cls(model_id: str): |
| if model_id == "Lightricks/LTX-Video": |
| from diffusers import LTXImageToVideoPipeline |
|
|
| return LTXImageToVideoPipeline |
| return DiffusionPipeline |
|
|
