| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import torch |
| from diffusers import DiffusionPipeline |
|
|
| from .base import TextToImageGenerator |
|
|
|
|
| class DiffusersTextToImageAdapter(TextToImageGenerator): |
| 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.float16 |
| pipeline = DiffusionPipeline.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, prompts: list[str], *, generator: Any | None = None, **kwargs) -> list[Any]: |
| out = self.pipe(prompt=prompts, generator=generator, **kwargs) |
| if hasattr(out, "images"): |
| return list(out.images) |
| if isinstance(out, list): |
| return out |
| raise TypeError("Text-to-image pipeline output does not expose `.images`.") |
|
|
| def generate_batch(self, prompts, *, generators=None, **kwargs): |
| out = self.pipe(prompt=list(prompts), generator=generators, **kwargs) |
| if hasattr(out, "images"): |
| return list(out.images) |
| if isinstance(out, list): |
| return out |
| raise TypeError("Text-to-image pipeline output does not expose `.images`.") |
|
|
