avinjcy/custom-diffusion

1from __future__ import annotations2 3import gc4import pathlib5import sys6 7import gradio as gr8import PIL.Image9import numpy as np10 11import torch12from diffusers import StableDiffusionPipeline13sys.path.insert(0, './custom-diffusion')14 15 16class InferencePipeline:17    def __init__(self):18        self.pipe = None19        self.device = torch.device(20            'cuda:0' if torch.cuda.is_available() else 'cpu')21        self.weight_path = None22 23    def clear(self) -> None:24        self.weight_path = None25        del self.pipe26        self.pipe = None27        torch.cuda.empty_cache()28        gc.collect()29 30    @staticmethod31    def get_weight_path(name: str) -> pathlib.Path:32        curr_dir = pathlib.Path(__file__).parent33        return curr_dir / name34 35    def load_pipe(self, model_id: str, filename: str) -> None:36        weight_path = self.get_weight_path(filename)37        if weight_path == self.weight_path:38            return39        self.weight_path = weight_path40        weight = torch.load(self.weight_path, map_location=self.device)41 42        if self.device.type == 'cpu':43            pipe = StableDiffusionPipeline.from_pretrained(model_id)44        else:45            pipe = StableDiffusionPipeline.from_pretrained(46                model_id, torch_dtype=torch.float16)47            pipe = pipe.to(self.device)48 49        from src import diffuser_training50        diffuser_training.load_model(pipe.text_encoder, pipe.tokenizer, pipe.unet, weight_path, compress=False)51 52        self.pipe = pipe53 54    def run(55        self,56        base_model: str,57        weight_name: str,58        prompt: str,59        seed: int,60        n_steps: int,61        guidance_scale: float,62        eta: float,63        batch_size: int,64        resolution: int,65    ) -> PIL.Image.Image:66        if not torch.cuda.is_available():67            raise gr.Error('CUDA is not available.')68 69        self.load_pipe(base_model, weight_name)70 71        generator = torch.Generator(device=self.device).manual_seed(seed)72        out = self.pipe([prompt]*batch_size,73                        num_inference_steps=n_steps,74                        guidance_scale=guidance_scale,75                        height=resolution, width=resolution,76                        eta = eta,77                        generator=generator)  # type: ignore78        torch.cuda.empty_cache()79        out = out.images80        out = PIL.Image.fromarray(np.hstack([np.array(x) for x in out]))81        return out82