File size: 7,259 Bytes
d5f497d 6c91ee7 66fd925 6c91ee7 d5f497d 6c91ee7 66fd925 d5f497d 66fd925 6c91ee7 66fd925 6c91ee7 d5f497d 66fd925 d5f497d 6c91ee7 66fd925 6c91ee7 d5f497d 66fd925 d5f497d 66fd925 6c91ee7 66fd925 6c91ee7 66fd925 d5f497d 8004741 d5f497d 66fd925 6c91ee7 9de30d4 66fd925 6c91ee7 2502de8 66fd925 6c91ee7 d5f497d 2502de8 66fd925 cd4f227 e9f3ef9 66fd925 e9f3ef9 66fd925 e9f3ef9 fad18b4 66fd925 78ad020 66fd925 d5f497d 66fd925 d5f497d d890da3 d5f497d 20c2217 83bde13 20c2217 d5f497d f92dc60 d5f497d d890da3 d5f497d 6c91ee7 d5f497d 66fd925 d5f497d 66fd925 d5f497d 78ad020 66fd925 d5f497d 66fd925 9de30d4 d5f497d 66fd925 fad18b4 7132521 d5f497d 66fd925 9de30d4 78ad020 8004741 |
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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
import spaces
import random
import torch
import cv2
import insightface
import gradio as gr
import numpy as np
from huggingface_hub import snapshot_download
from transformers import CLIPVisionModelWithProjection,CLIPImageProcessor
from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter_FaceID import StableDiffusionXLPipeline
from kolors.models.modeling_chatglm import ChatGLMModel
from kolors.models.tokenization_chatglm import ChatGLMTokenizer
from diffusers import AutoencoderKL
from kolors.models.unet_2d_condition import UNet2DConditionModel
from diffusers import EulerDiscreteScheduler
from PIL import Image
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
device = "cuda"
ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
ckpt_dir_faceid = snapshot_download(repo_id="Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus")
text_encoder = ChatGLMModel.from_pretrained(f'{ckpt_dir}/text_encoder', torch_dtype=torch.float16).half().to(device)
tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half().to(device)
scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().to(device)
clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(f'{ckpt_dir_faceid}/clip-vit-large-patch14-336', ignore_mismatched_sizes=True)
clip_image_encoder.to(device)
clip_image_processor = CLIPImageProcessor(size = 336, crop_size = 336)
pipe = StableDiffusionXLPipeline(
vae = vae,
text_encoder = text_encoder,
tokenizer = tokenizer,
unet = unet,
scheduler = scheduler,
face_clip_encoder = clip_image_encoder,
face_clip_processor = clip_image_processor,
force_zeros_for_empty_prompt = False,
)
class FaceInfoGenerator():
def __init__(self, root_dir = "./"):
self.app = FaceAnalysis(name = 'antelopev2', root = root_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
self.app.prepare(ctx_id = 0, det_size = (640, 640))
def get_faceinfo_one_img(self, face_image):
face_info = self.app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
if len(face_info) == 0:
face_info = None
else:
face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
return face_info
def face_bbox_to_square(bbox):
## l, t, r, b to square l, t, r, b
l,t,r,b = bbox
cent_x = (l + r) / 2
cent_y = (t + b) / 2
w, h = r - l, b - t
r = max(w, h) / 2
l0 = cent_x - r
r0 = cent_x + r
t0 = cent_y - r
b0 = cent_y + r
return [l0, t0, r0, b0]
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
@spaces.GPU
def infer(prompt,
image = None,
negative_prompt = "nsfw,脸部阴影,低分辨率,jpeg伪影、模糊、糟糕,黑脸,霓虹灯",
seed = 66,
randomize_seed = False,
guidance_scale = 5.0,
num_inference_steps = 50
):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator().manual_seed(seed)
global pipe
pipe = pipe.to(device)
pipe.load_ip_adapter_faceid_plus(f'{ckpt_dir_faceid}/ipa-faceid-plus.bin', device = device)
scale = 0.8
pipe.set_face_fidelity_scale(scale)
face_info_generator = FaceInfoGenerator(root_dir = "./")
face_info = face_info_generator.get_faceinfo_one_img(image)
face_bbox_square = face_bbox_to_square(face_info["bbox"])
crop_image = image.crop(face_bbox_square)
crop_image = crop_image.resize((336, 336))
crop_image = [crop_image]
face_embeds = torch.from_numpy(np.array([face_info["embedding"]]))
face_embeds = face_embeds.to(device, dtype = torch.float16)
image = pipe(
prompt = prompt,
negative_prompt = negative_prompt,
height = 1024,
width = 1024,
num_inference_steps= num_inference_steps,
guidance_scale = guidance_scale,
num_images_per_prompt = 1,
generator = generator,
face_crop_image = crop_image,
face_insightface_embeds = face_embeds
).images[0]
return image, seed
examples = [
["穿着晚礼服,在星光下的晚宴场景中,烛光闪闪,整个场景洋溢着浪漫而奢华的氛围", "image/image1.png"],
["西部牛仔,牛仔帽,荒野大镖客,背景是西部小镇,仙人掌,,日落余晖, 暖色调, 使用XT4胶片拍摄, 噪点, 晕影, 柯达胶卷,复古", "image/image2.png"]
]
css="""
#col-left {
margin: 0 auto;
max-width: 600px;
}
#col-right {
margin: 0 auto;
max-width: 750px;
}
#button {
color: blue;
}
"""
def load_description(fp):
with open(fp, 'r', encoding='utf-8') as f:
content = f.read()
return content
with gr.Blocks(css=css) as Kolors:
gr.HTML(load_description("assets/title.md"))
with gr.Row():
with gr.Column(elem_id="col-left"):
with gr.Row():
prompt = gr.Textbox(
label="Prompt",
placeholder="Enter your prompt",
lines=2
)
with gr.Row():
image = gr.Image(label="Image", type="pil")
with gr.Accordion("Advanced Settings", open=False):
negative_prompt = gr.Textbox(
label="Negative prompt",
placeholder="Enter a negative prompt",
visible=True,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=5.0,
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=10,
maximum=50,
step=1,
value=25,
)
with gr.Row():
button = gr.Button("Run", elem_id="button")
with gr.Column(elem_id="col-right"):
result = gr.Image(label="Result", show_label=False)
seed_used = gr.Number(label="Seed Used")
with gr.Row():
gr.Examples(
fn = infer,
examples = examples,
inputs = [prompt, image],
outputs = [result, seed_used],
)
button.click(
fn = infer,
inputs = [prompt, image, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps],
outputs = [result, seed_used]
)
Kolors.queue().launch(debug=True)
|