cleanup denspose

brycedrennan · Apr 5, 2024 · 1ce7be1 · 1ce7be1
1 parent 8d5fedf
commit 1ce7be1
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 29 deletions.
diff --git a/imaginairy/config.py b/imaginairy/config.py
@@ -168,10 +168,17 @@ def __post_init__(self):
         defaults={"negative_prompt": DEFAULT_NEGATIVE_PROMPT},
     ),
     ModelWeightsConfig(
-        name="Redshift Diffusion",
-        aliases=["redshift-diffusion", "red", "redshift-diffusion-15", "red15"],
+        name="Miniaturus Potentia V1.2",
+        aliases=[
+            "miniaturupotentia",
+            "potentia",
+            "miniaturupotentia12",
+            "mp12",
+            "mp",
+            "potentia12",
+        ],
         architecture=MODEL_ARCHITECTURE_LOOKUP["sd15"],
-        weights_location="https://huggingface.co/nitrosocke/redshift-diffusion/tree/80837fe18df05807861ab91c3bad3693c9342e4c/",
+        weights_location="https://huggingface.co/dataautogpt3/Miniaturus_PotentiaV1.2/tree/7ef539518ad5ad591c45f0b920050883f7e51e83/",
         defaults={"negative_prompt": DEFAULT_NEGATIVE_PROMPT},
     ),
     # SDXL Weights

diff --git a/imaginairy/img_processors/control_modes.py b/imaginairy/img_processors/control_modes.py
@@ -139,17 +139,24 @@ def create_pose_map(img_t: "Tensor"):
     return pose_t
 
 
-def create_densepose_map(img_t: "torch.Tensor"):
-    from imaginairy.img_processors.densepose import generate_densepose_image
+def create_densepose_map(img_t: "Tensor") -> "Tensor":
     import torch
+
+    from imaginairy.img_processors.densepose import generate_densepose_image
+
     img_np = generate_densepose_image(img_t)
 
-    img_t = torch.tensor(img_np, dtype=torch.float) if not isinstance(img_np, torch.Tensor) else img_np.float()
+    img_t = (
+        torch.tensor(img_np, dtype=torch.float)
+        if not isinstance(img_np, torch.Tensor)
+        else img_np.float()
+    )
     img_t /= 255.0
     img_t = img_t.permute(2, 0, 1).unsqueeze(0)
 
     return img_t
 
+
 def make_noise_disk(H: int, W: int, C: int, F: int) -> "np.ndarray":
     import cv2
     import numpy as np

diff --git a/imaginairy/img_processors/densepose.py b/imaginairy/img_processors/densepose.py
@@ -50,13 +50,13 @@ def resize_image_with_pad_torch(
         img = F.interpolate(img, size=(H_target, W_target), mode="area")
 
     H_pad, W_pad = pad64(H_target), pad64(W_target)
-    print(f"image after resize but before padding: {img.shape}")
+    # print(f"image after resize but before padding: {img.shape}")
     img_padded = F.pad(img, (0, W_pad, 0, H_pad), mode=mode)
 
     def remove_pad(x):
-        print(
-            f"remove_pad: x.shape: {x.shape}. H_target: {H_target}, W_target: {W_target}"
-        )
+        # print(
+        #     f"remove_pad: x.shape: {x.shape}. H_target: {H_target}, W_target: {W_target}"
+        # )
         return safer_memory(x[:H_target, :W_target, ...])
 
     return img_padded, remove_pad
@@ -79,6 +79,7 @@ def HWC3(x: np.ndarray) -> np.ndarray:
         y = color * alpha + 255.0 * (1.0 - alpha)
         y = y.clip(0, 255).astype(np.uint8)
         return y
+    raise RuntimeError("unreachable")
 
 
 @lru_cache(maxsize=1)
@@ -174,7 +175,7 @@ def generate_densepose_image(
     detect_resolution=512,
     upscale_method="bicubic",
     cmap="viridis",
-    double_pass=True,
+    double_pass=False,
 ):
     assert_tensor_float_11_bchw(img)
     input_h, input_w = img.shape[-2:]
@@ -185,23 +186,23 @@ def generate_densepose_image(
         first_densepose_img_t = _np_to_t(first_densepose_img_np)
         # convert the densepose image into a mask (every color other than black is part of the mask)
         densepose_img_mask = first_densepose_img_t[0].sum(dim=0) > 0
-        print(f"Mask shape: {densepose_img_mask.shape}")
+        # print(f"Mask shape: {densepose_img_mask.shape}")
         # bbox = masks_to_boxes(densepose_img_mask.unsqueeze(0)).to(torch.uint8)
         # crop image by bbox
         bbox = mask_to_bbox(densepose_img_mask)
-        print(f"bbox: {bbox}")
+        # print(f"bbox: {bbox}")
 
         if bbox is None:
             densepose_np = first_densepose_img_np
         else:
             bbox = pad_bbox(bbox, max_height=input_h, max_width=input_w, pad=10)
-            print(f"padded bbox: {bbox}")
+            # print(f"padded bbox: {bbox}")
             bbox = square_bbox(bbox, max_height=input_h, max_width=input_w)
-            print(f"boxed bbox: {bbox}")
+            # print(f"boxed bbox: {bbox}")
             x0, y0, x1, y1 = bbox
 
             cropped_img = img[:, :, y0:y1, x0:x1]
-            print(f"cropped_img shape: {cropped_img.shape}")
+            # print(f"cropped_img shape: {cropped_img.shape}")
 
             densepose_np = _generate_densepose_image(
                 cropped_img,
@@ -210,10 +211,10 @@ def generate_densepose_image(
                 cmap,
                 adapt_viridis_bg=False,
             )
-            print(f"cropped densepose_np shape: {densepose_np.shape}")
-            print(
-                f"pasting into first_densepose_img_np shape: {first_densepose_img_np.shape} at {y0}:{y1}, {x0}:{x1}"
-            )
+            # print(f"cropped densepose_np shape: {densepose_np.shape}")
+            # print(
+            #     f"pasting into first_densepose_img_np shape: {first_densepose_img_np.shape} at {y0}:{y1}, {x0}:{x1}"
+            # )
             # paste denspose_np back into first_densepose_img_np using bbox
             first_densepose_img_np[y0:y1, x0:x1] = densepose_np
             densepose_np = first_densepose_img_np
@@ -239,14 +240,14 @@ def _generate_densepose_image(
 ) -> np.ndarray:
     assert_tensor_float_11_bchw(img)
     input_h, input_w = img.shape[-2:]
-    print(f"input_h: {input_h}, input_w: {input_w}")
+    # print(f"input_h: {input_h}, input_w: {input_w}")
     img, remove_pad = resize_image_with_pad_torch(
         img, detect_resolution, upscale_method
     )
     img = ((img + 1.0) * 127.5).to(torch.uint8)
     assert_tensor_uint8_255_bchw(img)
     H, W = img.shape[-2:]
-    print(f"reduced input img size (with padding): h{H}xw{W}")
+    # print(f"reduced input img size (with padding): h{H}xw{W}")
     hint_image_canvas = np.zeros([H, W], dtype=np.uint8)
     hint_image_canvas = np.tile(hint_image_canvas[:, :, np.newaxis], [1, 1, 3])
     densepose_model = get_densepose_model()
@@ -276,10 +277,10 @@ def _generate_densepose_image(
         hint_image[:, :, 0][hint_image[:, :, 0] == 0] = 68
         hint_image[:, :, 1][hint_image[:, :, 1] == 0] = 1
         hint_image[:, :, 2][hint_image[:, :, 2] == 0] = 84
-    print(f"hint_image shape: {hint_image.shape}")
+    # print(f"hint_image shape: {hint_image.shape}")
     detected_map = remove_pad(HWC3(hint_image))
-    print(f"detected_map shape (padding removed): {detected_map.shape}")
-
+    # print(f"detected_map shape (padding removed): {detected_map.shape}")
+    # print(f"Resizing detected_map to original size: {input_w}x{input_h}")
     # if map is smaller than input size, scale it up
     if detected_map.shape[0] < input_h or detected_map.shape[1] < input_w:
         detected_map = cv2.resize(
@@ -290,7 +291,7 @@ def _generate_densepose_image(
         detected_map = cv2.resize(
             detected_map, (input_w, input_h), interpolation=cv2.INTER_AREA
         )
-    print(f"detected_map shape (resized to original): {detected_map.shape}")
+    # print(f"detected_map shape (resized to original): {detected_map.shape}")
     return detected_map
 
 
@@ -391,7 +392,7 @@ def convert(
             arr = torch.tensor(box)[None, :]
         else:
             # avoid modifying the input box
-            arr = torch.from_numpy(np.asarray(box)).clone() if is_numpy else box.clone()
+            arr = torch.from_numpy(np.asarray(box)).clone() if is_numpy else box.clone()  # type: ignore
 
         assert to_mode not in [
             BoxMode.XYXY_REL,
@@ -432,7 +433,7 @@ def convert(
             arr[:, 0] += arr[:, 2] / 2.0
             arr[:, 1] += arr[:, 3] / 2.0
             angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
-            arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
+            arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)  # type: ignore
         else:
             if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
                 arr[:, 2] += arr[:, 0]
@@ -578,7 +579,7 @@ def densepose_chart_predictor_output_to_result(
     v = v.unsqueeze(0)
     boxes_xyxy_abs = boxes.clone()
     boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-    box_xywh = make_int_box(boxes_xywh_abs[0])
+    box_xywh = make_int_box(boxes_xywh_abs[0])  # type: ignore
 
     labels = resample_fine_and_coarse_segm_tensors_to_bbox(
         fine_segm, coarse_segm, box_xywh

diff --git a/tests/expected_output/test_control_images[densepose-create_densepose_map]_.png b/tests/expected_output/test_control_images[densepose-create_densepose_map]_.png