Spaces:
Runtime error
Runtime error
| import os | |
| import os.path as osp | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from basicsr.utils import img2tensor, tensor2img | |
| from pytorch_lightning import seed_everything | |
| from ldm.models.diffusion.plms import PLMSSampler | |
| from ldm.modules.encoders.adapter import Adapter | |
| from ldm.util import instantiate_from_config | |
| from model_edge import pidinet | |
| import gradio as gr | |
| from omegaconf import OmegaConf | |
| import mmcv | |
| from mmdet.apis import inference_detector, init_detector | |
| from mmpose.apis import (inference_top_down_pose_model, init_pose_model, process_mmdet_results, vis_pose_result) | |
| skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10], | |
| [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]] | |
| pose_kpt_color = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], | |
| [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], | |
| [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0]] | |
| pose_link_color = [[0, 255, 0], [0, 255, 0], [255, 128, 0], [255, 128, 0], | |
| [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], [255, 128, 0], | |
| [0, 255, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], | |
| [51, 153, 255], [51, 153, 255], [51, 153, 255]] | |
| def imshow_keypoints(img, | |
| pose_result, | |
| skeleton=None, | |
| kpt_score_thr=0.1, | |
| pose_kpt_color=None, | |
| pose_link_color=None, | |
| radius=4, | |
| thickness=1): | |
| """Draw keypoints and links on an image. | |
| Args: | |
| img (ndarry): The image to draw poses on. | |
| pose_result (list[kpts]): The poses to draw. Each element kpts is | |
| a set of K keypoints as an Kx3 numpy.ndarray, where each | |
| keypoint is represented as x, y, score. | |
| kpt_score_thr (float, optional): Minimum score of keypoints | |
| to be shown. Default: 0.3. | |
| pose_kpt_color (np.array[Nx3]`): Color of N keypoints. If None, | |
| the keypoint will not be drawn. | |
| pose_link_color (np.array[Mx3]): Color of M links. If None, the | |
| links will not be drawn. | |
| thickness (int): Thickness of lines. | |
| """ | |
| img_h, img_w, _ = img.shape | |
| img = np.zeros(img.shape) | |
| for idx, kpts in enumerate(pose_result): | |
| if idx > 1: | |
| continue | |
| kpts = kpts['keypoints'] | |
| # print(kpts) | |
| kpts = np.array(kpts, copy=False) | |
| # draw each point on image | |
| if pose_kpt_color is not None: | |
| assert len(pose_kpt_color) == len(kpts) | |
| for kid, kpt in enumerate(kpts): | |
| x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2] | |
| if kpt_score < kpt_score_thr or pose_kpt_color[kid] is None: | |
| # skip the point that should not be drawn | |
| continue | |
| color = tuple(int(c) for c in pose_kpt_color[kid]) | |
| cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1) | |
| # draw links | |
| if skeleton is not None and pose_link_color is not None: | |
| assert len(pose_link_color) == len(skeleton) | |
| for sk_id, sk in enumerate(skeleton): | |
| pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1])) | |
| pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1])) | |
| if (pos1[0] <= 0 or pos1[0] >= img_w or pos1[1] <= 0 or pos1[1] >= img_h or pos2[0] <= 0 | |
| or pos2[0] >= img_w or pos2[1] <= 0 or pos2[1] >= img_h or kpts[sk[0], 2] < kpt_score_thr | |
| or kpts[sk[1], 2] < kpt_score_thr or pose_link_color[sk_id] is None): | |
| # skip the link that should not be drawn | |
| continue | |
| color = tuple(int(c) for c in pose_link_color[sk_id]) | |
| cv2.line(img, pos1, pos2, color, thickness=thickness) | |
| return img | |
| def load_model_from_config(config, ckpt, verbose=False): | |
| print(f"Loading model from {ckpt}") | |
| pl_sd = torch.load(ckpt, map_location="cpu") | |
| if "global_step" in pl_sd: | |
| print(f"Global Step: {pl_sd['global_step']}") | |
| if "state_dict" in pl_sd: | |
| sd = pl_sd["state_dict"] | |
| else: | |
| sd = pl_sd | |
| model = instantiate_from_config(config.model) | |
| m, u = model.load_state_dict(sd, strict=False) | |
| model.cuda() | |
| model.eval() | |
| return model | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| config = OmegaConf.load("configs/stable-diffusion/test_keypose.yaml") | |
| config.model.params.cond_stage_config.params.device = device | |
| model = load_model_from_config(config, "models/sd-v1-4.ckpt").to(device) | |
| current_base = 'sd-v1-4.ckpt' | |
| model_ad = Adapter(cin=int(3*64), channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device) | |
| model_ad.load_state_dict(torch.load("models/t2iadapter_keypose_sd14v1.pth")) | |
| sampler = PLMSSampler(model) | |
| ## mmpose | |
| det_config = 'models/faster_rcnn_r50_fpn_coco.py' | |
| det_checkpoint = 'models/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth' | |
| pose_config = 'models/hrnet_w48_coco_256x192.py' | |
| pose_checkpoint = 'models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth' | |
| det_cat_id = 1 | |
| bbox_thr = 0.2 | |
| ## detector | |
| det_config_mmcv = mmcv.Config.fromfile(det_config) | |
| det_model = init_detector(det_config_mmcv, det_checkpoint, device=device) | |
| pose_config_mmcv = mmcv.Config.fromfile(pose_config) | |
| pose_model = init_pose_model(pose_config_mmcv, pose_checkpoint, device=device) | |
| W, H = 512, 512 | |
| def process(input_img, type_in, prompt, neg_prompt, fix_sample, scale, con_strength, base_model): | |
| global current_base | |
| if current_base != base_model: | |
| ckpt = os.path.join("models", base_model) | |
| pl_sd = torch.load(ckpt, map_location="cpu") | |
| if "state_dict" in pl_sd: | |
| sd = pl_sd["state_dict"] | |
| else: | |
| sd = pl_sd | |
| model.load_state_dict(sd, strict=False) | |
| current_base = base_model | |
| con_strength = int((1-con_strength)*50) | |
| if fix_sample == 'True': | |
| seed_everything(42) | |
| im = cv2.resize(input_img,(W,H)) | |
| if type_in == 'Keypose': | |
| im_pose = im.copy() | |
| im = img2tensor(im).unsqueeze(0)/255. | |
| elif type_in == 'Image': | |
| image = im.copy() | |
| im = img2tensor(im).unsqueeze(0)/255. | |
| mmdet_results = inference_detector(det_model, image) | |
| # keep the person class bounding boxes. | |
| person_results = process_mmdet_results(mmdet_results, det_cat_id) | |
| # optional | |
| return_heatmap = False | |
| dataset = pose_model.cfg.data['test']['type'] | |
| # e.g. use ('backbone', ) to return backbone feature | |
| output_layer_names = None | |
| pose_results, returned_outputs = inference_top_down_pose_model( | |
| pose_model, | |
| image, | |
| person_results, | |
| bbox_thr=bbox_thr, | |
| format='xyxy', | |
| dataset=dataset, | |
| dataset_info=None, | |
| return_heatmap=return_heatmap, | |
| outputs=output_layer_names) | |
| # show the results | |
| im_pose = imshow_keypoints( | |
| image, | |
| pose_results, | |
| skeleton=skeleton, | |
| pose_kpt_color=pose_kpt_color, | |
| pose_link_color=pose_link_color, | |
| radius=2, | |
| thickness=2) | |
| im_pose = cv2.resize(im_pose,(W,H)) | |
| with torch.no_grad(): | |
| c = model.get_learned_conditioning([prompt]) | |
| nc = model.get_learned_conditioning([neg_prompt]) | |
| # extract condition features | |
| pose = img2tensor(im_pose, bgr2rgb=True, float32=True)/255. | |
| pose = pose.unsqueeze(0) | |
| features_adapter = model_ad(pose.to(device)) | |
| shape = [4, W//8, H//8] | |
| # sampling | |
| samples_ddim, _ = sampler.sample(S=50, | |
| conditioning=c, | |
| batch_size=1, | |
| shape=shape, | |
| verbose=False, | |
| unconditional_guidance_scale=scale, | |
| unconditional_conditioning=nc, | |
| eta=0.0, | |
| x_T=None, | |
| features_adapter1=features_adapter, | |
| mode = 'sketch', | |
| con_strength = con_strength) | |
| x_samples_ddim = model.decode_first_stage(samples_ddim) | |
| x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) | |
| x_samples_ddim = x_samples_ddim.to('cpu') | |
| x_samples_ddim = x_samples_ddim.permute(0, 2, 3, 1).numpy()[0] | |
| x_samples_ddim = 255.*x_samples_ddim | |
| x_samples_ddim = x_samples_ddim.astype(np.uint8) | |
| return [im_pose[:,:,::-1].astype(np.uint8), x_samples_ddim] | |
| DESCRIPTION = '''# T2I-Adapter (Keypose) | |
| [Paper](https://arxiv.org/abs/2302.08453) [GitHub](https://github.com/TencentARC/T2I-Adapter) | |
| This gradio demo is for keypose-guided generation. The current functions include: | |
| - Keypose to Image Generation | |
| - Image to Image Generation | |
| - Generation with **Anything** setting | |
| ''' | |
| block = gr.Blocks().queue() | |
| with block: | |
| with gr.Row(): | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(source='upload', type="numpy") | |
| prompt = gr.Textbox(label="Prompt") | |
| neg_prompt = gr.Textbox(label="Negative Prompt", | |
| value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face') | |
| with gr.Row(): | |
| type_in = gr.inputs.Radio(['Keypose', 'Image'], type="value", default='Image', label='Input Types\n (You can input an image or a keypose map)') | |
| fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed to produce a fixed output)') | |
| run_button = gr.Button(label="Run") | |
| con_strength = gr.Slider(label="Controling Strength (The guidance strength of the keypose to the result)", minimum=0, maximum=1, value=1, step=0.1) | |
| scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=9, step=0.1) | |
| base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use') | |
| with gr.Column(): | |
| result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto') | |
| ips = [input_img, type_in, prompt, neg_prompt, fix_sample, scale, con_strength, base_model] | |
| run_button.click(fn=process, inputs=ips, outputs=[result]) | |
| block.launch(server_name='0.0.0.0') | |