| import itertools |
| from typing import Dict, List, Any |
| import torch |
| from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline |
| from PIL import Image |
| from skimage.exposure import match_histograms, adjust_gamma |
| import numpy as np |
| import diffusers |
|
|
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| self.txt2img_pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16).to("cuda") |
| self.txt2img_pipe.register_to_config(requires_safety_checker=False) |
| self.txt2img_pipe.register_to_config(safety_checker = None) |
| self.img2img_pipe = StableDiffusionImg2ImgPipeline( |
| vae=self.txt2img_pipe.vae, |
| text_encoder=self.txt2img_pipe.text_encoder, |
| tokenizer=self.txt2img_pipe.tokenizer, |
| unet=self.txt2img_pipe.unet, |
| scheduler=self.txt2img_pipe.scheduler, |
| safety_checker=self.txt2img_pipe.safety_checker, |
| feature_extractor=self.txt2img_pipe.feature_extractor, |
| ).to("cuda") |
| |
| def generate_story(self, prompts, do_img2img, num_inference_steps, guidance_scale): |
| first_img = self.txt2img_pipe(prompts[0], num_inference_steps = num_inference_steps, guidance_scale = guidance_scale).images[0] |
| ret = [first_img] |
| for prompt in prompts[1:]: |
| if do_img2img: |
| out = self.img2img_pipe( |
| prompt=prompt, |
| strength=0.9, |
| image=first_img, |
| num_inference_steps=num_inference_steps, |
| guidance_scale=guidance_scale) |
| else: |
| out = self.txt2img_pipe( |
| prompt=prompt, |
| num_inference_steps=num_inference_steps, |
| guidance_scale=guidance_scale) |
| ret.append(out.images[0]) |
| return ret |
| |
|
|
|
|
| def __call__(self, data: Any) -> List[List[Dict[str, float]]]: |
| prompts = data["prompts"] |
| num_images_per_prompt = data.get("num_images_per_prompt", 1) |
| do_img2img = data.get("do_img2img", True) |
| match_histogram = data.get("match_histogram", True) |
| adjust_contrast = data.get("adjust_contrast", True) |
| num_inference_steps = data.get("num_inference_steps", 25) |
| guidance_scale = data.pop("guidance_scale", 7.5) |
| all_stories = [] |
| for _ in range(num_images_per_prompt): |
| story = self.generate_story(prompts, do_img2img, num_inference_steps, guidance_scale) |
| if match_histogram: |
| story = [match_colors(story[0], img) for img in story] |
| if adjust_contrast: |
| story = [add_contrast(img) for img in story] |
| all_stories.append(story) |
| all_stories = [self.generate_story(prompts, do_img2img, num_inference_steps, guidance_scale) for _ in range(num_images_per_prompt)] |
| return stories_to_grid(all_stories) |
| |
| def match_colors(src,tar): |
| ret = match_histograms(np.array(tar), np.array(src), multichannel=True) |
| return Image.fromarray(ret.astype('uint8')) |
|
|
| def add_contrast(image, gamma_cor=1.2): |
| ar = np.array(image) |
| adjusted = adjust_gamma(ar, gamma_cor).astype(np.uint8) |
| return Image.fromarray(adjusted) |
|
|
| def stories_to_grid(stories): |
| rows = len(stories) |
| cols = len(stories[0]) |
| imgs = list(itertools.chain(*stories)) |
| return image_grid(imgs, rows, cols) |
|
|
| def image_grid(imgs, rows, cols): |
| assert len(imgs) == rows*cols |
|
|
| w, h = imgs[0].size |
| grid = Image.new('RGB', size=(cols*w, rows*h)) |
| |
| for i, img in enumerate(imgs): |
| grid.paste(img, box=(i%cols*w, i//cols*h)) |
| return grid |
