File size: 3,673 Bytes
d3fb87a |
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 |
import PIL
from PIL import Image
from PIL import ImageDraw
import numpy as np
import dlib
import cv2
import torch
import diffusers
from diffusers import StableDiffusionPipeline, DiffusionPipeline
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionControlNetImg2ImgPipeline, StableDiffusionImg2ImgPipeline
from chat_anything.face_generator.pipelines.lpw_stable_diffusion import StableDiffusionLongPromptWeightingPipeline, get_weighted_text_embeddings
from diffusers.schedulers import EulerAncestralDiscreteScheduler,DPMSolverMultistepScheduler # DPM++ SDE Karras
from chat_anything.face_generator.utils.generate import generate
class LongPromptGenerator():
prompt_template = "A portrait of a {}, fine face, nice looking"
negative_prompt = "easynegative,Low resolution,Low quality, Opened Mouth"
# negative_prompt = "(((sexy))),paintings,loli,,big head,sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, age spot, glans, nsfw, nipples,extra fingers, ((extra arms)), (extra legs), mutated hands, (fused fingers), (too many fingers), (long neck:1.3)"
def __init__(self, model_dir, lora_path=None, prompt_template="{}", negative_prompt=""):
self.model_dir = model_dir
self.lora_path = lora_path
self.prompt_template = prompt_template
self.negative_prompt = negative_prompt
def load_model(self, *args, **kwargs):
# load model
try:
pipe = DiffusionPipeline.from_pretrained(self.model_dir, torch_dtype=torch.float16, **kwargs)
except:
pipe = StableDiffusionPipeline.from_pretrained(self.model_dir, torch_dtype=torch.float16, **kwargs)
pipe = pipe.to('cuda')
sche_conf = dict(pipe.scheduler.config)
fk_kwargs = ["skip_prk_steps","steps_offset","clip_sample","clip_sample_range","rescale_betas_zero_snr","timestep_spacing", "set_alpha_to_one"]
for k in fk_kwargs:
if k in sche_conf:
sche_conf.pop(k)
scheduler = DPMSolverMultistepScheduler(**sche_conf)
pipe.scheduler=scheduler
pipe_longprompt = StableDiffusionLongPromptWeightingPipeline(**pipe.components)
self.pipe, self.pipe_longprompt = pipe, pipe_longprompt
if self.lora_path is not None:
pipe.load_lora_weights(self.lora_path)
self.pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(self.model_dir, **pipe.components)
def generate(
self,
prompt,
do_inversion=False,
**kwargs,
):
"""
Face control generating.
"""
print('GENERATING:', prompt)
if not do_inversion:
generating_conf = {
"pipe": self.pipe,
"prompt": prompt,
"negative_prompt": self.negative_prompt,
"num_inference_steps": 25,
"guidance_scale": 7,
}
else:
assert 'image' in kwargs, 'doing inversion, prepare the init image please PIL Image'
init_image = kwargs['image']
generating_conf = {
"pipe": self.pipe_img2img,
"prompt": prompt,
"negative_prompt": self.negative_prompt,
"image": init_image,
"num_inference_steps": 25,
"guidance_scale": 7,
"strength": kwargs.pop('strength', 0.9),
}
pipe_out = generate(**generating_conf)
generated_img = pipe_out[0][0]
return generated_img |