FlexTheAi/Flexstorydiff
0
1# Prediction interface for Cog ⚙️2# https://cog.run/python3 4import os5import copy6import random7import subprocess8import numpy as np9import time10import torch11import torch.nn.functional as F12from PIL import ImageFont13from cog import BasePredictor, Input, Path, BaseModel14from diffusers import StableDiffusionXLPipeline, DDIMScheduler15from diffusers.utils import load_image16 17from utils import PhotoMakerStableDiffusionXLPipeline18from utils.style_template import styles19from utils.gradio_utils import (20 AttnProcessor2_0 as AttnProcessor,21) # with torch2 installed22from utils.gradio_utils import cal_attn_mask_xl23from utils.utils import get_comic24 25MODEL_URL = "https://weights.replicate.delivery/default/HVision_NKU/StoryDiffusion.tar"26MODEL_CACHE = "model_weights"27STYLE_NAMES = list(styles.keys())28DEFAULT_STYLE_NAME = "Japanese Anime"29 30global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet31global sa32, sa6432global write33global height, width34 35 36"""37# load and upload the weights to replicate.delivery for faster booting on Replicate38models_dict = {39 "RealVision": "SG161222/RealVisXL_V4.0",40 "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y",41}42# photomaker_path = hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")43photomaker_path = f"{MODEL_CACHE}/PhotoMaker/photomaker-v1.bin"44 45pipe_unstable = PhotoMakerStableDiffusionXLPipeline.from_pretrained(46 models_dict["Unstable"],47 torch_dtype=torch.float16,48 use_safetensors=False,49)50pipe_unstable.save_pretrained(f"{MODEL_CACHE}/Unstable/stablediffusionapi/sdxl-unstable-diffusers-y")51 52pipe_realvision = PhotoMakerStableDiffusionXLPipeline.from_pretrained(53 models_dict["RealVision"], torch_dtype=torch.float16, use_safetensors=True54)55pipe_realvision.save_pretrained(f"{MODEL_CACHE}/RealVision/SG161222/RealVisXL_V4.0")56"""57 58 59class ModelOutput(BaseModel):60 comic: Path61 individual_images: list[Path]62 63 64def download_weights(url, dest):65 start = time.time()66 print("downloading url: ", url)67 print("downloading to: ", dest)68 subprocess.check_call(["pget", "-x", url, dest], close_fds=False)69 print("downloading took: ", time.time() - start)70 71 72def setup_seed(seed):73 torch.manual_seed(seed)74 torch.cuda.manual_seed_all(seed)75 np.random.seed(seed)76 random.seed(seed)77 torch.backends.cudnn.deterministic = True78 79 80def apply_style_positive(style_name: str, positive: str):81 p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])82 return p.replace("{prompt}", positive)83 84 85def apply_style(style_name: str, positives: list, negative: str = ""):86 p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])87 return [88 p.replace("{prompt}", positive) for positive in positives89 ], n + " " + negative90 91 92def set_attention_processor(unet, id_length, is_ipadapter=False):93 global total_count94 total_count = 095 attn_procs = {}96 for name in unet.attn_processors.keys():97 cross_attention_dim = (98 None99 if name.endswith("attn1.processor")100 else unet.config.cross_attention_dim101 )102 if name.startswith("mid_block"):103 hidden_size = unet.config.block_out_channels[-1]104 elif name.startswith("up_blocks"):105 block_id = int(name[len("up_blocks.")])106 hidden_size = list(reversed(unet.config.block_out_channels))[block_id]107 elif name.startswith("down_blocks"):108 block_id = int(name[len("down_blocks.")])109 hidden_size = unet.config.block_out_channels[block_id]110 if cross_attention_dim is None:111 if name.startswith("up_blocks"):112 attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)113 total_count += 1114 else:115 attn_procs[name] = AttnProcessor()116 else:117 if is_ipadapter:118 attn_procs[name] = IPAttnProcessor2_0(119 hidden_size=hidden_size,120 cross_attention_dim=cross_attention_dim,121 scale=1,122 num_tokens=4,123 ).to(unet.device, dtype=torch.float16)124 else:125 attn_procs[name] = AttnProcessor()126 127 unet.set_attn_processor(copy.deepcopy(attn_procs))128 print("Successfully load paired self-attention")129 print(f"Number of the processor : {total_count}")130 131 132#################################################133########Consistent Self-Attention################134#################################################135class SpatialAttnProcessor2_0(torch.nn.Module):136 r"""137 Attention processor for IP-Adapater for PyTorch 2.0.138 Args:139 hidden_size (`int`):140 The hidden size of the attention layer.141 cross_attention_dim (`int`):142 The number of channels in the `encoder_hidden_states`.143 text_context_len (`int`, defaults to 77):144 The context length of the text features.145 scale (`float`, defaults to 1.0):146 the weight scale of image prompt.147 """148 149 def __init__(150 self,151 hidden_size=None,152 cross_attention_dim=None,153 id_length=4,154 device="cuda",155 dtype=torch.float16,156 ):157 super().__init__()158 if not hasattr(F, "scaled_dot_product_attention"):159 raise ImportError(160 "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."161 )162 self.device = device163 self.dtype = dtype164 self.hidden_size = hidden_size165 self.cross_attention_dim = cross_attention_dim166 self.total_length = id_length + 1167 self.id_length = id_length168 self.id_bank = {}169 170 def __call__(171 self,172 attn,173 hidden_states,174 encoder_hidden_states=None,175 attention_mask=None,176 temb=None,177 ):178 global total_count, attn_count, cur_step, mask1024, mask4096179 global sa32, sa64180 global write181 global height, width182 if write:183 self.id_bank[cur_step] = [184 hidden_states[: self.id_length],185 hidden_states[self.id_length :],186 ]187 else:188 encoder_hidden_states = torch.cat(189 (190 self.id_bank[cur_step][0].to(self.device),191 hidden_states[:1],192 self.id_bank[cur_step][1].to(self.device),193 hidden_states[1:],194 )195 )196 # skip in early step197 if cur_step < 5:198 hidden_states = self.__call2__(199 attn, hidden_states, encoder_hidden_states, attention_mask, temb200 )201 else: # 256 1024 4096202 random_number = random.random()203 if cur_step < 20:204 rand_num = 0.3205 else:206 rand_num = 0.1207 if random_number > rand_num:208 if not write:209 if hidden_states.shape[1] == (height // 32) * (width // 32):210 attention_mask = mask1024[211 mask1024.shape[0] // self.total_length * self.id_length :212 ]213 else:214 attention_mask = mask4096[215 mask4096.shape[0] // self.total_length * self.id_length :216 ]217 else:218 if hidden_states.shape[1] == (height // 32) * (width // 32):219 attention_mask = mask1024[220 : mask1024.shape[0] // self.total_length * self.id_length,221 : mask1024.shape[0] // self.total_length * self.id_length,222 ]223 else:224 attention_mask = mask4096[225 : mask4096.shape[0] // self.total_length * self.id_length,226 : mask4096.shape[0] // self.total_length * self.id_length,227 ]228 hidden_states = self.__call1__(229 attn, hidden_states, encoder_hidden_states, attention_mask, temb230 )231 else:232 hidden_states = self.__call2__(233 attn, hidden_states, None, attention_mask, temb234 )235 attn_count += 1236 if attn_count == total_count:237 attn_count = 0238 cur_step += 1239 mask1024, mask4096 = cal_attn_mask_xl(240 self.total_length,241 self.id_length,242 sa32,243 sa64,244 height,245 width,246 device=self.device,247 dtype=self.dtype,248 )249 250 return hidden_states251 252 def __call1__(253 self,254 attn,255 hidden_states,256 encoder_hidden_states=None,257 attention_mask=None,258 temb=None,259 ):260 residual = hidden_states261 if attn.spatial_norm is not None:262 hidden_states = attn.spatial_norm(hidden_states, temb)263 input_ndim = hidden_states.ndim264 265 if input_ndim == 4:266 total_batch_size, channel, height, width = hidden_states.shape267 hidden_states = hidden_states.view(268 total_batch_size, channel, height * width269 ).transpose(1, 2)270 total_batch_size, nums_token, channel = hidden_states.shape271 img_nums = total_batch_size // 2272 hidden_states = hidden_states.view(-1, img_nums, nums_token, channel).reshape(273 -1, img_nums * nums_token, channel274 )275 276 batch_size, sequence_length, _ = hidden_states.shape277 278 if attn.group_norm is not None:279 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(280 1, 2281 )282 283 query = attn.to_q(hidden_states)284 285 if encoder_hidden_states is None:286 encoder_hidden_states = hidden_states # B, N, C287 else:288 encoder_hidden_states = encoder_hidden_states.view(289 -1, self.id_length + 1, nums_token, channel290 ).reshape(-1, (self.id_length + 1) * nums_token, channel)291 292 key = attn.to_k(encoder_hidden_states)293 value = attn.to_v(encoder_hidden_states)294 295 inner_dim = key.shape[-1]296 head_dim = inner_dim // attn.heads297 298 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)299 300 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)301 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)302 hidden_states = F.scaled_dot_product_attention(303 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False304 )305 306 hidden_states = hidden_states.transpose(1, 2).reshape(307 total_batch_size, -1, attn.heads * head_dim308 )309 hidden_states = hidden_states.to(query.dtype)310 311 # linear proj312 hidden_states = attn.to_out[0](hidden_states)313 # dropout314 hidden_states = attn.to_out[1](hidden_states)315 316 if input_ndim == 4:317 hidden_states = hidden_states.transpose(-1, -2).reshape(318 total_batch_size, channel, height, width319 )320 if attn.residual_connection:321 hidden_states = hidden_states + residual322 hidden_states = hidden_states / attn.rescale_output_factor323 # print(hidden_states.shape)324 return hidden_states325 326 def __call2__(327 self,328 attn,329 hidden_states,330 encoder_hidden_states=None,331 attention_mask=None,332 temb=None,333 ):334 residual = hidden_states335 336 if attn.spatial_norm is not None:337 hidden_states = attn.spatial_norm(hidden_states, temb)338 339 input_ndim = hidden_states.ndim340 341 if input_ndim == 4:342 batch_size, channel, height, width = hidden_states.shape343 hidden_states = hidden_states.view(344 batch_size, channel, height * width345 ).transpose(1, 2)346 347 batch_size, sequence_length, channel = hidden_states.shape348 # print(hidden_states.shape)349 if attention_mask is not None:350 attention_mask = attn.prepare_attention_mask(351 attention_mask, sequence_length, batch_size352 )353 # scaled_dot_product_attention expects attention_mask shape to be354 # (batch, heads, source_length, target_length)355 attention_mask = attention_mask.view(356 batch_size, attn.heads, -1, attention_mask.shape[-1]357 )358 359 if attn.group_norm is not None:360 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(361 1, 2362 )363 364 query = attn.to_q(hidden_states)365 366 if encoder_hidden_states is None:367 encoder_hidden_states = hidden_states # B, N, C368 else:369 encoder_hidden_states = encoder_hidden_states.view(370 -1, self.id_length + 1, sequence_length, channel371 ).reshape(-1, (self.id_length + 1) * sequence_length, channel)372 373 key = attn.to_k(encoder_hidden_states)374 value = attn.to_v(encoder_hidden_states)375 376 inner_dim = key.shape[-1]377 head_dim = inner_dim // attn.heads378 379 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)380 381 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)382 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)383 384 hidden_states = F.scaled_dot_product_attention(385 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False386 )387 388 hidden_states = hidden_states.transpose(1, 2).reshape(389 batch_size, -1, attn.heads * head_dim390 )391 hidden_states = hidden_states.to(query.dtype)392 393 # linear proj394 hidden_states = attn.to_out[0](hidden_states)395 # dropout396 hidden_states = attn.to_out[1](hidden_states)397 398 if input_ndim == 4:399 hidden_states = hidden_states.transpose(-1, -2).reshape(400 batch_size, channel, height, width401 )402 403 if attn.residual_connection:404 hidden_states = hidden_states + residual405 406 hidden_states = hidden_states / attn.rescale_output_factor407 408 return hidden_states409 410 411class Predictor(BasePredictor):412 def setup(self) -> None:413 """Load the model into memory to make running multiple predictions efficient"""414 415 models_dict = {416 "RealVision": "SG161222/RealVisXL_V4.0",417 "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y",418 }419 420 if not os.path.exists(MODEL_CACHE):421 download_weights(MODEL_URL, MODEL_CACHE)422 423 photomaker_path = f"{MODEL_CACHE}/PhotoMaker/photomaker-v1.bin"424 425 self.sdxl_pipe_unstable = StableDiffusionXLPipeline.from_pretrained(426 f"{MODEL_CACHE}/Unstable/sdxl/stablediffusionapi/sdxl-unstable-diffusers-y",427 torch_dtype=torch.float16,428 )429 self.sdxl_pipe_realvision = StableDiffusionXLPipeline.from_pretrained(430 f"{MODEL_CACHE}/RealVision/sdxl/SG161222/RealVisXL_V4.0",431 torch_dtype=torch.float16,432 )433 434 self.pipe_unstable = PhotoMakerStableDiffusionXLPipeline.from_pretrained(435 f"{MODEL_CACHE}/Unstable/stablediffusionapi/sdxl-unstable-diffusers-y",436 torch_dtype=torch.float16,437 use_safetensors=False,438 )439 self.pipe_unstable.load_photomaker_adapter(440 os.path.dirname(photomaker_path),441 subfolder="",442 weight_name=os.path.basename(photomaker_path),443 trigger_word="img", # define the trigger word444 )445 446 self.pipe_realvision = PhotoMakerStableDiffusionXLPipeline.from_pretrained(447 f"{MODEL_CACHE}/RealVision/SG161222/RealVisXL_V4.0",448 torch_dtype=torch.float16,449 use_safetensors=True,450 )451 self.pipe_realvision.load_photomaker_adapter(452 os.path.dirname(photomaker_path),453 subfolder="",454 weight_name=os.path.basename(photomaker_path),455 trigger_word="img", # define the trigger word456 )457 self.pipe_realvision.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)458 self.pipe_realvision.fuse_lora()459 460 @torch.inference_mode()461 def predict(462 self,463 sd_model: str = Input(464 description="Choose a model",465 choices=["Unstable", "RealVision"],466 default="Unstable",467 ),468 ref_image: Path = Input(469 description="Reference image for the character",470 default=None,471 ),472 character_description: str = Input(473 description="General description of the character. If ref_image above is provided, making sure to follow the class word you want to customize with the trigger word 'img', such as: 'man img' or 'woman img' or 'girl img'",474 default="a man, wearing black suit",475 ),476 negative_prompt: str = Input(477 description="Describe things you do not want to see in the output",478 default="bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",479 ),480 comic_description: str = Input(481 description="Comic Description. Each frame is divided by a new line. Only the first 10 prompts are valid for demo speed! For comic_description NOT using ref_image: (1) Support Typesetting Style and Captioning. By default, the prompt is used as the caption for each image. If you need to change the caption, add a '#' at the end of each line. Only the part after the '#' will be added as a caption to the image. (2) The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the '[NC]' at the beginning of the line.",482 default="at home, read new paper #at home, The newspaper says there is a treasure house in the forest.\non the road, near the forest\n[NC] The car on the road, near the forest #He drives to the forest in search of treasure.\n[NC]A tiger appeared in the forest, at night \nvery frightened, open mouth, in the forest, at night\nrunning very fast, in the forest, at night\n[NC] A house in the forest, at night #Suddenly, he discovers the treasure house!\nin the house filled with treasure, laughing, at night #He is overjoyed inside the house.",483 ),484 style_name: str = Input(485 description="Style template",486 choices=STYLE_NAMES,487 default=DEFAULT_STYLE_NAME,488 ),489 comic_style: str = Input(490 description="Select the comic style for the combined comic",491 choices=["Four Pannel", "Classic Comic Style"],492 default="Classic Comic Style",493 ),494 style_strength_ratio: int = Input(495 description="Style strength of Ref Image (%), only used if ref_image is provided",496 default=20,497 ge=15,498 le=50,499 ),500 image_width: int = Input(501 description="Width of output image",502 choices=[503 256,504 288,505 320,506 352,507 384,508 416,509 448,510 480,511 512,512 544,513 576,514 608,515 640,516 672,517 704,518 736,519 768,520 800,521 832,522 864,523 896,524 928,525 960,526 992,527 1024,528 ],529 default=768,530 ),531 image_height: int = Input(532 description="Height of output image",533 choices=[534 256,535 288,536 320,537 352,538 384,539 416,540 448,541 480,542 512,543 544,544 576,545 608,546 640,547 672,548 704,549 736,550 768,551 800,552 832,553 864,554 896,555 928,556 960,557 992,558 1024,559 ],560 default=768,561 ),562 num_steps: int = Input(563 description="Number of sample steps", ge=20, le=50, default=25564 ),565 guidance_scale: float = Input(566 description="Scale for classifier-free guidance", ge=0.1, le=10, default=5567 ),568 seed: int = Input(569 description="Random seed. Leave blank to randomize the seed", default=None570 ),571 sa32_setting: float = Input(572 description="The degree of Paired Attention at 32 x 32 self-attention layers",573 default=0.5,574 ge=0,575 le=1.0,576 ),577 sa64_setting: float = Input(578 description="The degree of Paired Attention at 64 x 64 self-attention layers",579 default=0.5,580 ge=0,581 le=1.0,582 ),583 num_ids: int = Input(584 description="Number of id images in total images. This should not exceed total number of line-separated prompts",585 default=3,586 ),587 output_format: str = Input(588 description="Format of the output images",589 choices=["webp", "jpg", "png"],590 default="webp",591 ),592 output_quality: int = Input(593 description="Quality of the output images, from 0 to 100. 100 is best quality, 0 is lowest quality",594 default=80,595 ge=0,596 le=100,597 ),598 ) -> ModelOutput:599 """Run a single prediction on the model"""600 601 global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet602 global sa32, sa64603 global write604 global height, width605 606 assert (607 len(character_description.strip()) > 0608 ), "Please provide the description of the character."609 610 if ref_image is not None:611 assert (612 "img" in character_description613 ), f"When using ref_image, please add the trigger word 'img' behind the class word you want to customize, such as: man img or woman img"614 assert (615 "[NC]" not in comic_description616 ), "You should not use trigger word [NC] when ref_image is provided."617 618 height = image_height619 width = image_width620 id_length = num_ids621 sa32 = sa32_setting622 sa64 = sa64_setting623 624 clipped_prompts = comic_description.splitlines()[:10]625 print(clipped_prompts)626 prompts = [627 (628 character_description + "," + prompt629 if "[NC]" not in prompt630 else prompt.replace("[NC]", "")631 )632 for prompt in clipped_prompts633 ]634 print(prompts)635 prompts = [636 prompt.rpartition("#")[0].strip() if "#" in prompt else prompt.strip()637 for prompt in prompts638 ]639 print(prompts)640 assert id_length <= len(641 prompts642 ), "id_length should not exceed total number of line-separated prompts"643 644 id_prompts = prompts[:id_length]645 real_prompts = prompts[id_length:]646 647 if seed is None:648 seed = int.from_bytes(os.urandom(2), "big")649 print(f"Using seed: {seed}")650 651 device = "cuda:0"652 setup_seed(seed)653 generator = torch.Generator(device=device).manual_seed(seed)654 655 torch.cuda.empty_cache()656 657 model_type = "original" if ref_image is None else "Photomaker"658 659 if model_type == "original":660 pipe = (661 self.sdxl_pipe_realvision662 if style_name == "(No style)"663 else self.sdxl_pipe_unstable664 )665 pipe = pipe.to(device)666 pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)667 else:668 if sd_model != "RealVision" and style_name != "(No style)":669 pipe = self.pipe_unstable.to(device)670 else:671 pipe = self.pipe_realvision.to(device)672 pipe.id_encoder.to(device)673 674 write = True675 cur_step = 0676 attn_count = 0677 678 set_attention_processor(pipe.unet, id_length, is_ipadapter=False)679 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)680 pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)681 curmodel_type = sd_model + "-" + model_type + "" + str(id_length)682 683 id_prompts, negative_prompt = apply_style(684 style_name, id_prompts, negative_prompt685 )686 687 total_results = []688 if model_type == "original":689 id_images = pipe(690 id_prompts,691 num_inference_steps=num_steps,692 guidance_scale=guidance_scale,693 height=height,694 width=width,695 negative_prompt=negative_prompt,696 generator=generator,697 ).images698 else:699 input_id_images = [load_image(str(ref_image))]700 start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)701 id_images = pipe(702 id_prompts,703 input_id_images=input_id_images,704 num_inference_steps=num_steps,705 guidance_scale=guidance_scale,706 start_merge_step=start_merge_step,707 height=height,708 width=width,709 negative_prompt=negative_prompt,710 generator=generator,711 ).images712 713 total_results = id_images + total_results714 715 real_images = []716 write = False717 for real_prompt in real_prompts:718 cur_step = 0719 real_prompt = apply_style_positive(style_name, real_prompt)720 if model_type == "original":721 real_images.append(722 pipe(723 real_prompt,724 num_inference_steps=num_steps,725 guidance_scale=guidance_scale,726 height=height,727 width=width,728 negative_prompt=negative_prompt,729 generator=generator,730 ).images[0]731 )732 else:733 real_images.append(734 pipe(735 real_prompt,736 input_id_images=input_id_images,737 num_inference_steps=num_steps,738 guidance_scale=guidance_scale,739 start_merge_step=start_merge_step,740 height=height,741 width=width,742 negative_prompt=negative_prompt,743 generator=generator,744 ).images[0]745 )746 747 total_results = [real_images[-1]] + total_results748 749 captions = clipped_prompts750 captions = [caption.replace("[NC]", "") for caption in captions]751 captions = [752 caption.split("#")[-1].strip() if "#" in caption else caption.strip()753 for caption in captions754 ]755 756 comic = get_comic(757 id_images + real_images,758 comic_style,759 captions=captions,760 font=ImageFont.truetype("./fonts/Inkfree.ttf", int(45)),761 )762 763 extension = output_format.lower()764 extension = "jpeg" if extension == "jpg" else extension765 comic_out = f"/tmp/comic.{extension}"766 comic[0].save(comic_out)767 768 save_params = {"format": extension.upper()}769 if not output_format == "png":770 save_params["quality"] = output_quality771 save_params["optimize"] = True772 773 output_paths = []774 for index, sample in enumerate(total_results[::-1]):775 output_filename = f"/tmp/out-{index}.{extension}"776 sample.save(output_filename, **save_params)777 output_paths.append(Path(output_filename))778 779 del pipe780 781 return ModelOutput(comic=Path(comic_out), individual_images=output_paths)782