app.py

import spaces
import numpy as np
from PIL import Image
import gradio as gr
import open3d as o3d
import trimesh
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, EulerAncestralDiscreteScheduler
import torch
from collections import Counter
import random
from controlnet_aux import OpenposeDetector

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.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
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

@spaces.GPU
def generate_(prompt, negative_prompt, pose_image, input_image, num_steps, controlnet_conditioning_scale, seed):
    generator = torch.Generator("cuda").manual_seed(seed)    
    images = pipe(
    prompt, negative_prompt=negative_prompt, image=pose_image, num_inference_steps=num_steps, 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, num_steps, controlnet_conditioning_scale, seed):
    # 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, 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)
  
#   # 一時的に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:
  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")
      num_steps = gr.Slider(label="Number of steps", minimum=25, maximum=100, value=50, step=1)
      controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05)
      seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True,)
      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)

  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")
  
  # 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)