load_model

app.py

@@ -13,15 +13,15 @@ import spaces
 pipe = None
 
 def load_model():
-    global pipe
-    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-        "yeq6x/animagine_position_map",
-        controlnet=ControlNetModel.from_pretrained("yeq6x/Image2PositionColor_v3"),
-    ).to("cuda")
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+  global pipe
+  pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+      "yeq6x/animagine_position_map",
+      controlnet=ControlNetModel.from_pretrained("yeq6x/Image2PositionColor_v3"),
+  ).to("cuda")
+  pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
 def convert_pil_to_opencv(pil_image):
-    return np.array(pil_image)
+  return np.array(pil_image)
 
 def inv_func(y,
   c = -712.380100,
@@ -30,103 +30,103 @@ def inv_func(y,
   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)
+  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
+  # 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)
+  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
+  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))  # 下端
+  # 外周の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')
+  # 外周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
+    return None
   resized_img = scale_image(image)
   image = paste_image(resized_img)
   
@@ -134,29 +134,30 @@ def outpaint_image(image):
 
 @spaces.GPU
 def predict_image(cond_image, prompt, negative_prompt):
-    generator = torch.Generator()
-    generator.manual_seed(random.randint(0, 2147483647))
-
-    prompt = 'position map, 1girl, white background'
-    negative_prompt = "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"
-
-    image = pipe(
-        prompt,
-        prompt,
-        cond_image,
-        negative_prompt=negative_prompt,
-        width=1024,
-        height=1024,
-        guidance_scale=8,
-        num_inference_steps=20,
-        generator=generator,
-        guess_mode = True,
-        controlnet_conditioning_scale = 0.6,
-    ).images[0]
-    
-    return image
+  if pipe is None:
+    load_model()
+
+  generator = torch.Generator()
+  generator.manual_seed(random.randint(0, 2147483647))
+
+  prompt = 'position map, 1girl, white background'
+  negative_prompt = "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"
+
+  image = pipe(
+      prompt,
+      prompt,
+      cond_image,
+      negative_prompt=negative_prompt,
+      width=1024,
+      height=1024,
+      guidance_scale=8,
+      num_inference_steps=20,
+      generator=generator,
+      guess_mode = True,
+      controlnet_conditioning_scale = 0.6,
+  ).images[0]
   
-load_model()
+  return image
 
 # Gradioアプリケーション
 with gr.Blocks() as demo: