Spaces:
Runtime error
Runtime error
| from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline, StableDiffusionInstructPix2PixPipeline | |
| from diffusers import EulerAncestralDiscreteScheduler | |
| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler | |
| from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPSegProcessor, CLIPSegForImageSegmentation | |
| from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering | |
| from transformers import AutoImageProcessor, UperNetForSemanticSegmentation | |
| import os | |
| import random | |
| import torch | |
| import cv2 | |
| import re | |
| import uuid | |
| from PIL import Image, ImageOps, ImageDraw, ImageFont | |
| import numpy as np | |
| import math | |
| import inspect | |
| import tempfile | |
| from langchain.llms.openai import OpenAI | |
| # Grounding DINO | |
| import groundingdino.datasets.transforms as T | |
| from groundingdino.models import build_model | |
| from groundingdino.util import box_ops | |
| from groundingdino.util.slconfig import SLConfig | |
| from groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap | |
| # segment anything | |
| from segment_anything import build_sam, SamPredictor, SamAutomaticMaskGenerator | |
| import matplotlib.pyplot as plt | |
| import wget | |
| def prompts(name, description): | |
| def decorator(func): | |
| func.name = name | |
| func.description = description | |
| return func | |
| return decorator | |
| def blend_gt2pt(old_image, new_image, sigma=0.15, steps=100): | |
| new_size = new_image.size | |
| old_size = old_image.size | |
| easy_img = np.array(new_image) | |
| gt_img_array = np.array(old_image) | |
| pos_w = (new_size[0] - old_size[0]) // 2 | |
| pos_h = (new_size[1] - old_size[1]) // 2 | |
| kernel_h = cv2.getGaussianKernel(old_size[1], old_size[1] * sigma) | |
| kernel_w = cv2.getGaussianKernel(old_size[0], old_size[0] * sigma) | |
| kernel = np.multiply(kernel_h, np.transpose(kernel_w)) | |
| kernel[steps:-steps, steps:-steps] = 1 | |
| kernel[:steps, :steps] = kernel[:steps, :steps] / kernel[steps - 1, steps - 1] | |
| kernel[:steps, -steps:] = kernel[:steps, -steps:] / kernel[steps - 1, -(steps)] | |
| kernel[-steps:, :steps] = kernel[-steps:, :steps] / kernel[-steps, steps - 1] | |
| kernel[-steps:, -steps:] = kernel[-steps:, -steps:] / kernel[-steps, -steps] | |
| kernel = np.expand_dims(kernel, 2) | |
| kernel = np.repeat(kernel, 3, 2) | |
| weight = np.linspace(0, 1, steps) | |
| top = np.expand_dims(weight, 1) | |
| top = np.repeat(top, old_size[0] - 2 * steps, 1) | |
| top = np.expand_dims(top, 2) | |
| top = np.repeat(top, 3, 2) | |
| weight = np.linspace(1, 0, steps) | |
| down = np.expand_dims(weight, 1) | |
| down = np.repeat(down, old_size[0] - 2 * steps, 1) | |
| down = np.expand_dims(down, 2) | |
| down = np.repeat(down, 3, 2) | |
| weight = np.linspace(0, 1, steps) | |
| left = np.expand_dims(weight, 0) | |
| left = np.repeat(left, old_size[1] - 2 * steps, 0) | |
| left = np.expand_dims(left, 2) | |
| left = np.repeat(left, 3, 2) | |
| weight = np.linspace(1, 0, steps) | |
| right = np.expand_dims(weight, 0) | |
| right = np.repeat(right, old_size[1] - 2 * steps, 0) | |
| right = np.expand_dims(right, 2) | |
| right = np.repeat(right, 3, 2) | |
| kernel[:steps, steps:-steps] = top | |
| kernel[-steps:, steps:-steps] = down | |
| kernel[steps:-steps, :steps] = left | |
| kernel[steps:-steps, -steps:] = right | |
| pt_gt_img = easy_img[pos_h:pos_h + old_size[1], pos_w:pos_w + old_size[0]] | |
| gaussian_gt_img = kernel * gt_img_array + (1 - kernel) * pt_gt_img # gt img with blur img | |
| gaussian_gt_img = gaussian_gt_img.astype(np.int64) | |
| easy_img[pos_h:pos_h + old_size[1], pos_w:pos_w + old_size[0]] = gaussian_gt_img | |
| gaussian_img = Image.fromarray(easy_img) | |
| return gaussian_img | |
| def get_new_image_name(org_img_name, func_name="update"): | |
| head_tail = os.path.split(org_img_name) | |
| head = head_tail[0] | |
| tail = head_tail[1] | |
| name_split = tail.split('.')[0].split('_') | |
| this_new_uuid = str(uuid.uuid4())[0:4] | |
| if len(name_split) == 1: | |
| most_org_file_name = name_split[0] | |
| recent_prev_file_name = name_split[0] | |
| new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name) | |
| else: | |
| assert len(name_split) == 4 | |
| most_org_file_name = name_split[3] | |
| recent_prev_file_name = name_split[0] | |
| new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name) | |
| return os.path.join(head, new_file_name) | |
| def seed_everything(seed): | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed_all(seed) | |
| return seed | |
| class InstructPix2Pix: | |
| def __init__(self, device): | |
| print(f"Initializing InstructPix2Pix to {device}") | |
| self.device = device | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", | |
| safety_checker=None, | |
| torch_dtype=self.torch_dtype).to(device) | |
| self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) | |
| def inference(self, inputs): | |
| """Change style of image.""" | |
| print("===>Starting InstructPix2Pix Inference") | |
| image_path, text = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) | |
| original_image = Image.open(image_path) | |
| image = self.pipe(text, image=original_image, num_inference_steps=40, image_guidance_scale=1.2).images[0] | |
| updated_image_path = get_new_image_name(image_path, func_name="pix2pix") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed InstructPix2Pix, Input Image: {image_path}, Instruct Text: {text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Text2Image: | |
| def __init__(self, device): | |
| print(f"Initializing Text2Image to {device}") | |
| self.device = device | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe.to(device) | |
| 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' | |
| def inference(self, text): | |
| image_filename = os.path.join('image', f"{str(uuid.uuid4())[:8]}.png") | |
| prompt = text + ', ' + self.a_prompt | |
| image = self.pipe(prompt, negative_prompt=self.n_prompt).images[0] | |
| image.save(image_filename) | |
| print( | |
| f"\nProcessed Text2Image, Input Text: {text}, Output Image: {image_filename}") | |
| return image_filename | |
| class ImageCaptioning: | |
| def __init__(self, device): | |
| print(f"Initializing ImageCaptioning to {device}") | |
| self.device = device | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
| self.model = BlipForConditionalGeneration.from_pretrained( | |
| "Salesforce/blip-image-captioning-base", torch_dtype=self.torch_dtype).to(self.device) | |
| def inference(self, image_path): | |
| inputs = self.processor(Image.open(image_path), return_tensors="pt").to(self.device, self.torch_dtype) | |
| out = self.model.generate(**inputs) | |
| captions = self.processor.decode(out[0], skip_special_tokens=True) | |
| print(f"\nProcessed ImageCaptioning, Input Image: {image_path}, Output Text: {captions}") | |
| return captions | |
| class Image2Canny: | |
| def __init__(self, device): | |
| print("Initializing Image2Canny") | |
| self.low_threshold = 100 | |
| self.high_threshold = 200 | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| image = np.array(image) | |
| canny = cv2.Canny(image, self.low_threshold, self.high_threshold) | |
| canny = canny[:, :, None] | |
| canny = np.concatenate([canny, canny, canny], axis=2) | |
| canny = Image.fromarray(canny) | |
| updated_image_path = get_new_image_name(inputs, func_name="edge") | |
| canny.save(updated_image_path) | |
| print(f"\nProcessed Image2Canny, Input Image: {inputs}, Output Text: {updated_image_path}") | |
| return updated_image_path | |
| class CannyText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing CannyText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-canny", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="canny2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed CannyText2Image, Input Canny: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Text: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Line: | |
| def __init__(self, device): | |
| print("Initializing Image2Line") | |
| self.detector = MLSDdetector.from_pretrained('lllyasviel/ControlNet') | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| mlsd = self.detector(image) | |
| updated_image_path = get_new_image_name(inputs, func_name="line-of") | |
| mlsd.save(updated_image_path) | |
| print(f"\nProcessed Image2Line, Input Image: {inputs}, Output Line: {updated_image_path}") | |
| return updated_image_path | |
| class LineText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing LineText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-mlsd", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="line2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed LineText2Image, Input Line: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Text: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Hed: | |
| def __init__(self, device): | |
| print("Initializing Image2Hed") | |
| self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet') | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| hed = self.detector(image) | |
| updated_image_path = get_new_image_name(inputs, func_name="hed-boundary") | |
| hed.save(updated_image_path) | |
| print(f"\nProcessed Image2Hed, Input Image: {inputs}, Output Hed: {updated_image_path}") | |
| return updated_image_path | |
| class HedText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing HedText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-hed", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="hed2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed HedText2Image, Input Hed: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Scribble: | |
| def __init__(self, device): | |
| print("Initializing Image2Scribble") | |
| self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet') | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| scribble = self.detector(image, scribble=True) | |
| updated_image_path = get_new_image_name(inputs, func_name="scribble") | |
| scribble.save(updated_image_path) | |
| print(f"\nProcessed Image2Scribble, Input Image: {inputs}, Output Scribble: {updated_image_path}") | |
| return updated_image_path | |
| class ScribbleText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing ScribbleText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-scribble", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="scribble2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed ScribbleText2Image, Input Scribble: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Pose: | |
| def __init__(self, device): | |
| print("Initializing Image2Pose") | |
| self.detector = OpenposeDetector.from_pretrained('lllyasviel/ControlNet') | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| pose = self.detector(image) | |
| updated_image_path = get_new_image_name(inputs, func_name="human-pose") | |
| pose.save(updated_image_path) | |
| print(f"\nProcessed Image2Pose, Input Image: {inputs}, Output Pose: {updated_image_path}") | |
| return updated_image_path | |
| class PoseText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing PoseText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-openpose", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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.num_inference_steps = 20 | |
| self.seed = -1 | |
| self.unconditional_guidance_scale = 9.0 | |
| 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' | |
| 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 = f'{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="pose2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed PoseText2Image, Input Pose: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class SegText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing SegText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-seg", | |
| torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="segment2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed SegText2Image, Input Seg: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Depth: | |
| def __init__(self, device): | |
| print("Initializing Image2Depth") | |
| self.depth_estimator = pipeline('depth-estimation') | |
| def inference(self, inputs): | |
| image = Image.open(inputs) | |
| depth = self.depth_estimator(image)['depth'] | |
| depth = np.array(depth) | |
| depth = depth[:, :, None] | |
| depth = np.concatenate([depth, depth, depth], axis=2) | |
| depth = Image.fromarray(depth) | |
| updated_image_path = get_new_image_name(inputs, func_name="depth") | |
| depth.save(updated_image_path) | |
| print(f"\nProcessed Image2Depth, Input Image: {inputs}, Output Depth: {updated_image_path}") | |
| return updated_image_path | |
| class DepthText2Image: | |
| def __init__(self, device): | |
| print(f"Initializing DepthText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained( | |
| "fusing/stable-diffusion-v1-5-controlnet-depth", torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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="depth2image") | |
| image.save(updated_image_path) | |
| print(f"\nProcessed DepthText2Image, Input Depth: {image_path}, Input Text: {instruct_text}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Image2Normal: | |
| def __init__(self, device): | |
| print("Initializing Image2Normal") | |
| self.depth_estimator = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas") | |
| self.bg_threhold = 0.4 | |
| 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): | |
| print(f"Initializing NormalText2Image to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.controlnet = ControlNetModel.from_pretrained( | |
| "fusing/stable-diffusion-v1-5-controlnet-normal", torch_dtype=self.torch_dtype) | |
| self.pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/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' | |
| 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 = f'{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 | |
| class VisualQuestionAnswering: | |
| def __init__(self, device): | |
| print(f"Initializing VisualQuestionAnswering to {device}") | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.device = device | |
| self.processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base") | |
| self.model = BlipForQuestionAnswering.from_pretrained( | |
| "Salesforce/blip-vqa-base", torch_dtype=self.torch_dtype).to(self.device) | |
| def inference(self, inputs): | |
| image_path, question = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) | |
| raw_image = Image.open(image_path).convert('RGB') | |
| inputs = self.processor(raw_image, question, return_tensors="pt").to(self.device, self.torch_dtype) | |
| out = self.model.generate(**inputs) | |
| answer = self.processor.decode(out[0], skip_special_tokens=True) | |
| print(f"\nProcessed VisualQuestionAnswering, Input Image: {image_path}, Input Question: {question}, " | |
| f"Output Answer: {answer}") | |
| return answer | |
| class Segmenting: | |
| def __init__(self, device): | |
| print(f"Inintializing Segmentation to {device}") | |
| self.device = device | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.model_checkpoint_path = os.path.join("checkpoints", "sam") | |
| self.download_parameters() | |
| self.sam = build_sam(checkpoint=self.model_checkpoint_path).to(device) | |
| self.sam_predictor = SamPredictor(self.sam) | |
| self.mask_generator = SamAutomaticMaskGenerator(self.sam) | |
| def download_parameters(self): | |
| url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" | |
| if not os.path.exists(self.model_checkpoint_path): | |
| wget.download(url, out=self.model_checkpoint_path) | |
| def show_mask(self, mask, ax, random_color=False): | |
| if random_color: | |
| color = np.concatenate([np.random.random(3), np.array([1])], axis=0) | |
| else: | |
| color = np.array([30 / 255, 144 / 255, 255 / 255, 1]) | |
| h, w = mask.shape[-2:] | |
| mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) | |
| ax.imshow(mask_image) | |
| def show_box(self, box, ax, label): | |
| x0, y0 = box[0], box[1] | |
| w, h = box[2] - box[0], box[3] - box[1] | |
| ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2)) | |
| ax.text(x0, y0, label) | |
| def get_mask_with_boxes(self, image_pil, image, boxes_filt): | |
| size = image_pil.size | |
| H, W = size[1], size[0] | |
| for i in range(boxes_filt.size(0)): | |
| boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H]) | |
| boxes_filt[i][:2] -= boxes_filt[i][2:] / 2 | |
| boxes_filt[i][2:] += boxes_filt[i][:2] | |
| boxes_filt = boxes_filt.cpu() | |
| transformed_boxes = self.sam_predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(self.device) | |
| masks, _, _ = self.sam_predictor.predict_torch( | |
| point_coords=None, | |
| point_labels=None, | |
| boxes=transformed_boxes.to(self.device), | |
| multimask_output=False, | |
| ) | |
| return masks | |
| def segment_image_with_boxes(self, image_pil, image_path, boxes_filt, pred_phrases): | |
| image = cv2.imread(image_path) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| self.sam_predictor.set_image(image) | |
| masks = self.get_mask_with_boxes(image_pil, image, boxes_filt) | |
| # draw output image | |
| plt.figure(figsize=(10, 10)) | |
| plt.imshow(image) | |
| for mask in masks: | |
| self.show_mask(mask.cpu().numpy(), plt.gca(), random_color=True) | |
| updated_image_path = get_new_image_name(image_path, func_name="segmentation") | |
| plt.axis('off') | |
| plt.savefig( | |
| updated_image_path, | |
| bbox_inches="tight", dpi=300, pad_inches=0.0 | |
| ) | |
| return updated_image_path | |
| def inference_all(self, image_path): | |
| image = cv2.imread(image_path) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| masks = self.mask_generator.generate(image) | |
| plt.figure(figsize=(20, 20)) | |
| plt.imshow(image) | |
| if len(masks) == 0: | |
| return | |
| sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True) | |
| ax = plt.gca() | |
| ax.set_autoscale_on(False) | |
| polygons = [] | |
| color = [] | |
| for ann in sorted_anns: | |
| m = ann['segmentation'] | |
| img = np.ones((m.shape[0], m.shape[1], 3)) | |
| color_mask = np.random.random((1, 3)).tolist()[0] | |
| for i in range(3): | |
| img[:, :, i] = color_mask[i] | |
| ax.imshow(np.dstack((img, m))) | |
| updated_image_path = get_new_image_name(image_path, func_name="segment-image") | |
| plt.axis('off') | |
| plt.savefig( | |
| updated_image_path, | |
| bbox_inches="tight", dpi=300, pad_inches=0.0 | |
| ) | |
| return updated_image_path | |
| class Text2Box: | |
| def __init__(self, device): | |
| print(f"Initializing ObjectDetection to {device}") | |
| self.device = device | |
| self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 | |
| self.model_checkpoint_path = os.path.join("checkpoints", "groundingdino") | |
| self.model_config_path = os.path.join("checkpoints", "grounding_config.py") | |
| self.download_parameters() | |
| self.box_threshold = 0.3 | |
| self.text_threshold = 0.25 | |
| self.grounding = (self.load_model()).to(self.device) | |
| def download_parameters(self): | |
| url = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth" | |
| if not os.path.exists(self.model_checkpoint_path): | |
| wget.download(url, out=self.model_checkpoint_path) | |
| config_url = "https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/main/groundingdino/config/GroundingDINO_SwinT_OGC.py" | |
| if not os.path.exists(self.model_config_path): | |
| wget.download(config_url, out=self.model_config_path) | |
| def load_image(self, image_path): | |
| # load image | |
| image_pil = Image.open(image_path).convert("RGB") # load image | |
| transform = T.Compose( | |
| [ | |
| T.RandomResize([512], max_size=1333), | |
| T.ToTensor(), | |
| T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ] | |
| ) | |
| image, _ = transform(image_pil, None) # 3, h, w | |
| return image_pil, image | |
| def load_model(self): | |
| args = SLConfig.fromfile(self.model_config_path) | |
| args.device = self.device | |
| model = build_model(args) | |
| checkpoint = torch.load(self.model_checkpoint_path, map_location="cpu") | |
| load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) | |
| print(load_res) | |
| _ = model.eval() | |
| return model | |
| def get_grounding_boxes(self, image, caption, with_logits=True): | |
| caption = caption.lower() | |
| caption = caption.strip() | |
| if not caption.endswith("."): | |
| caption = caption + "." | |
| image = image.to(self.device) | |
| with torch.no_grad(): | |
| outputs = self.grounding(image[None], captions=[caption]) | |
| logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256) | |
| boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4) | |
| logits.shape[0] | |
| # filter output | |
| logits_filt = logits.clone() | |
| boxes_filt = boxes.clone() | |
| filt_mask = logits_filt.max(dim=1)[0] > self.box_threshold | |
| logits_filt = logits_filt[filt_mask] # num_filt, 256 | |
| boxes_filt = boxes_filt[filt_mask] # num_filt, 4 | |
| logits_filt.shape[0] | |
| # get phrase | |
| tokenlizer = self.grounding.tokenizer | |
| tokenized = tokenlizer(caption) | |
| # build pred | |
| pred_phrases = [] | |
| for logit, box in zip(logits_filt, boxes_filt): | |
| pred_phrase = get_phrases_from_posmap(logit > self.text_threshold, tokenized, tokenlizer) | |
| if with_logits: | |
| pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})") | |
| else: | |
| pred_phrases.append(pred_phrase) | |
| return boxes_filt, pred_phrases | |
| def plot_boxes_to_image(self, image_pil, tgt): | |
| H, W = tgt["size"] | |
| boxes = tgt["boxes"] | |
| labels = tgt["labels"] | |
| assert len(boxes) == len(labels), "boxes and labels must have same length" | |
| draw = ImageDraw.Draw(image_pil) | |
| mask = Image.new("L", image_pil.size, 0) | |
| mask_draw = ImageDraw.Draw(mask) | |
| # draw boxes and masks | |
| for box, label in zip(boxes, labels): | |
| # from 0..1 to 0..W, 0..H | |
| box = box * torch.Tensor([W, H, W, H]) | |
| # from xywh to xyxy | |
| box[:2] -= box[2:] / 2 | |
| box[2:] += box[:2] | |
| # random color | |
| color = tuple(np.random.randint(0, 255, size=3).tolist()) | |
| # draw | |
| x0, y0, x1, y1 = box | |
| x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1) | |
| draw.rectangle([x0, y0, x1, y1], outline=color, width=6) | |
| # draw.text((x0, y0), str(label), fill=color) | |
| font = ImageFont.load_default() | |
| if hasattr(font, "getbbox"): | |
| bbox = draw.textbbox((x0, y0), str(label), font) | |
| else: | |
| w, h = draw.textsize(str(label), font) | |
| bbox = (x0, y0, w + x0, y0 + h) | |
| # bbox = draw.textbbox((x0, y0), str(label)) | |
| draw.rectangle(bbox, fill=color) | |
| draw.text((x0, y0), str(label), fill="white") | |
| mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=2) | |
| return image_pil, mask | |
| def inference(self, inputs): | |
| image_path, det_prompt = inputs.split(",") | |
| print(f"image_path={image_path}, text_prompt={det_prompt}") | |
| image_pil, image = self.load_image(image_path) | |
| boxes_filt, pred_phrases = self.get_grounding_boxes(image, det_prompt) | |
| size = image_pil.size | |
| pred_dict = { | |
| "boxes": boxes_filt, | |
| "size": [size[1], size[0]], # H,W | |
| "labels": pred_phrases, } | |
| image_with_box = self.plot_boxes_to_image(image_pil, pred_dict)[0] | |
| updated_image_path = get_new_image_name(image_path, func_name="detect-something") | |
| updated_image = image_with_box.resize(size) | |
| updated_image.save(updated_image_path) | |
| print( | |
| f"\nProcessed ObejectDetecting, Input Image: {image_path}, Object to be Detect {det_prompt}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class Inpainting: | |
| def __init__(self, device): | |
| self.device = device | |
| self.revision = 'fp16' if 'cuda' in self.device else None | |
| self.torch_dtype = torch.float16 if 'cuda' in self.device else torch.float32 | |
| self.inpaint = StableDiffusionInpaintPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-inpainting", revision=self.revision, torch_dtype=self.torch_dtype).to(device) | |
| def __call__(self, prompt, image, mask_image, height=512, width=512, num_inference_steps=50): | |
| update_image = self.inpaint(prompt=prompt, image=image.resize((width, height)), | |
| mask_image=mask_image.resize((width, height)), height=height, width=width, | |
| num_inference_steps=num_inference_steps).images[0] | |
| return update_image | |
| class InfinityOutPainting: | |
| template_model = True # Add this line to show this is a template model. | |
| def __init__(self, ImageCaptioning, Inpainting, VisualQuestionAnswering): | |
| self.ImageCaption = ImageCaptioning | |
| self.inpaint = Inpainting | |
| self.ImageVQA = VisualQuestionAnswering | |
| 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' | |
| def get_BLIP_vqa(self, image, question): | |
| inputs = self.ImageVQA.processor(image, question, return_tensors="pt").to(self.ImageVQA.device, | |
| self.ImageVQA.torch_dtype) | |
| out = self.ImageVQA.model.generate(**inputs) | |
| answer = self.ImageVQA.processor.decode(out[0], skip_special_tokens=True) | |
| print(f"\nProcessed VisualQuestionAnswering, Input Question: {question}, Output Answer: {answer}") | |
| return answer | |
| def get_BLIP_caption(self, image): | |
| inputs = self.ImageCaption.processor(image, return_tensors="pt").to(self.ImageCaption.device, | |
| self.ImageCaption.torch_dtype) | |
| out = self.ImageCaption.model.generate(**inputs) | |
| BLIP_caption = self.ImageCaption.processor.decode(out[0], skip_special_tokens=True) | |
| return BLIP_caption | |
| def get_imagine_caption(self, image, imagine): | |
| BLIP_caption = self.get_BLIP_caption(image) | |
| caption = BLIP_caption | |
| print(f'Prompt: {caption}') | |
| return caption | |
| def resize_image(self, image, max_size=1000000, multiple=8): | |
| aspect_ratio = image.size[0] / image.size[1] | |
| new_width = int(math.sqrt(max_size * aspect_ratio)) | |
| new_height = int(new_width / aspect_ratio) | |
| new_width, new_height = new_width - (new_width % multiple), new_height - (new_height % multiple) | |
| return image.resize((new_width, new_height)) | |
| def dowhile(self, original_img, tosize, expand_ratio, imagine, usr_prompt): | |
| old_img = original_img | |
| while (old_img.size != tosize): | |
| prompt = self.check_prompt(usr_prompt) if usr_prompt else self.get_imagine_caption(old_img, imagine) | |
| crop_w = 15 if old_img.size[0] != tosize[0] else 0 | |
| crop_h = 15 if old_img.size[1] != tosize[1] else 0 | |
| old_img = ImageOps.crop(old_img, (crop_w, crop_h, crop_w, crop_h)) | |
| temp_canvas_size = (expand_ratio * old_img.width if expand_ratio * old_img.width < tosize[0] else tosize[0], | |
| expand_ratio * old_img.height if expand_ratio * old_img.height < tosize[1] else tosize[ | |
| 1]) | |
| temp_canvas, temp_mask = Image.new("RGB", temp_canvas_size, color="white"), Image.new("L", temp_canvas_size, | |
| color="white") | |
| x, y = (temp_canvas.width - old_img.width) // 2, (temp_canvas.height - old_img.height) // 2 | |
| temp_canvas.paste(old_img, (x, y)) | |
| temp_mask.paste(0, (x, y, x + old_img.width, y + old_img.height)) | |
| resized_temp_canvas, resized_temp_mask = self.resize_image(temp_canvas), self.resize_image(temp_mask) | |
| image = self.inpaint(prompt=prompt, image=resized_temp_canvas, mask_image=resized_temp_mask, | |
| height=resized_temp_canvas.height, width=resized_temp_canvas.width, | |
| num_inference_steps=50).resize( | |
| (temp_canvas.width, temp_canvas.height), Image.ANTIALIAS) | |
| image = blend_gt2pt(old_img, image) | |
| old_img = image | |
| return old_img | |
| def inference(self, inputs): | |
| image_path, resolution = inputs.split(',') | |
| width, height = resolution.split('x') | |
| tosize = (int(width), int(height)) | |
| image = Image.open(image_path) | |
| image = ImageOps.crop(image, (10, 10, 10, 10)) | |
| out_painted_image = self.dowhile(image, tosize, 4, True, False) | |
| updated_image_path = get_new_image_name(image_path, func_name="outpainting") | |
| out_painted_image.save(updated_image_path) | |
| print(f"\nProcessed InfinityOutPainting, Input Image: {image_path}, Input Resolution: {resolution}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class ObjectSegmenting: | |
| template_model = True # Add this line to show this is a template model. | |
| def __init__(self, Text2Box: Text2Box, Segmenting: Segmenting): | |
| # self.llm = OpenAI(temperature=0) | |
| self.grounding = Text2Box | |
| self.sam = Segmenting | |
| def inference(self, inputs): | |
| image_path, det_prompt = inputs.split(",") | |
| print(f"image_path={image_path}, text_prompt={det_prompt}") | |
| image_pil, image = self.grounding.load_image(image_path) | |
| boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, det_prompt) | |
| updated_image_path = self.sam.segment_image_with_boxes(image_pil, image_path, boxes_filt, pred_phrases) | |
| print( | |
| f"\nProcessed ObejectSegmenting, Input Image: {image_path}, Object to be Segment {det_prompt}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path | |
| class ImageEditing: | |
| template_model = True | |
| def __init__(self, Text2Box: Text2Box, Segmenting: Segmenting, Inpainting: Inpainting): | |
| print(f"Initializing ImageEditing") | |
| self.sam = Segmenting | |
| self.grounding = Text2Box | |
| self.inpaint = Inpainting | |
| def pad_edge(self, mask, padding): | |
| # mask Tensor [H,W] | |
| mask = mask.numpy() | |
| true_indices = np.argwhere(mask) | |
| mask_array = np.zeros_like(mask, dtype=bool) | |
| for idx in true_indices: | |
| padded_slice = tuple(slice(max(0, i - padding), i + padding + 1) for i in idx) | |
| mask_array[padded_slice] = True | |
| new_mask = (mask_array * 255).astype(np.uint8) | |
| # new_mask | |
| return new_mask | |
| def inference_remove(self, inputs): | |
| image_path, to_be_removed_txt = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) | |
| return self.inference_replace_sam(f"{image_path},{to_be_removed_txt},background") | |
| def inference_replace_sam(self, inputs): | |
| image_path, to_be_replaced_txt, replace_with_txt = inputs.split(",") | |
| print(f"image_path={image_path}, to_be_replaced_txt={to_be_replaced_txt}") | |
| image_pil, image = self.grounding.load_image(image_path) | |
| boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, to_be_replaced_txt) | |
| image = cv2.imread(image_path) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| self.sam.sam_predictor.set_image(image) | |
| masks = self.sam.get_mask_with_boxes(image_pil, image, boxes_filt) | |
| mask = torch.sum(masks, dim=0).unsqueeze(0) | |
| mask = torch.where(mask > 0, True, False) | |
| mask = mask.squeeze(0).squeeze(0).cpu() # tensor | |
| mask = self.pad_edge(mask, padding=20) # numpy | |
| mask_image = Image.fromarray(mask) | |
| updated_image = self.inpaint(prompt=replace_with_txt, image=image_pil, | |
| mask_image=mask_image) | |
| updated_image_path = get_new_image_name(image_path, func_name="replace-something") | |
| updated_image = updated_image.resize(image_pil.size) | |
| updated_image.save(updated_image_path) | |
| print( | |
| f"\nProcessed ImageEditing, Input Image: {image_path}, Replace {to_be_replaced_txt} to {replace_with_txt}, " | |
| f"Output Image: {updated_image_path}") | |
| return updated_image_path |