import open3d_zerogpu_fix import spaces from diffusers import ControlNetModel from diffusers import StableDiffusionXLControlNetPipeline from diffusers import EulerAncestralDiscreteScheduler from PIL import Image import torch import numpy as np import cv2 import gradio as gr from torchvision import transforms from controlnet_aux import OpenposeDetector import random import open3d as o3d from collections import Counter import trimesh ratios_map = { 0.5:{"width":704,"height":1408}, 0.57:{"width":768,"height":1344}, 0.68:{"width":832,"height":1216}, 0.72:{"width":832,"height":1152}, 0.78:{"width":896,"height":1152}, 0.82:{"width":896,"height":1088}, 0.88:{"width":960,"height":1088}, 0.94:{"width":960,"height":1024}, 1.00:{"width":1024,"height":1024}, 1.13:{"width":1088,"height":960}, 1.21:{"width":1088,"height":896}, 1.29:{"width":1152,"height":896}, 1.38:{"width":1152,"height":832}, 1.46:{"width":1216,"height":832}, 1.67:{"width":1280,"height":768}, 1.75:{"width":1344,"height":768}, 2.00:{"width":1408,"height":704} } ratios = np.array(list(ratios_map.keys())) openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet') controlnet = ControlNetModel.from_pretrained( "yeq6x/Image2PositionColor_v3", torch_dtype=torch.float16 ).to('cuda') pipe = StableDiffusionXLControlNetPipeline.from_pretrained( "yeq6x/animagine_position_map", controlnet=controlnet, torch_dtype=torch.float16, low_cpu_mem_usage=True, offload_state_dict=True, ).to('cuda').to(torch.float16) pipe.scheduler = EulerAncestralDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, steps_offset=1 ) # pipe.enable_freeu(b1=1.1, b2=1.1, s1=0.5, s2=0.7) # pipe.enable_xformers_memory_efficient_attention() pipe.force_zeros_for_empty_prompt = False def get_size(init_image): w,h=init_image.size curr_ratio = w/h ind = np.argmin(np.abs(curr_ratio-ratios)) ratio = ratios[ind] chosen_ratio = ratios_map[ratio] w,h = chosen_ratio['width'], chosen_ratio['height'] return w,h def resize_image(image): image = image.convert('RGB') w,h = get_size(image) resized_image = image.resize((w, h)) return resized_image def resize_image_old(image): image = image.convert('RGB') current_size = image.size if current_size[0] > current_size[1]: center_cropped_image = transforms.functional.center_crop(image, (current_size[1], current_size[1])) else: center_cropped_image = transforms.functional.center_crop(image, (current_size[0], current_size[0])) resized_image = transforms.functional.resize(center_cropped_image, (1024, 1024)) return resized_image @spaces.GPU def generate_(prompt, negative_prompt, pose_image, input_image, controlnet_conditioning_scale): generator = torch.Generator() generator.manual_seed(random.randint(0, 2147483647)) images = pipe( prompt, negative_prompt=negative_prompt, image=pose_image, num_inference_steps=20, controlnet_conditioning_scale=float(controlnet_conditioning_scale), generator=generator, height=input_image.size[1], width=input_image.size[0], ).images return images @spaces.GPU def process(input_image, prompt, negative_prompt, controlnet_conditioning_scale): # resize input_image to 1024x1024 input_image = resize_image(input_image) pose_image = openpose(input_image, include_body=True, include_hand=True, include_face=True) images = generate_(prompt, negative_prompt, pose_image, input_image, controlnet_conditioning_scale) return [pose_image,images[0]] @spaces.GPU def predict_image(cond_image, prompt, negative_prompt, controlnet_conditioning_scale): print("predict position map") global pipe generator = torch.Generator() generator.manual_seed(random.randint(0, 2147483647)) image = pipe( prompt, negative_prompt=negative_prompt, image = cond_image, width=1024, height=1024, guidance_scale=8, num_inference_steps=20, generator=generator, guess_mode = True, controlnet_conditioning_scale = controlnet_conditioning_scale ).images[0] return image def convert_pil_to_opencv(pil_image): return np.array(pil_image) def inv_func(y, c = -712.380100, a = 137.375240, b = 192.435866): return (np.exp((y - c) / a) - np.exp(-c/a)) / 964.8468371292845 def create_point_cloud(img1, img2): if img1.shape != img2.shape: raise ValueError("Both images must have the same dimensions.") h, w, _ = img1.shape points = [] colors = [] for y in range(h): for x in range(w): # ピクセル位置 (x, y) のRGBをXYZとして取得 r, g, b = img1[y, x] r = inv_func(r) * 0.9 g = inv_func(g) / 1.7 * 0.6 b = inv_func(b) r *= 150 g *= 150 b *= 150 points.append([g, b, r]) # X, Y, Z # 対応するピクセル位置の画像2の色を取得 colors.append(img2[y, x] / 255.0) # 色は0〜1にスケール return np.array(points), np.array(colors) def point_cloud_to_glb(points, colors): # Open3Dでポイントクラウドを作成 pc = o3d.geometry.PointCloud() pc.points = o3d.utility.Vector3dVector(points) pc.colors = o3d.utility.Vector3dVector(colors) # 一時的にPLY形式で保存 temp_ply_file = "temp_output.ply" o3d.io.write_point_cloud(temp_ply_file, pc) # PLYをGLBに変換 mesh = trimesh.load(temp_ply_file) glb_file = "output.glb" mesh.export(glb_file) return glb_file def visualize_3d(image1, image2): print("Processing...") # PIL画像をOpenCV形式に変換 img1 = convert_pil_to_opencv(image1) img2 = convert_pil_to_opencv(image2) # ポイントクラウド生成 points, colors = create_point_cloud(img1, img2) # GLB形式に変換 glb_file = point_cloud_to_glb(points, colors) return glb_file def scale_image(original_image): aspect_ratio = original_image.width / original_image.height if original_image.width > original_image.height: new_width = 1024 new_height = round(new_width / aspect_ratio) else: new_height = 1024 new_width = round(new_height * aspect_ratio) resized_original = original_image.resize((new_width, new_height), Image.LANCZOS) return resized_original def get_edge_mode_color(img, edge_width=10): # 外周の10ピクセル領域を取得 left = img.crop((0, 0, edge_width, img.height)) # 左端 right = img.crop((img.width - edge_width, 0, img.width, img.height)) # 右端 top = img.crop((0, 0, img.width, edge_width)) # 上端 bottom = img.crop((0, img.height - edge_width, img.width, img.height)) # 下端 # 各領域のピクセルデータを取得して結合 colors = list(left.getdata()) + list(right.getdata()) + list(top.getdata()) + list(bottom.getdata()) # 最頻値（mode）を計算 mode_color = Counter(colors).most_common(1)[0][0] # 最も頻繁に出現する色を取得 return mode_color def paste_image(resized_img): # 外周10pxの最頻値を背景色に設定 mode_color = get_edge_mode_color(resized_img, edge_width=10) mode_background = Image.new("RGBA", (1024, 1024), mode_color) mode_background = mode_background.convert('RGB') x = (1024 - resized_img.width) // 2 y = (1024 - resized_img.height) // 2 mode_background.paste(resized_img, (x, y)) return mode_background def outpaint_image(image): if type(image) == type(None): return None resized_img = scale_image(image) image = paste_image(resized_img) return image block = gr.Blocks().queue() with block: gr.Markdown("## BRIA 2.3 ControlNet Pose") with gr.Row(): with gr.Column(): input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam prompt = gr.Textbox(label="Prompt") negative_prompt = gr.Textbox(label="Negative prompt", value="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers") controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05) run_button = gr.Button(value="Run") with gr.Column(): with gr.Row(): pose_image_output = gr.Image(label="Pose Image", type="pil", interactive=False) generated_image_output = gr.Image(label="Generated Image", type="pil", interactive=False) run_button.click(fn=process, inputs=[input_image, prompt, negative_prompt, controlnet_conditioning_scale], outputs=[pose_image_output, generated_image_output]) block.launch(debug = True)