de

app.py

@@ -8,6 +8,7 @@ from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, Eule
 import torch
 from collections import Counter
 import random
+from controlnet_aux import OpenposeDetector
 
 ratios_map =  {
     0.5:{"width":704,"height":1408},
@@ -30,6 +31,7 @@ ratios_map =  {
 }
 ratios = np.array(list(ratios_map.keys()))
 
+openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
 controlnet = ControlNetModel.from_pretrained(
     "yeq6x/Image2PositionColor_v3",
     torch_dtype=torch.float16
@@ -43,13 +45,14 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     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.scheduler = EulerAncestralDiscreteScheduler(
+#     beta_start=0.00085,
+#     beta_end=0.012,
+#     beta_schedule="scaled_linear",
+#     num_train_timesteps=1000,
+#     steps_offset=1
+# )
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.force_zeros_for_empty_prompt = False
 
 def get_size(init_image):
@@ -78,147 +81,148 @@ def generate_(prompt, negative_prompt, pose_image, input_image, num_steps, contr
 
 @spaces.GPU
 def process(input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed):
-    
     # resize input_image to 1024x1024
     input_image = resize_image(input_image)
     
-    images = generate_(prompt, negative_prompt, input_image, input_image, num_steps, controlnet_conditioning_scale, seed)
-
-    return [input_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]
+    pose_image = openpose(input_image, include_body=True, include_hand=True, include_face=True)
+  
+    images = generate_(prompt, negative_prompt, pose_image, input_image, num_steps, controlnet_conditioning_scale, seed)
+
+    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)
+#   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形式で保存
+#   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)
+#   # PLYをGLBに変換
+#   mesh = trimesh.load(temp_ply_file)
+#   glb_file = "output.glb"
+#   mesh.export(glb_file)
 
-  return 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)
+# 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)
+#   # ポイントクラウド生成
+#   points, colors = create_point_cloud(img1, img2)
 
-  # GLB形式に変換
-  glb_file = point_cloud_to_glb(points, colors)
+#   # GLB形式に変換
+#   glb_file = point_cloud_to_glb(points, colors)
 
-  return glb_file
+#   return glb_file
 
-def scale_image(original_image):
-  aspect_ratio = original_image.width / original_image.height
+# 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)
+#   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)
+#   resized_original = original_image.resize((new_width, new_height), Image.LANCZOS)
 
-  return resized_original
+#   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))  # 下端
+# 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())
+#   # 各領域のピクセルデータを取得して結合
+#   colors = list(left.getdata()) + list(right.getdata()) + list(top.getdata()) + list(bottom.getdata())
 
-  # 最頻値（mode）を計算
-  mode_color = Counter(colors).most_common(1)[0][0]  # 最も頻繁に出現する色を取得
+#   # 最頻値（mode）を計算
+#   mode_color = Counter(colors).most_common(1)[0][0]  # 最も頻繁に出現する色を取得
 
-  return mode_color
+#   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')
+# 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))
+#   x = (1024 - resized_img.width) // 2
+#   y = (1024 - resized_img.height) // 2
+#   mode_background.paste(resized_img, (x, y))
 
-  return mode_background
+#   return mode_background
 
-def outpaint_image(image):
-  if type(image) == type(None):
-    return None
-  resized_img = scale_image(image)
-  image = paste_image(resized_img)
+# def outpaint_image(image):
+#   if type(image) == type(None):
+#     return None
+#   resized_img = scale_image(image)
+#   image = paste_image(resized_img)
   
-  return image
+#   return image
 
 block = gr.Blocks().queue()
 
@@ -241,30 +245,30 @@ with block:
   ips = [input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed]
   run_button.click(fn=process, inputs=ips, outputs=[pose_image_output, generated_image_output])
 
-  gr.Markdown("## Position Map Visualizer")
+  # gr.Markdown("## Position Map Visualizer")
   
-  with gr.Row():
-    with gr.Column():
-      with gr.Row():
-        img1 = gr.Image(type="pil", label="color Image", height=300)
-        img2 = gr.Image(type="pil", label="map Image", height=300)
-      prompt = gr.Textbox("position map, 1girl, white background", label="Prompt")
-      negative_prompt = gr.Textbox("lowres, bad anatomy, bad hands, bad feet, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry", label="Negative Prompt")
-      controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=0.6, step=0.05)
-      predict_map_btn = gr.Button("Predict Position Map")
-      visualize_3d_btn = gr.Button("Generate 3D Point Cloud")
-    with gr.Column():
-      reconstruction_output = gr.Model3D(label="3D Viewer", height=600)
-      gr.Examples(
-          examples=[
-          ["resources/source/000006.png", "resources/target/000006.png"],
-          ["resources/source/006420.png", "resources/target/006420.png"],
-      ],
-          inputs=[img1, img2]
-      )
-
-  img1.input(outpaint_image, inputs=img1, outputs=img1)
-  predict_map_btn.click(predict_image, inputs=[img1, prompt, negative_prompt, controlnet_conditioning_scale], outputs=img2)
-  visualize_3d_btn.click(visualize_3d, inputs=[img2, img1], outputs=reconstruction_output)
+  # with gr.Row():
+  #   with gr.Column():
+  #     with gr.Row():
+  #       img1 = gr.Image(type="pil", label="color Image", height=300)
+  #       img2 = gr.Image(type="pil", label="map Image", height=300)
+  #     prompt = gr.Textbox("position map, 1girl, white background", label="Prompt")
+  #     negative_prompt = gr.Textbox("lowres, bad anatomy, bad hands, bad feet, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry", label="Negative Prompt")
+  #     controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=0.6, step=0.05)
+  #     predict_map_btn = gr.Button("Predict Position Map")
+  #     visualize_3d_btn = gr.Button("Generate 3D Point Cloud")
+  #   with gr.Column():
+  #     reconstruction_output = gr.Model3D(label="3D Viewer", height=600)
+  #     gr.Examples(
+  #         examples=[
+  #         ["resources/source/000006.png", "resources/target/000006.png"],
+  #         ["resources/source/006420.png", "resources/target/006420.png"],
+  #     ],
+  #         inputs=[img1, img2]
+  #     )
+
+  # img1.input(outpaint_image, inputs=img1, outputs=img1)
+  # predict_map_btn.click(predict_image, inputs=[img1, prompt, negative_prompt, controlnet_conditioning_scale], outputs=img2)
+  # visualize_3d_btn.click(visualize_3d, inputs=[img2, img1], outputs=reconstruction_output)
 
 block.launch(debug = True)