Spaces:
Running on Zero
Running on Zero
File size: 1,704 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 | 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
if dtype is None and str(device).startswith("mps"):
dtype = torch.float32
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`.")
|