182 lines
8.0 KiB
Python
182 lines
8.0 KiB
Python
import os
|
|
now_dir = os.path.dirname(os.path.abspath(__file__))
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import cuda_malloc
|
|
|
|
from huggingface_hub import snapshot_download
|
|
from diffusers import StableDiffusionPipeline,UNet2DConditionModel, ControlNetModel,StableDiffusionAdapterPipeline, T2IAdapter
|
|
|
|
import cv2
|
|
from .styleshot.annotator.hed import SOFT_HEDdetector
|
|
from .styleshot.annotator.lineart import LineartDetector
|
|
from .styleshot.ip_adapter.ip_adapter import StyleShot,StyleContentStableDiffusionControlNetPipeline,StableDiffusionControlNetPipeline
|
|
|
|
prtrained_dir = os.path.join(now_dir,"prtrained_models")
|
|
device = "cuda" if cuda_malloc.cuda_malloc_supported() else "cpu"
|
|
|
|
|
|
|
|
def load_weights(use_case="text_driven",preprocessor = "Contour",seed=42):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed(seed)
|
|
base_model_path = "runwayml/stable-diffusion-v1-5"
|
|
transformer_block_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
|
|
|
|
if preprocessor == "Lineart":
|
|
# detector = LineartDetector()
|
|
styleshot_model_path = "Gaojunyao/StyleShot_lineart"
|
|
elif preprocessor == "Contour":
|
|
# detector = SOFT_HEDdetector()
|
|
styleshot_model_path = "Gaojunyao/StyleShot"
|
|
else:
|
|
raise ValueError("Invalid preprocessor")
|
|
|
|
if not os.path.isdir(base_model_path):
|
|
base_model_path = snapshot_download(base_model_path,
|
|
allow_patterns=["*fp16.safetensors","*.json","*yaml","*.txt"],
|
|
local_dir=os.path.join(prtrained_dir,base_model_path.split("/")[-1]))
|
|
print(f"Downloaded model to {base_model_path}")
|
|
if not os.path.isdir(transformer_block_path):
|
|
transformer_block_path = snapshot_download(transformer_block_path,
|
|
# allow_patterns=["*.safetensors","*.json"],
|
|
ignore_patterns=["open_clip*","*.bin"],
|
|
local_dir=os.path.join(prtrained_dir,transformer_block_path.split("/")[-1]))
|
|
print(f"Downloaded model to {transformer_block_path}")
|
|
if not os.path.isdir(styleshot_model_path):
|
|
styleshot_model_path = snapshot_download(styleshot_model_path,
|
|
local_dir=os.path.join(prtrained_dir,styleshot_model_path.split("/")[-1]))
|
|
print(f"Downloaded model to {styleshot_model_path}")
|
|
|
|
ip_ckpt = os.path.join(styleshot_model_path, "pretrained_weight/ip.bin")
|
|
style_aware_encoder_path = os.path.join(styleshot_model_path, "pretrained_weight/style_aware_encoder.bin")
|
|
|
|
if use_case == "text_driven":
|
|
pipe = StableDiffusionPipeline.from_pretrained(base_model_path,variant="fp16")
|
|
styleshot = StyleShot(device, pipe, ip_ckpt, style_aware_encoder_path, transformer_block_path)
|
|
if use_case == "image_driven":
|
|
unet = UNet2DConditionModel.from_pretrained(base_model_path, subfolder="unet",variant="fp16")
|
|
content_fusion_encoder = ControlNetModel.from_unet(unet)
|
|
|
|
pipe = StyleContentStableDiffusionControlNetPipeline.from_pretrained(base_model_path,variant="fp16",controlnet=content_fusion_encoder)
|
|
styleshot = StyleShot(device, pipe, ip_ckpt, style_aware_encoder_path, transformer_block_path)
|
|
|
|
if use_case == "t2i-adapter":
|
|
adapter_model_path = "TencentARC/t2iadapter_depth_sd15v2"
|
|
if not os.path.isdir(adapter_model_path):
|
|
adapter_model_path = snapshot_download(adapter_model_path,
|
|
ignore_patterns=["*.png"],
|
|
local_dir=os.path.join(prtrained_dir,adapter_model_path.split("/")[-1]))
|
|
print(f"Downloaded model to {adapter_model_path}")
|
|
adapter = T2IAdapter.from_pretrained(adapter_model_path, torch_dtype=torch.float16)
|
|
pipe = StableDiffusionAdapterPipeline.from_pretrained(base_model_path, adapter=adapter,variant="fp16")
|
|
|
|
styleshot = StyleShot(device, pipe, ip_ckpt, style_aware_encoder_path, transformer_block_path)
|
|
|
|
if use_case == "controlnet":
|
|
controlnet_model_path = "lllyasviel/control_v11f1p_sd15_depth"
|
|
if not os.path.isdir(controlnet_model_path):
|
|
controlnet_model_path = snapshot_download(controlnet_model_path,
|
|
allow_patterns=["*.json","*.fp16.safetensors"],
|
|
local_dir=os.path.join(prtrained_dir,controlnet_model_path.split("/")[-1]))
|
|
print(f"Downloaded model to {controlnet_model_path}")
|
|
controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch.float16)
|
|
pipe = StableDiffusionControlNetPipeline.from_pretrained(base_model_path, controlnet=controlnet,variant="fp16")
|
|
|
|
styleshot = StyleShot(device, pipe, ip_ckpt, style_aware_encoder_path, transformer_block_path)
|
|
|
|
return styleshot
|
|
|
|
class TextNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"text": ("STRING", {"multiline": True, "dynamicPrompts": True})}}
|
|
RETURN_TYPES = ("TEXT",)
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "AIFSH_StyleShot"
|
|
|
|
def encode(self, text):
|
|
return (text, )
|
|
|
|
class StyleShotNode:
|
|
def __init__(self) -> None:
|
|
self.use_case = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":{
|
|
"style":("IMAGE",),
|
|
"prompt":("TEXT",),
|
|
"use_case":(["text_driven","image_driven","controlnet","t2i-adapter"],),
|
|
"seed":("INT",{
|
|
"default":42
|
|
})
|
|
},
|
|
"optional":{
|
|
"content":("IMAGE",),
|
|
"preprocessor":(["Contour", "Lineart"],),
|
|
"condition":("IMAGE",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
#RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "generate"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "AIFSH_StyleShot"
|
|
|
|
def comfyimage2Image(self,comfyimage):
|
|
comfyimage = comfyimage.numpy()[0] * 255
|
|
image_np = comfyimage.astype(np.uint8)
|
|
image = Image.fromarray(image_np)
|
|
return image
|
|
|
|
def generate(self,style,prompt,use_case,seed,
|
|
content=None,preprocessor="Contour",condition=None):
|
|
if self.use_case != use_case:
|
|
self.use_case = use_case
|
|
self.styleshot = load_weights(self.use_case,preprocessor,seed)
|
|
torch.cuda.empty_cache()
|
|
style_image = self.comfyimage2Image(style)
|
|
if self.use_case == "image_driven":
|
|
annotator_ckpts_path = os.path.join(prtrained_dir,"Annotators")
|
|
snapshot_download("lllyasviel/Annotators",
|
|
allow_patterns=["sk_model*","ControlNetHED*"],
|
|
local_dir=annotator_ckpts_path)
|
|
|
|
if preprocessor == "Lineart":
|
|
detector = LineartDetector(annotator_ckpts_path)
|
|
else:
|
|
detector = SOFT_HEDdetector(annotator_ckpts_path)
|
|
|
|
content_image = self.comfyimage2Image(content)
|
|
content_image = cv2.cvtColor(np.asarray(content_image), cv2.COLOR_BGR2RGB)
|
|
content_image = detector(content_image)
|
|
content_image = Image.fromarray(content_image)
|
|
|
|
else:
|
|
content_image = None
|
|
|
|
negative_prompt = "nsfw,bare,naked"
|
|
if use_case in ["t2i-adapter","controlnet"]:
|
|
condition_image = self.comfyimage2Image(condition)
|
|
generation = self.styleshot.generate(style_image=style_image, prompt=[[prompt]],negative_prompt=negative_prompt, image=[condition_image])
|
|
else:
|
|
generation = self.styleshot.generate(style_image=style_image, prompt=[[prompt]],negative_prompt=negative_prompt,content_image=content_image)
|
|
# generation[0][0].save("test.png")
|
|
out_image = torch.from_numpy(np.array(generation[0][0]) / 255.0).unsqueeze(0)
|
|
return (out_image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TextNode":TextNode,
|
|
"StyleShotNode": StyleShotNode
|
|
}
|
|
|