Spaces:
Running on Zero
Running on Zero
File size: 3,419 Bytes
97f62f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | 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
if dtype is None and str(device).startswith("mps"):
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
# Fall back to per-item generation (e.g. model can't batch, or OOM).
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
|