Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import logging | |
| import os | |
| import random | |
| import sys | |
| import numpy as np | |
| import PIL.Image | |
| import torch | |
| from diffusers import (DDIMPipeline, DDIMScheduler, DDPMPipeline, | |
| DiffusionPipeline, PNDMPipeline, PNDMScheduler) | |
| HF_TOKEN = os.environ['HF_TOKEN'] | |
| formatter = logging.Formatter( | |
| '[%(asctime)s] %(name)s %(levelname)s: %(message)s', | |
| datefmt='%Y-%m-%d %H:%M:%S') | |
| stream_handler = logging.StreamHandler(stream=sys.stdout) | |
| stream_handler.setLevel(logging.INFO) | |
| stream_handler.setFormatter(formatter) | |
| logger = logging.getLogger(__name__) | |
| logger.setLevel(logging.INFO) | |
| logger.propagate = False | |
| logger.addHandler(stream_handler) | |
| class Model: | |
| MODEL_NAMES = [ | |
| 'ddpm-128-exp000', | |
| ] | |
| def __init__(self, device: str | torch.device): | |
| self.device = torch.device(device) | |
| self._download_all_models() | |
| self.model_name = self.MODEL_NAMES[0] | |
| self.scheduler_type = 'DDIM' | |
| self.pipeline = self._load_pipeline(self.model_name, | |
| self.scheduler_type) | |
| self.rng = random.Random() | |
| def _load_pipeline(self, model_name: str, | |
| scheduler_type: str) -> DiffusionPipeline: | |
| repo_id = f'hysts/diffusers-anime-faces-{model_name}' | |
| if scheduler_type == 'DDPM': | |
| pipeline = DDPMPipeline.from_pretrained(repo_id, | |
| use_auth_token=HF_TOKEN) | |
| elif scheduler_type == 'DDIM': | |
| pipeline = DDIMPipeline.from_pretrained(repo_id, | |
| use_auth_token=HF_TOKEN) | |
| config, _ = DDIMScheduler.extract_init_dict( | |
| dict(pipeline.scheduler.config)) | |
| pipeline.scheduler = DDIMScheduler(**config) | |
| elif scheduler_type == 'PNDM': | |
| pipeline = PNDMPipeline.from_pretrained(repo_id, | |
| use_auth_token=HF_TOKEN) | |
| config, _ = PNDMScheduler.extract_init_dict( | |
| dict(pipeline.scheduler.config)) | |
| pipeline.scheduler = PNDMScheduler(**config) | |
| else: | |
| raise ValueError | |
| return pipeline | |
| def set_pipeline(self, model_name: str, scheduler_type: str) -> None: | |
| logger.info('--- set_pipeline ---') | |
| logger.info(f'{model_name=}, {scheduler_type=}') | |
| if model_name == self.model_name and scheduler_type == self.scheduler_type: | |
| logger.info('Skipping') | |
| logger.info('--- done ---') | |
| return | |
| self.model_name = model_name | |
| self.scheduler_type = scheduler_type | |
| self.pipeline = self._load_pipeline(model_name, scheduler_type) | |
| logger.info('--- done ---') | |
| def _download_all_models(self) -> None: | |
| for name in self.MODEL_NAMES: | |
| self._load_pipeline(name, 'DDPM') | |
| def generate(self, | |
| seed: int, | |
| num_steps: int, | |
| num_images: int = 1) -> list[PIL.Image.Image]: | |
| logger.info('--- generate ---') | |
| logger.info(f'{seed=}, {num_steps=}') | |
| torch.manual_seed(seed) | |
| if self.scheduler_type == 'DDPM': | |
| res = self.pipeline(batch_size=num_images, | |
| torch_device=self.device)['sample'] | |
| elif self.scheduler_type in ['DDIM', 'PNDM']: | |
| res = self.pipeline(batch_size=num_images, | |
| torch_device=self.device, | |
| num_inference_steps=num_steps)['sample'] | |
| else: | |
| raise ValueError | |
| logger.info('--- done ---') | |
| return res | |
| def run( | |
| self, | |
| model_name: str, | |
| scheduler_type: str, | |
| num_steps: int, | |
| seed: int, | |
| ) -> PIL.Image.Image: | |
| self.set_pipeline(model_name, scheduler_type) | |
| if scheduler_type == 'PNDM': | |
| num_steps = max(4, min(num_steps, 100)) | |
| return self.generate(seed, num_steps)[0] | |
| def to_grid(images: list[PIL.Image.Image], | |
| ncols: int = 2) -> PIL.Image.Image: | |
| images = [np.asarray(image) for image in images] | |
| nrows = (len(images) + ncols - 1) // ncols | |
| h, w = images[0].shape[:2] | |
| d = nrows * ncols - len(images) | |
| if d > 0: | |
| images += [np.full((h, w, 3), 255, dtype=np.uint8)] * d | |
| grid = np.asarray(images).reshape(2, 2, h, w, 3).transpose( | |
| 0, 2, 1, 3, 4).reshape(2 * h, 2 * w, 3) | |
| return PIL.Image.fromarray(grid) | |
| def run_simple(self) -> PIL.Image.Image: | |
| self.set_pipeline(self.MODEL_NAMES[0], 'DDIM') | |
| seed = self.rng.randint(0, 100000) | |
| images = self.generate(seed, num_steps=10, num_images=4) | |
| return self.to_grid(images, 2) | |