File size: 5,251 Bytes
8222fd4 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
from skimage.filters import threshold_otsu
from PIL import Image
from diffusers import AutoencoderKL, StableDiffusionPipeline, UNet2DConditionModel,UniPCMultistepScheduler
from accelerate import Accelerator
from accelerate.utils import ProjectConfiguration
from transformers import CLIPTextModel,CLIPTokenizer
import torch
import numpy as np
from skimage.filters import threshold_otsu
from model.controlnet_SPADE import ControlNetModel
from model.ControlnetPipeline import StableDiffusionControlNetPipeline_SPADE
class Pipeline_demo:
def __init__(self,
unet_ckpt="pretrain_weights/unet",
controlnet_ckpt="pretrain_weights/controlnet",
pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
is_cpu=False,
mixed_precision = "fp16",
enable_xformers_memory_efficient_attention=True,
output_dir = "output"
) -> None:
accelerator_project_config = ProjectConfiguration(project_dir=output_dir, logging_dir=output_dir)
self.is_cpu = is_cpu
if is_cpu:
mixed_precision = "no"
self.accelerator = Accelerator(
gradient_accumulation_steps=1,
mixed_precision=mixed_precision,
log_with="tensorboard",
cpu= True if is_cpu else False,
project_config=accelerator_project_config,
)
weight_dtype = torch.float32
if not is_cpu and self.accelerator.mixed_precision == "fp16":
weight_dtype = torch.float16
self.text_encoder = CLIPTextModel.from_pretrained(
pretrained_model_name_or_path, subfolder="text_encoder", revision=None
)
self.tokenizer = CLIPTokenizer.from_pretrained(
pretrained_model_name_or_path, subfolder="tokenizer", revision=None
)
self.unet = UNet2DConditionModel.from_pretrained(unet_ckpt)
self.vae = AutoencoderKL.from_pretrained(
pretrained_model_name_or_path, subfolder="vae", revision=None
)
self.controlnet = ControlNetModel.from_pretrained(controlnet_ckpt)
self.pipeline1 = StableDiffusionPipeline.from_pretrained(
pretrained_model_name_or_path,
vae=self.accelerator.unwrap_model(self.vae),
text_encoder=self.accelerator.unwrap_model(self.text_encoder),
tokenizer=self.tokenizer,
unet=self.accelerator.unwrap_model(self.unet),
safety_checker=None,
revision=None,
torch_dtype=weight_dtype,
)
self.pipeline2 = StableDiffusionControlNetPipeline_SPADE.from_pretrained(
pretrained_model_name_or_path,
vae=self.accelerator.unwrap_model(self.vae),
text_encoder=self.accelerator.unwrap_model(self.text_encoder),
tokenizer=self.tokenizer,
unet=self.accelerator.unwrap_model(self.unet),
controlnet=self.accelerator.unwrap_model(self.controlnet),
safety_checker=None,
revision=None,
torch_dtype=weight_dtype,
)
self.pipeline1 = self.pipeline1.to(self.accelerator.device)
self.pipeline1.set_progress_bar_config(disable=False)
self.pipeline2 = self.pipeline2.to(self.accelerator.device)
self.pipeline2.set_progress_bar_config(disable=False)
self.pipeline2.scheduler = UniPCMultistepScheduler.from_config(self.pipeline2.scheduler.config)
if not is_cpu and enable_xformers_memory_efficient_attention:
self.pipeline1.enable_xformers_memory_efficient_attention()
self.pipeline2.enable_xformers_memory_efficient_attention()
def binarize(self,image_array):
otsu_thresh = threshold_otsu(image_array)
# Apply Otsu's threshold to binarize the image
binarized_otsu_array = np.where(image_array > otsu_thresh, 255, 0)
# Convert the binarized array back into an image
binarized_otsu_image = Image.fromarray(binarized_otsu_array.astype(np.uint8)).convert("RGB")
return binarized_otsu_image
def generate_semantic_map(self, semantic_map_prompt,seed=None):
if seed is None:
generator = None
else:
generator = torch.Generator(device=self.accelerator.device).manual_seed(seed)
device_type = "cpu" if self.is_cpu else "cuda"
with torch.autocast(device_type):
image_anno = self.pipeline1(semantic_map_prompt, num_inference_steps=20, generator=generator).images[0]
bin_img = self.binarize(np.array(image_anno))
return bin_img
def generate_image(self, image_prompt, bin_img,seed=None):
if seed is None:
generator = None
else:
generator = torch.Generator(device=self.accelerator.device).manual_seed(seed)
device_type = "cpu" if self.is_cpu else "cuda"
bin_img = Image.fromarray(bin_img)
with torch.autocast(device_type):
image_img = self.pipeline2(image_prompt, bin_img, num_inference_steps=20, generator=generator).images[0]
return image_img
|