Small cleanup

This commit is contained in:
Ean Garvey
2024-02-02 11:25:41 -06:00
parent a43c5590d5
commit 019ba7051d
2 changed files with 44 additions and 40 deletions

View File

@@ -374,46 +374,6 @@ class StableDiffusion(SharkPipelineBase):
pil_images = self.image_processor.numpy_to_pil(images)
return pil_images
# def process_sd_init_image(self, sd_init_image, resample_type):
# if isinstance(sd_init_image, list):
# images = []
# for img in sd_init_image:
# img, _ = self.process_sd_init_image(img, resample_type)
# images.append(img)
# is_img2img = True
# return images, is_img2img
# if isinstance(sd_init_image, str):
# if os.path.isfile(sd_init_image):
# sd_init_image = Image.open(sd_init_image, mode="r").convert("RGB")
# image, is_img2img = self.process_sd_init_image(
# sd_init_image, resample_type
# )
# else:
# image = None
# is_img2img = False
# elif isinstance(sd_init_image, Image.Image):
# image = sd_init_image.convert("RGB")
# elif sd_init_image:
# image = sd_init_image["image"].convert("RGB")
# else:
# image = None
# is_img2img = False
# if image:
# resample_type = (
# resamplers[resample_type]
# if resample_type in resampler_list
# # Fallback to Lanczos
# else Image.Resampling.LANCZOS
# )
# image = image.resize((self.width, self.height), resample=resample_type)
# image_arr = np.stack([np.array(i) for i in (image,)], axis=0)
# image_arr = image_arr / 255.0
# image_arr = torch.from_numpy(image_arr).permute(0, 3, 1, 2).to(self.dtype)
# image_arr = 2 * (image_arr - 0.5)
# is_img2img = True
# image = image_arr
# return image, is_img2img
def generate_images(
self,
prompt,

View File

@@ -1,6 +1,8 @@
import os
import re
import json
import torch
import numpy as np
from csv import DictWriter
from PIL import Image, PngImagePlugin
@@ -8,6 +10,7 @@ from pathlib import Path
from datetime import datetime as dt
from base64 import decode
resamplers = {
"Lanczos": Image.Resampling.LANCZOS,
"Nearest Neighbor": Image.Resampling.NEAREST,
@@ -158,3 +161,44 @@ def resize_stencil(image: Image.Image, width, height, resampler_type=None):
resampler = resamplers["Nearest Neighbor"]
new_image = image.resize((n_width, n_height), resampler=resampler)
return new_image, n_width, n_height
def process_sd_init_image(self, sd_init_image, resample_type):
if isinstance(sd_init_image, list):
images = []
for img in sd_init_image:
img, _ = self.process_sd_init_image(img, resample_type)
images.append(img)
is_img2img = True
return images, is_img2img
if isinstance(sd_init_image, str):
if os.path.isfile(sd_init_image):
sd_init_image = Image.open(sd_init_image, mode="r").convert("RGB")
image, is_img2img = self.process_sd_init_image(
sd_init_image, resample_type
)
else:
image = None
is_img2img = False
elif isinstance(sd_init_image, Image.Image):
image = sd_init_image.convert("RGB")
elif sd_init_image:
image = sd_init_image["image"].convert("RGB")
else:
image = None
is_img2img = False
if image:
resample_type = (
resamplers[resample_type]
if resample_type in resampler_list
# Fallback to Lanczos
else Image.Resampling.LANCZOS
)
image = image.resize((self.width, self.height), resample=resample_type)
image_arr = np.stack([np.array(i) for i in (image,)], axis=0)
image_arr = image_arr / 255.0
image_arr = torch.from_numpy(image_arr).permute(0, 3, 1, 2).to(self.dtype)
image_arr = 2 * (image_arr - 0.5)
is_img2img = True
image = image_arr
return image, is_img2img