Upload inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -0,0 +1,428 @@
+import argparse
+
+import numpy as np
+import torch
+from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
+from PIL import Image
+from torchvision import transforms
+from tqdm import tqdm
+from transformers import AutoModelForImageSegmentation
+
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
+
+
+from mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
+from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
+from mvadapter.utils import (
+    get_orthogonal_camera,
+    get_plucker_embeds_from_cameras_ortho,
+    make_image_grid,
+)
+
+
+def prepare_pipeline(
+    base_model,
+    vae_model,
+    unet_model,
+    lora_model,
+    adapter_path,
+    scheduler,
+    num_views,
+    device,
+    dtype,
+):
+    # Load vae and unet if provided
+    pipe_kwargs = {}
+    if vae_model is not None:
+        pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
+    if unet_model is not None:
+        pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
+
+    # Prepare pipeline
+    pipe: MVAdapterI2MVSDXLPipeline
+    pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
+
+    # Load scheduler if provided
+    scheduler_class = None
+    if scheduler == "ddpm":
+        scheduler_class = DDPMScheduler
+    elif scheduler == "lcm":
+        scheduler_class = LCMScheduler
+
+    pipe.scheduler = ShiftSNRScheduler.from_scheduler(
+        pipe.scheduler,
+        shift_mode="interpolated",
+        shift_scale=8.0,
+        scheduler_class=scheduler_class,
+    )
+    pipe.init_custom_adapter(num_views=num_views)
+    pipe.load_custom_adapter(
+        adapter_path, weight_name="mvadapter_i2mv_sdxl.safetensors"
+    )
+
+    pipe.to(device=device, dtype=dtype)
+    pipe.cond_encoder.to(device=device, dtype=dtype)
+
+    # load lora if provided
+    if lora_model is not None:
+        model_, name_ = lora_model.rsplit("/", 1)
+        pipe.load_lora_weights(model_, weight_name=name_)
+
+    # vae slicing for lower memory usage
+    pipe.enable_vae_slicing()
+
+    return pipe
+
+def remove_bg(image: Image.Image, net, transform, device, mask: Image.Image = None):
+    """
+    Applies a pre-existing mask to an image to make the background transparent.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        net: Pre-trained neural network (not used but kept for compatibility).
+        transform: Image transformation object (not used but kept for compatibility).
+        device: Device used for inference (not used but kept for compatibility).
+        mask (PIL.Image.Image, optional): The mask to use. Should be the same size
+                                 as the input image, with values between 0 and 255 (or 0-1).
+                                 If None, will return image with no changes.
+
+    Returns:
+        PIL.Image.Image: The modified image with transparent background.
+    """
+    if mask is None:
+      return image
+    
+    image_size = image.size
+    if mask.size != image_size:
+        mask = mask.resize(image_size) # Resizing the mask if it is not the same size as image
+    
+    image.putalpha(mask)
+    return image
+
+
+# def remove_bg(image, net, transform, device):
+#     image_size = image.size
+#     input_images = transform(image).unsqueeze(0).to(device)
+#     with torch.no_grad():
+#         preds = net(input_images)[0].sigmoid().cpu()
+#         #preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
+#     pred = preds[0].squeeze()
+#     pred_pil = transforms.ToPILImage()(pred)
+#     mask = pred_pil.resize(image_size)
+#     image.putalpha(mask)
+#     return image
+
+
+# def remove_bg(image: Image.Image, net, transform, device):
+#     """
+#     Applies a pre-existing mask to an image to make the background transparent.
+#     Args:
+#         image (PIL.Image.Image): The input image.
+#         net: Pre-trained neural network (not used but kept for compatibility).
+#         transform: Image transformation object (not used but kept for compatibility).
+#         device: Device used for inference (not used but kept for compatibility).
+#     Returns:
+#         PIL.Image.Image: The modified image with transparent background.
+#     """
+#     image_size = image.size
+#     input_images = transform(image).unsqueeze(0).to(device)
+    
+#     with torch.no_grad():
+#         preds = net(input_images)[-1].sigmoid().cpu()
+    
+#     pred = preds[0].squeeze()
+#     pred_pil = transforms.ToPILImage()(pred)
+    
+#     # Resize the mask to match the original image size
+#     mask = pred_pil.resize(image_size, Image.LANCZOS)
+    
+#     # Create a new image with the same size and mode as the original
+#     output_image = Image.new("RGBA", image_size)
+    
+#     # Apply the mask to the original image
+#     image.putalpha(mask)
+    
+#     # Composite the original image with the mask
+#     output_image.paste(image, (0, 0), image)
+    
+#     return output_image
+
+
+def remove_bg(image: Image.Image, net, transform, device, mask: np.ndarray = None):
+    """
+    Applies a pre-existing mask to an image to make the background transparent.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        net: Pre-trained neural network (not used but kept for compatibility).
+        transform: Image transformation object (not used but kept for compatibility).
+        device: Device used for inference (not used but kept for compatibility).
+        mask (np.ndarray, optional): The mask to use. Should be the same size
+                                     as the input image, with values between 0 and 255.
+                                     If None, will return image with no changes.
+
+    Returns:
+        PIL.Image.Image: The modified image with transparent background.
+    """
+    if mask is None:
+        return image
+
+    # Ensure the mask is in the correct format
+    if mask.ndim == 2:  # If mask is 2D (H, W)
+        mask = mask.astype(np.uint8)  # Ensure mask is uint8
+        mask = np.expand_dims(mask, axis=-1)  # Add channel dimension
+
+    # Convert the mask to PIL Image
+    mask_pil = Image.fromarray(mask.squeeze(2) * 255)  # Convert to binary mask
+
+    # Resize the mask to match the original image size
+    mask_pil = mask_pil.resize(image.size, Image.LANCZOS)
+
+    # Create a new image with the same size and mode as the original
+    output_image = Image.new("RGBA", image.size)
+
+    # Apply the mask to the original image
+    image.putalpha(mask_pil)
+
+    # Composite the original image with the mask
+    output_image.paste(image, (0, 0), image)
+
+    return output_image
+
+
+# def preprocess_image(image: Image.Image, height, width):
+    
+#     alpha = image[..., 3] > 0
+#     # alpha = image
+
+#     #if image.mode in ("RGBA", "LA"):
+#     #    image = np.array(image)
+#     #    alpha = image[..., 3]  # Extract the alpha channel
+#     #elif image.mode in ("RGB"):
+#     #    image = np.array(image)
+#         # Create default alpha for non-alpha images
+#     #    alpha = np.ones(image[..., 0].shape, dtype=np.uint8) * 255 # Create
+#     H, W = alpha.shape
+#     # get the bounding box of alpha
+#     y, x = np.where(alpha)
+#     y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+#     x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
+#     image_center = image[y0:y1, x0:x1]
+#     # resize the longer side to H * 0.9
+#     H, W, _ = image_center.shape
+#     if H > W:
+#         W = int(W * (height * 0.9) / H)
+#         H = int(height * 0.9)
+#     else:
+#         H = int(H * (width * 0.9) / W)
+#         W = int(width * 0.9)
+#     image_center = np.array(Image.fromarray(image_center).resize((W, H)))
+#     # pad to H, W
+#     start_h = (height - H) // 2
+#     start_w = (width - W) // 2
+#     image = np.zeros((height, width, 4), dtype=np.uint8)
+#     image[start_h : start_h + H, start_w : start_w + W] = image_center
+#     image = image.astype(np.float32) / 255.0
+#     image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+#     image = (image * 255).clip(0, 255).astype(np.uint8)
+#     image = Image.fromarray(image)
+
+#     return image
+
+def preprocess_image(image: Image.Image, height, width):
+    # Convert image to numpy array
+    image_np = np.array(image)
+
+    # Extract the alpha channel if present
+    if image_np.shape[-1] == 4:
+        alpha = image_np[..., 3] > 0  # Create a binary mask from the alpha channel
+    else:
+        alpha = np.ones(image_np[..., 0].shape, dtype=bool)  # Default to all true for RGB images
+
+    H, W = alpha.shape
+    # Get the bounding box of the alpha
+    y, x = np.where(alpha)
+    y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+    x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
+    image_center = image_np[y0:y1, x0:x1]
+
+    # Resize the longer side to H * 0.9
+    H, W, _ = image_center.shape
+    if H > W:
+        W = int(W * (height * 0.9) / H)
+        H = int(height * 0.9)
+    else:
+        H = int(H * (width * 0.9) / W)
+        W = int(width * 0.9)
+
+    image_center = np.array(Image.fromarray(image_center).resize((W, H)))
+
+    # Pad to H, W
+    start_h = (height - H) // 2
+    start_w = (width - W) // 2
+    padded_image = np.zeros((height, width, 4), dtype=np.uint8)
+    padded_image[start_h:start_h + H, start_w:start_w + W] = image_center
+
+    # Convert back to PIL Image
+    return Image.fromarray(padded_image)
+
+
+def run_pipeline(
+    pipe,
+    num_views,
+    text,
+    image,
+    height,
+    width,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    remove_bg_fn=None,
+    reference_conditioning_scale=1.0,
+    negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
+    lora_scale=1.0,
+    device="cuda",
+):
+    # Prepare cameras
+    cameras = get_orthogonal_camera(
+        elevation_deg=[0, 0, 0, 0, 0, 0],
+        distance=[1.8] * num_views,
+        left=-0.55,
+        right=0.55,
+        bottom=-0.55,
+        top=0.55,
+        azimuth_deg=[x - 90 for x in [0, 45, 90, 180, 270, 315]],
+        device=device,
+    )
+
+    plucker_embeds = get_plucker_embeds_from_cameras_ortho(
+        cameras.c2w, [1.1] * num_views, width
+    )
+    control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
+
+    # Prepare image
+    # reference_image = Image.open(image) if isinstance(image, str) else image
+    # if remove_bg_fn is not None:
+    #     reference_image = remove_bg_fn(reference_image)
+    #     reference_image = preprocess_image(reference_image, height, width)
+    # elif reference_image.mode == "RGBA":
+    #     reference_image = preprocess_image(reference_image, height, width)
+    reference_image = Image.open(image) if isinstance(image, str) else image
+    logging.info(f"Initial reference_image mode: {reference_image.mode}")
+    
+    if remove_bg_fn is not None:
+        logging.info("Using remove_bg_fn")
+        reference_image = remove_bg_fn(reference_image)
+        reference_image = preprocess_image(reference_image, height, width)
+    elif reference_image.mode == "RGBA":
+        logging.info("Image is RGBA, preprocessing directly")
+        reference_image = preprocess_image(reference_image, height, width)
+    
+    logging.info(f"Final reference_image mode: {reference_image.mode}")
+
+    pipe_kwargs = {}
+    if seed != -1 and isinstance(seed, int):
+        pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
+
+    images = pipe(
+        text,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_views,
+        control_image=control_images,
+        control_conditioning_scale=1.0,
+        reference_image=reference_image,
+        reference_conditioning_scale=reference_conditioning_scale,
+        negative_prompt=negative_prompt,
+        cross_attention_kwargs={"scale": lora_scale},
+        **pipe_kwargs,
+    ).images
+
+    return images, reference_image
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Models
+    parser.add_argument(
+        "--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
+    )
+    parser.add_argument(
+        "--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
+    )
+    parser.add_argument("--unet_model", type=str, default=None)
+    parser.add_argument("--scheduler", type=str, default=None)
+    parser.add_argument("--lora_model", type=str, default=None)
+    parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
+    parser.add_argument("--num_views", type=int, default=6)
+    # Device
+    parser.add_argument("--device", type=str, default="cuda")
+    # Inference
+    parser.add_argument("--image", type=str, required=True)
+    parser.add_argument("--text", type=str, default="high quality")
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--guidance_scale", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=-1)
+    parser.add_argument("--lora_scale", type=float, default=1.0)
+    parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default="watermark, ugly, deformed, noisy, blurry, low contrast",
+    )
+    parser.add_argument("--output", type=str, default="output.png")
+    # Extra
+    parser.add_argument("--remove_bg", action="store_true", help="Remove background")
+    args = parser.parse_args()
+
+    pipe = prepare_pipeline(
+        base_model=args.base_model,
+        vae_model=args.vae_model,
+        unet_model=args.unet_model,
+        lora_model=args.lora_model,
+        adapter_path=args.adapter_path,
+        scheduler=args.scheduler,
+        num_views=args.num_views,
+        device=args.device,
+        dtype=torch.float16,
+    )
+
+    if args.remove_bg:
+        birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet", trust_remote_code=True
+        )
+        birefnet.to(args.device)
+        transform_image = transforms.Compose(
+            [
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+        remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
+    else:
+        remove_bg_fn = None
+
+    images, reference_image = run_pipeline(
+        pipe,
+        num_views=args.num_views,
+        text=args.text,
+        image=args.image,
+        height=768,
+        width=768,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        seed=args.seed,
+        lora_scale=args.lora_scale,
+        reference_conditioning_scale=args.reference_conditioning_scale,
+        negative_prompt=args.negative_prompt,
+        device=args.device,
+        remove_bg_fn=remove_bg_fn,
+    )
+    make_image_grid(images, rows=1).save(args.output)
+    reference_image.save(args.output.rsplit(".", 1)[0] + "_reference.png")