Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import torch | |
| torch.jit.script = lambda f: f | |
| from download import OMG_download | |
| download = OMG_download() | |
| import sys | |
| sys.path.append('./') | |
| import argparse | |
| import hashlib | |
| import json | |
| import os.path | |
| import numpy as np | |
| import torch | |
| from typing import Tuple, List | |
| from diffusers import DPMSolverMultistepScheduler | |
| from diffusers.models import T2IAdapter | |
| from PIL import Image | |
| import copy | |
| from diffusers import ControlNetModel, StableDiffusionXLPipeline | |
| try: | |
| from insightface.app import FaceAnalysis | |
| except: | |
| print("insightface can not be load") | |
| import gradio as gr | |
| import random | |
| from PIL import Image, ImageOps | |
| from transformers import DPTFeatureExtractor, DPTForDepthEstimation | |
| from controlnet_aux import OpenposeDetector | |
| from controlnet_aux.open_pose.body import Body | |
| try: | |
| from inference.models import YOLOWorld | |
| from src.efficientvit.models.efficientvit.sam import EfficientViTSamPredictor | |
| from src.efficientvit.sam_model_zoo import create_sam_model | |
| import supervision as sv | |
| except: | |
| print("YoloWorld can not be load") | |
| try: | |
| 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 | |
| from groundingdino.util.inference import annotate, predict | |
| from segment_anything import build_sam, SamPredictor | |
| import groundingdino.datasets.transforms as T | |
| except: | |
| print("groundingdino can not be load") | |
| from src.pipelines.instantid_pipeline import InstantidMultiConceptPipeline | |
| from src.pipelines.instantid_single_pieline import InstantidSingleConceptPipeline | |
| from src.prompt_attention.p2p_attention import AttentionReplace | |
| from src.pipelines.instantid_pipeline import revise_regionally_controlnet_forward | |
| import cv2 | |
| import math | |
| import PIL.Image | |
| from gradio_demo.character_template import styles, lorapath_styles | |
| STYLE_NAMES = list(styles.keys()) | |
| MAX_SEED = np.iinfo(np.int32).max | |
| title = r""" | |
| <h1 align="center"> OMG + InstantID </h1> | |
| """ | |
| description = r""" | |
| <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/kongzhecn/OMG/' target='_blank'><b>OMG: Occlusion-friendly Personalized Multi-concept Generation In Diffusion Models</b></a>.<be><br> | |
| <br> | |
| <a href='https://kongzhecn.github.io/omg-project/' target='_blank'><b>[Project]</b></a><a href='https://github.com/kongzhecn/OMG/' target='_blank'><b>[Code]</b></a><a href='https://arxiv.org/abs/2403.10983/' target='_blank'><b>[Arxiv]</b></a><br> | |
| <br> | |
| ❗️<b>Related demos<b>:<a href='https://huggingface.co/spaces/Fucius/OMG/' target='_blank'><b> OMG + LoRAs </b></a>❗️<br> | |
| <br> | |
| How to use:<br> | |
| 1. Upload two character images of a man and a woman. | |
| 2. Enter a text prompt as done in normal text-to-image models. | |
| 3. Click the <b>Submit</b> button to start customizing. | |
| 4. Enjoy the generated image😊! | |
| """ | |
| article = r""" | |
| --- | |
| 📝 **Citation** | |
| <br> | |
| If our work is helpful for your research or applications, please cite us via: | |
| ```bibtex | |
| @article{kong2024omg, | |
| title={OMG: Occlusion-friendly Personalized Multi-concept Generation in Diffusion Models}, | |
| author={Kong, Zhe and Zhang, Yong and Yang, Tianyu and Wang, Tao and Zhang, Kaihao and Wu, Bizhu and Chen, Guanying and Liu, Wei and Luo, Wenhan}, | |
| journal={arXiv preprint arXiv:2403.10983}, | |
| year={2024} | |
| } | |
| ``` | |
| """ | |
| tips = r""" | |
| ### Usage tips of OMG | |
| 1. Input text prompts to describe a man and a woman | |
| """ | |
| css = ''' | |
| .gradio-container {width: 85% !important} | |
| ''' | |
| def build_dino_segment_model(ckpt_repo_id, sam_checkpoint): | |
| ckpt_filenmae = "groundingdino_swinb_cogcoor.pth" | |
| ckpt_config_filename = os.path.join(ckpt_repo_id, "GroundingDINO_SwinB.cfg.py") | |
| groundingdino_model = load_model_hf(ckpt_repo_id, ckpt_filenmae, ckpt_config_filename) | |
| sam = build_sam(checkpoint=sam_checkpoint) | |
| sam.cuda() | |
| sam_predictor = SamPredictor(sam) | |
| return groundingdino_model, sam_predictor | |
| def load_model_hf(repo_id, filename, ckpt_config_filename, device='cpu'): | |
| args = SLConfig.fromfile(ckpt_config_filename) | |
| model = build_model(args) | |
| args.device = device | |
| checkpoint = torch.load(os.path.join(repo_id, filename), map_location='cpu') | |
| log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False) | |
| print("Model loaded from {} \n => {}".format(filename, log)) | |
| _ = model.eval() | |
| return model | |
| def build_yolo_segment_model(sam_path, device): | |
| yolo_world = YOLOWorld(model_id="yolo_world/l") | |
| sam = EfficientViTSamPredictor( | |
| create_sam_model(name="xl1", weight_url=sam_path).to(device).eval() | |
| ) | |
| return yolo_world, sam | |
| def sample_image(pipe, | |
| input_prompt, | |
| input_neg_prompt=None, | |
| generator=None, | |
| concept_models=None, | |
| num_inference_steps=50, | |
| guidance_scale=3.0, | |
| controller=None, | |
| face_app=None, | |
| image=None, | |
| stage=None, | |
| region_masks=None, | |
| controlnet_conditioning_scale=None, | |
| **extra_kargs | |
| ): | |
| if image is not None: | |
| image_condition = [image] | |
| else: | |
| image_condition = None | |
| images = pipe( | |
| prompt=input_prompt, | |
| concept_models=concept_models, | |
| negative_prompt=input_neg_prompt, | |
| generator=generator, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, | |
| cross_attention_kwargs={"scale": 0.8}, | |
| controller=controller, | |
| image=image_condition, | |
| face_app=face_app, | |
| stage=stage, | |
| controlnet_conditioning_scale = controlnet_conditioning_scale, | |
| region_masks=region_masks, | |
| **extra_kargs).images | |
| return images | |
| def load_image_yoloworld(image_source) -> Tuple[np.array, torch.Tensor]: | |
| image = np.asarray(image_source) | |
| return image | |
| def load_image_dino(image_source) -> Tuple[np.array, torch.Tensor]: | |
| transform = T.Compose( | |
| [ | |
| T.RandomResize([800], max_size=1333), | |
| T.ToTensor(), | |
| T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ] | |
| ) | |
| image = np.asarray(image_source) | |
| image_transformed, _ = transform(image_source, None) | |
| return image, image_transformed | |
| def draw_kps_multi(image_pil, kps_list, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]): | |
| stickwidth = 4 | |
| limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) | |
| w, h = image_pil.size | |
| out_img = np.zeros([h, w, 3]) | |
| for kps in kps_list: | |
| kps = np.array(kps) | |
| for i in range(len(limbSeq)): | |
| index = limbSeq[i] | |
| color = color_list[index[0]] | |
| x = kps[index][:, 0] | |
| y = kps[index][:, 1] | |
| length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 | |
| angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) | |
| polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, | |
| 360, 1) | |
| out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) | |
| out_img = (out_img * 0.6).astype(np.uint8) | |
| for idx_kp, kp in enumerate(kps): | |
| color = color_list[idx_kp] | |
| x, y = kp | |
| out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) | |
| out_img_pil = PIL.Image.fromarray(out_img.astype(np.uint8)) | |
| return out_img_pil | |
| def predict_mask(segmentmodel, sam, image, TEXT_PROMPT, segmentType, confidence = 0.2, threshold = 0.5): | |
| if segmentType=='GroundingDINO': | |
| image_source, image = load_image_dino(image) | |
| boxes, logits, phrases = predict( | |
| model=segmentmodel, | |
| image=image, | |
| caption=TEXT_PROMPT, | |
| box_threshold=0.3, | |
| text_threshold=0.25 | |
| ) | |
| sam.set_image(image_source) | |
| H, W, _ = image_source.shape | |
| boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H]) | |
| transformed_boxes = sam.transform.apply_boxes_torch(boxes_xyxy, image_source.shape[:2]).cuda() | |
| masks, _, _ = sam.predict_torch( | |
| point_coords=None, | |
| point_labels=None, | |
| boxes=transformed_boxes, | |
| multimask_output=False, | |
| ) | |
| masks=masks[0].squeeze(0) | |
| else: | |
| image_source = load_image_yoloworld(image) | |
| segmentmodel.set_classes([TEXT_PROMPT]) | |
| results = segmentmodel.infer(image_source, confidence=confidence) | |
| detections = sv.Detections.from_inference(results).with_nms( | |
| class_agnostic=True, threshold=threshold | |
| ) | |
| masks = None | |
| if len(detections) != 0: | |
| print(TEXT_PROMPT + " detected!") | |
| sam.set_image(image_source, image_format="RGB") | |
| masks, _, _ = sam.predict(box=detections.xyxy[0], multimask_output=False) | |
| masks = torch.from_numpy(masks.squeeze()) | |
| return masks | |
| def build_model_sd(pretrained_model, controlnet_path, face_adapter, device, prompts, antelopev2_path, width, height, style_lora): | |
| controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16).to(device) | |
| pipe = InstantidMultiConceptPipeline.from_pretrained( | |
| pretrained_model, controlnet=controlnet, torch_dtype=torch.float16).to(device) | |
| controller = AttentionReplace(prompts, 50, cross_replace_steps={"default_": 1.}, | |
| self_replace_steps=0.4, tokenizer=pipe.tokenizer, device=device, width=width, height=height, | |
| dtype=torch.float16) | |
| revise_regionally_controlnet_forward(pipe.unet, controller) | |
| controlnet_concept = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16) | |
| pipe_concept = InstantidSingleConceptPipeline.from_pretrained( | |
| pretrained_model, | |
| controlnet=controlnet_concept, | |
| torch_dtype=torch.float16 | |
| ) | |
| pipe_concept.load_ip_adapter_instantid(face_adapter) | |
| pipe_concept.set_ip_adapter_scale(0.8) | |
| pipe_concept.to(device) | |
| pipe_concept.image_proj_model.to(pipe_concept._execution_device) | |
| if style_lora is not None and os.path.exists(style_lora): | |
| pipe.load_lora_weights(style_lora, weight_name="pytorch_lora_weights.safetensors", adapter_name='style') | |
| pipe_concept.load_lora_weights(style_lora, weight_name="pytorch_lora_weights.safetensors", adapter_name='style') | |
| # modify | |
| app = FaceAnalysis(name='antelopev2', root=antelopev2_path, | |
| providers=['CPUExecutionProvider']) | |
| app.prepare(ctx_id=0, det_size=(640, 640)) | |
| return pipe, controller, pipe_concept, app | |
| def prepare_text(prompt, region_prompts): | |
| ''' | |
| Args: | |
| prompt_entity: [subject1]-*-[attribute1]-*-[Location1]|[subject2]-*-[attribute2]-*-[Location2]|[global text] | |
| Returns: | |
| full_prompt: subject1, attribute1 and subject2, attribute2, global text | |
| context_prompt: subject1 and subject2, global text | |
| entity_collection: [(subject1, attribute1), Location1] | |
| ''' | |
| region_collection = [] | |
| regions = region_prompts.split('|') | |
| for region in regions: | |
| if region == '': | |
| break | |
| prompt_region, neg_prompt_region, ref_img = region.split('-*-') | |
| prompt_region = prompt_region.replace('[', '').replace(']', '') | |
| neg_prompt_region = neg_prompt_region.replace('[', '').replace(']', '') | |
| region_collection.append((prompt_region, neg_prompt_region, ref_img)) | |
| return (prompt, region_collection) | |
| def build_model_lora(pipe, pipe_concept, style_path, condition, condition_img): | |
| if condition == "Human pose" and condition_img is not None: | |
| controlnet = ControlNetModel.from_pretrained(args.openpose_checkpoint, torch_dtype=torch.float16).to(device) | |
| pipe.controlnet2 = controlnet | |
| elif condition == "Canny Edge" and condition_img is not None: | |
| controlnet = ControlNetModel.from_pretrained(args.canny_checkpoint, torch_dtype=torch.float16, variant="fp16").to(device) | |
| pipe.controlnet2 = controlnet | |
| elif condition == "Depth" and condition_img is not None: | |
| controlnet = ControlNetModel.from_pretrained(args.depth_checkpoint, torch_dtype=torch.float16).to(device) | |
| pipe.controlnet2 = controlnet | |
| if style_path is not None and os.path.exists(style_path): | |
| pipe_concept.load_lora_weights(style_path, weight_name="pytorch_lora_weights.safetensors", adapter_name='style') | |
| pipe.load_lora_weights(style_path, weight_name="pytorch_lora_weights.safetensors", adapter_name='style') | |
| def resize_and_center_crop(image, output_size=(1024, 576)): | |
| width, height = image.size | |
| aspect_ratio = width / height | |
| new_height = output_size[1] | |
| new_width = int(aspect_ratio * new_height) | |
| resized_image = image.resize((new_width, new_height), Image.LANCZOS) | |
| if new_width < output_size[0] or new_height < output_size[1]: | |
| padding_color = "gray" | |
| resized_image = ImageOps.expand(resized_image, | |
| ((output_size[0] - new_width) // 2, | |
| (output_size[1] - new_height) // 2, | |
| (output_size[0] - new_width + 1) // 2, | |
| (output_size[1] - new_height + 1) // 2), | |
| fill=padding_color) | |
| left = (resized_image.width - output_size[0]) / 2 | |
| top = (resized_image.height - output_size[1]) / 2 | |
| right = (resized_image.width + output_size[0]) / 2 | |
| bottom = (resized_image.height + output_size[1]) / 2 | |
| cropped_image = resized_image.crop((left, top, right, bottom)) | |
| return cropped_image | |
| def main(device, segment_type): | |
| pipe, controller, pipe_concepts, face_app = build_model_sd(args.pretrained_model, args.controlnet_path, | |
| args.face_adapter_path, device, prompts_tmp, | |
| args.antelopev2_path, width // 32, height // 32, | |
| args.style_lora) | |
| if segment_type == 'GroundingDINO': | |
| detect_model, sam = build_dino_segment_model(args.dino_checkpoint, args.sam_checkpoint) | |
| else: | |
| detect_model, sam = build_yolo_segment_model(args.efficientViT_checkpoint, device) | |
| resolution_list = ["1440*728", | |
| "1344*768", | |
| "1216*832", | |
| "1152*896", | |
| "1024*1024", | |
| "896*1152", | |
| "832*1216", | |
| "768*1344", | |
| "728*1440"] | |
| ratio_list = [1440 / 728, 1344 / 768, 1216 / 832, 1152 / 896, 1024 / 1024, 896 / 1152, 832 / 1216, 768 / 1344, | |
| 728 / 1440] | |
| condition_list = ["None", | |
| "Human pose", | |
| "Canny Edge", | |
| "Depth"] | |
| depth_estimator = DPTForDepthEstimation.from_pretrained(args.dpt_checkpoint).to("cuda") | |
| feature_extractor = DPTFeatureExtractor.from_pretrained(args.dpt_checkpoint) | |
| body_model = Body(args.pose_detector_checkpoint) | |
| openpose = OpenposeDetector(body_model) | |
| prompts_rewrite = [args.prompt_rewrite] | |
| input_prompt_test = [prepare_text(p, p_w) for p, p_w in zip(prompts, prompts_rewrite)] | |
| input_prompt_test = [prompts, input_prompt_test[0][1]] | |
| def remove_tips(): | |
| return gr.update(visible=False) | |
| def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| return seed | |
| def get_humanpose(img): | |
| openpose_image = openpose(img) | |
| return openpose_image | |
| def get_cannyedge(image): | |
| image = np.array(image) | |
| image = cv2.Canny(image, 100, 200) | |
| image = image[:, :, None] | |
| image = np.concatenate([image, image, image], axis=2) | |
| canny_image = Image.fromarray(image) | |
| return canny_image | |
| def get_depth(image, height, weight): | |
| image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda") | |
| with torch.no_grad(), torch.autocast("cuda"): | |
| depth_map = depth_estimator(image).predicted_depth | |
| depth_map = torch.nn.functional.interpolate( | |
| depth_map.unsqueeze(1), | |
| size=(height, weight), | |
| mode="bicubic", | |
| align_corners=False, | |
| ) | |
| depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) | |
| depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) | |
| depth_map = (depth_map - depth_min) / (depth_max - depth_min) | |
| image = torch.cat([depth_map] * 3, dim=1) | |
| image = image.permute(0, 2, 3, 1).cpu().numpy()[0] | |
| image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8)) | |
| return image | |
| def generate_image(prompt1, negative_prompt, reference_1, reference_2, resolution, local_prompt1, local_prompt2, seed, style, identitynet_strength_ratio, adapter_strength_ratio, condition, condition_img, controlnet_ratio, cfg_scale): | |
| identitynet_strength_ratio = float(identitynet_strength_ratio) | |
| adapter_strength_ratio = float(adapter_strength_ratio) | |
| controlnet_ratio = float(controlnet_ratio) | |
| cfg_scale = float(cfg_scale) | |
| if lorapath_styles[style] is not None and os.path.exists(lorapath_styles[style]): | |
| styleL = True | |
| else: | |
| styleL = False | |
| width, height = int(resolution.split("*")[0]), int(resolution.split("*")[1]) | |
| kwargs = { | |
| 'height': height, | |
| 'width': width, | |
| 't2i_controlnet_conditioning_scale': controlnet_ratio, | |
| } | |
| if condition == 'Human pose' and condition_img is not None: | |
| index = ratio_list.index( | |
| min(ratio_list, key=lambda x: abs(x - condition_img.shape[1] / condition_img.shape[0]))) | |
| resolution = resolution_list[index] | |
| width, height = int(resolution.split("*")[0]), int(resolution.split("*")[1]) | |
| kwargs['height'] = height | |
| kwargs['width'] = width | |
| condition_img = resize_and_center_crop(Image.fromarray(condition_img), (width, height)) | |
| spatial_condition = get_humanpose(condition_img) | |
| elif condition == 'Canny Edge' and condition_img is not None: | |
| index = ratio_list.index( | |
| min(ratio_list, key=lambda x: abs(x - condition_img.shape[1] / condition_img.shape[0]))) | |
| resolution = resolution_list[index] | |
| width, height = int(resolution.split("*")[0]), int(resolution.split("*")[1]) | |
| kwargs['height'] = height | |
| kwargs['width'] = width | |
| condition_img = resize_and_center_crop(Image.fromarray(condition_img), (width, height)) | |
| spatial_condition = get_cannyedge(condition_img) | |
| elif condition == 'Depth' and condition_img is not None: | |
| index = ratio_list.index( | |
| min(ratio_list, key=lambda x: abs(x - condition_img.shape[1] / condition_img.shape[0]))) | |
| resolution = resolution_list[index] | |
| width, height = int(resolution.split("*")[0]), int(resolution.split("*")[1]) | |
| kwargs['height'] = height | |
| kwargs['width'] = width | |
| condition_img = resize_and_center_crop(Image.fromarray(condition_img), (width, height)) | |
| spatial_condition = get_depth(condition_img, height, width) | |
| else: | |
| spatial_condition = None | |
| kwargs['t2i_image'] = spatial_condition | |
| pipe.unload_lora_weights() | |
| pipe_concepts.unload_lora_weights() | |
| build_model_lora(pipe, pipe_concepts, lorapath_styles[style], condition, condition_img) | |
| pipe_concepts.set_ip_adapter_scale(adapter_strength_ratio) | |
| input_list = [prompt1] | |
| for prompt in input_list: | |
| if prompt != '': | |
| input_prompt = [] | |
| p = '{prompt}, 35mm photograph, film, professional, 4k, highly detailed.' | |
| if styleL: | |
| p = styles[style] + p | |
| input_prompt.append([p.replace('{prompt}', prompt), p.replace("{prompt}", prompt)]) | |
| if styleL: | |
| input_prompt.append([(styles[style] + local_prompt1, 'noisy, blurry, soft, deformed, ugly', | |
| PIL.Image.fromarray(reference_1)), | |
| (styles[style] + local_prompt2, 'noisy, blurry, soft, deformed, ugly', | |
| PIL.Image.fromarray(reference_2))]) | |
| else: | |
| input_prompt.append( | |
| [(local_prompt1, 'noisy, blurry, soft, deformed, ugly', PIL.Image.fromarray(reference_1)), | |
| (local_prompt2, 'noisy, blurry, soft, deformed, ugly', PIL.Image.fromarray(reference_2))]) | |
| controller.reset() | |
| image = sample_image( | |
| pipe, | |
| input_prompt=input_prompt, | |
| concept_models=pipe_concepts, | |
| input_neg_prompt=[negative_prompt] * len(input_prompt), | |
| generator=torch.Generator(device).manual_seed(seed), | |
| controller=controller, | |
| face_app=face_app, | |
| controlnet_conditioning_scale=identitynet_strength_ratio, | |
| stage=1, | |
| guidance_scale=cfg_scale, | |
| **kwargs) | |
| controller.reset() | |
| if pipe.tokenizer("man")["input_ids"][1] in pipe.tokenizer(args.prompt)["input_ids"][1:-1]: | |
| mask1 = predict_mask(detect_model, sam, image[0], 'man', args.segment_type, confidence=0.05, | |
| threshold=0.5) | |
| else: | |
| mask1 = None | |
| if pipe.tokenizer("woman")["input_ids"][1] in pipe.tokenizer(args.prompt)["input_ids"][1:-1]: | |
| mask2 = predict_mask(detect_model, sam, image[0], 'woman', args.segment_type, confidence=0.05, | |
| threshold=0.5) | |
| else: | |
| mask2 = None | |
| if mask1 is not None or mask2 is not None: | |
| face_info = face_app.get(cv2.cvtColor(np.array(image[0]), cv2.COLOR_RGB2BGR)) | |
| face_kps = draw_kps_multi(image[0], [face['kps'] for face in face_info]) | |
| image = sample_image( | |
| pipe, | |
| input_prompt=input_prompt, | |
| concept_models=pipe_concepts, | |
| input_neg_prompt=[negative_prompt] * len(input_prompt), | |
| generator=torch.Generator(device).manual_seed(seed), | |
| controller=controller, | |
| face_app=face_app, | |
| image=face_kps, | |
| stage=2, | |
| controlnet_conditioning_scale=identitynet_strength_ratio, | |
| region_masks=[mask1, mask2], | |
| guidance_scale=cfg_scale, | |
| **kwargs) | |
| return [image[1], spatial_condition] | |
| # return image | |
| with gr.Blocks(css=css) as demo: | |
| # description | |
| gr.Markdown(title) | |
| gr.Markdown(description) | |
| with gr.Row(): | |
| gallery = gr.Image(label="Generated Images", height=512, width=512) | |
| gallery1 = gr.Image(label="Input Condition", height=512, width=512) | |
| usage_tips = gr.Markdown(label="Usage tips of OMG", value=tips, visible=False) | |
| with gr.Row(): | |
| reference_1 = gr.Image(label="Input an RGB image for Character man", height=128, width=128) | |
| reference_2 = gr.Image(label="Input an RGB image for Character woman", height=128, width=128) | |
| condition_img1 = gr.Image(label="Input an RGB image for condition (Optional)", height=128, width=128) | |
| with gr.Row(): | |
| local_prompt1 = gr.Textbox(label="Character1_prompt", | |
| info="Describe the Character 1", | |
| value="Close-up photo of the a man, 35mm photograph, professional, 4k, highly detailed.") | |
| local_prompt2 = gr.Textbox(label="Character2_prompt", | |
| info="Describe the Character 2", | |
| value="Close-up photo of the a woman, 35mm photograph, professional, 4k, highly detailed.") | |
| with gr.Row(): | |
| identitynet_strength_ratio = gr.Slider( | |
| label="IdentityNet strength (for fidelity)", | |
| minimum=0, | |
| maximum=1.5, | |
| step=0.05, | |
| value=0.80, | |
| ) | |
| adapter_strength_ratio = gr.Slider( | |
| label="Image adapter strength (for detail)", | |
| minimum=0, | |
| maximum=1.5, | |
| step=0.05, | |
| value=0.80, | |
| ) | |
| controlnet_ratio = gr.Slider( | |
| label="ControlNet strength", | |
| minimum=0, | |
| maximum=1.5, | |
| step=0.05, | |
| value=1, | |
| ) | |
| cfg_ratio = gr.Slider( | |
| label="CFG scale ", | |
| minimum=0.5, | |
| maximum=10, | |
| step=0.5, | |
| value=3.0, | |
| ) | |
| resolution = gr.Dropdown(label="Image Resolution (width*height)", choices=resolution_list, | |
| value="1024*1024") | |
| style = gr.Dropdown(label="style", choices=STYLE_NAMES, value="None") | |
| condition = gr.Dropdown(label="Input condition type", choices=condition_list, value="None") | |
| # prompt | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Prompt 1", | |
| info="Give a simple prompt to describe the first image content", | |
| placeholder="Required", | |
| value="close-up shot, photography, a man and a woman on the street, facing the camera smiling") | |
| with gr.Accordion(open=False, label="Advanced Options"): | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=42, | |
| ) | |
| negative_prompt = gr.Textbox(label="Negative Prompt", | |
| placeholder="noisy, blurry, soft, deformed, ugly", | |
| value="noisy, blurry, soft, deformed, ugly") | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| submit = gr.Button("Submit", variant="primary") | |
| submit.click( | |
| fn=remove_tips, | |
| outputs=usage_tips, | |
| ).then( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=generate_image, | |
| inputs=[prompt, negative_prompt, reference_1, reference_2, resolution, local_prompt1, local_prompt2, seed, style, identitynet_strength_ratio, adapter_strength_ratio, condition, condition_img1, controlnet_ratio, cfg_ratio], | |
| outputs=[gallery, gallery1] | |
| ) | |
| gr.Markdown(article) | |
| demo.launch(share=True) | |
| def parse_args(): | |
| parser = argparse.ArgumentParser('', add_help=False) | |
| parser.add_argument('--pretrained_model', default='wangqixun/YamerMIX_v8', type=str) | |
| parser.add_argument('--controlnet_path', default='/home/user/app/checkpoint/InstantID/ControlNetModel', type=str) | |
| parser.add_argument('--face_adapter_path', default='/home/user/app/checkpoint/InstantID/ip-adapter.bin', type=str) | |
| parser.add_argument('--openpose_checkpoint', default='thibaud/controlnet-openpose-sdxl-1.0', type=str) | |
| parser.add_argument('--canny_checkpoint', default='diffusers/controlnet-canny-sdxl-1.0', type=str) | |
| parser.add_argument('--depth_checkpoint', default='diffusers/controlnet-depth-sdxl-1.0', type=str) | |
| parser.add_argument('--dpt_checkpoint', default='Intel/dpt-hybrid-midas', type=str) | |
| parser.add_argument('--pose_detector_checkpoint', | |
| default='/home/user/app/checkpoint/ControlNet/annotator/ckpts/body_pose_model.pth', type=str) | |
| parser.add_argument('--efficientViT_checkpoint', default='/home/user/app/checkpoint/sam/xl1.pt', type=str) | |
| parser.add_argument('--dino_checkpoint', default='/home/user/app/checkpoint/GroundingDINO', type=str) | |
| parser.add_argument('--sam_checkpoint', default='/home/user/app/checkpoint/sam/sam_vit_h_4b8939.pth', type=str) | |
| parser.add_argument('--antelopev2_path', default='/home/user/app/checkpoint/antelopev2', type=str) | |
| parser.add_argument('--save_dir', default='results/instantID', type=str) | |
| parser.add_argument('--prompt', default='Close-up photo of the cool man and beautiful woman as they accidentally discover a mysterious island while on vacation by the sea, facing the camera smiling, 35mm photograph, film, professional, 4k, highly detailed.', type=str) | |
| parser.add_argument('--negative_prompt', default='noisy, blurry, soft, deformed, ugly', type=str) | |
| parser.add_argument('--prompt_rewrite', | |
| default='[Close-up photo of a man, 35mm photograph, professional, 4k, highly detailed.]-*' | |
| '-[noisy, blurry, soft, deformed, ugly]-*-' | |
| '../example/chris-evans.jpg|' | |
| '[Close-up photo of a woman, 35mm photograph, professional, 4k, highly detailed.]-' | |
| '*-[noisy, blurry, soft, deformed, ugly]-*-' | |
| '../example/TaylorSwift.png', | |
| type=str) | |
| parser.add_argument('--seed', default=0, type=int) | |
| parser.add_argument('--suffix', default='', type=str) | |
| parser.add_argument('--segment_type', default='yoloworld', help='GroundingDINO or yoloworld', type=str) | |
| parser.add_argument('--style_lora', default='', type=str) | |
| return parser.parse_args() | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| prompts = [args.prompt] * 2 | |
| prompts_tmp = copy.deepcopy(prompts) | |
| width, height = 1024, 1024 | |
| kwargs = { | |
| 'height': height, | |
| 'width': width, | |
| } | |
| device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') | |
| main(device, args.segment_type) | |