|
from modules.utils import * |
|
|
|
class Image2Normal: |
|
def __init__(self, device, pretrained_model_dir): |
|
print("Initializing Image2Normal") |
|
self.depth_estimator = pipeline("depth-estimation", model=f"{pretrained_model_dir}/dpt-hybrid-midas") |
|
self.bg_threhold = 0.4 |
|
|
|
@prompts(name="Predict Normal Map On Image", |
|
description="useful when you want to detect norm map of the image. " |
|
"like: generate normal map from this image, or predict normal map of this image. " |
|
"The input to this tool should be a string, representing the image_path") |
|
def inference(self, inputs): |
|
image = Image.open(inputs) |
|
original_size = image.size |
|
image = self.depth_estimator(image)['predicted_depth'][0] |
|
image = image.numpy() |
|
image_depth = image.copy() |
|
image_depth -= np.min(image_depth) |
|
image_depth /= np.max(image_depth) |
|
x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3) |
|
x[image_depth < self.bg_threhold] = 0 |
|
y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3) |
|
y[image_depth < self.bg_threhold] = 0 |
|
z = np.ones_like(x) * np.pi * 2.0 |
|
image = np.stack([x, y, z], axis=2) |
|
image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5 |
|
image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8) |
|
image = Image.fromarray(image) |
|
image = image.resize(original_size) |
|
updated_image_path = get_new_image_name(inputs, func_name="normal-map") |
|
image.save(updated_image_path) |
|
print(f"\nProcessed Image2Normal, Input Image: {inputs}, Output Depth: {updated_image_path}") |
|
return updated_image_path |
|
|
|
|
|
class NormalText2Image: |
|
def __init__(self, device, pretrained_model_dir): |
|
print("Initializing NormalText2Image to %s" % device) |
|
self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
|
self.controlnet = ControlNetModel.from_pretrained( |
|
f"{pretrained_model_dir}/sd-controlnet-normal", torch_dtype=self.torch_dtype) |
|
self.pipe = StableDiffusionControlNetPipeline.from_pretrained( |
|
f"{pretrained_model_dir}/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, |
|
torch_dtype=self.torch_dtype) |
|
self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config) |
|
self.pipe.to(device) |
|
self.seed = -1 |
|
self.a_prompt = 'best quality, extremely detailed' |
|
self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \ |
|
' fewer digits, cropped, worst quality, low quality' |
|
|
|
@prompts(name="Generate Image Condition On Normal Map", |
|
description="useful when you want to generate a new real image from both the user desciption and normal map. " |
|
"like: generate a real image of a object or something from this normal map, " |
|
"or generate a new real image of a object or something from the normal map. " |
|
"The input to this tool should be a comma seperated string of two, " |
|
"representing the image_path and the user description") |
|
def inference(self, inputs): |
|
image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) |
|
image = Image.open(image_path) |
|
self.seed = random.randint(0, 65535) |
|
seed_everything(self.seed) |
|
prompt = instruct_text + ', ' + self.a_prompt |
|
image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, |
|
guidance_scale=9.0).images[0] |
|
updated_image_path = get_new_image_name(image_path, func_name="normal2image") |
|
image.save(updated_image_path) |
|
print(f"\nProcessed NormalText2Image, Input Normal: {image_path}, Input Text: {instruct_text}, " |
|
f"Output Image: {updated_image_path}") |
|
return updated_image_path |
