import os import warnings from typing import Callable, List, Optional, Union, Dict, Any import PIL.Image import trimesh import rembg import torch import numpy as np from huggingface_hub import hf_hub_download from diffusers.utils import BaseOutput import craftsman from craftsman.utils.config import ExperimentConfig, load_config class MeshPipelineOutput(BaseOutput): """ Output class for image pipelines. Args: images (`List[trimesh.Trimesh]` or `np.ndarray`) List of denoised trimesh meshes of length `batch_size` or a tuple of NumPy array with shape `((vertices, 3), (faces, 3)) of length `batch_size``. """ meshes: Union[List[trimesh.Trimesh], np.ndarray] class CraftsManPipeline(): """ Pipeline for text-guided image to image generation using CraftsMan(https://github.com/wyysf-98/CraftsMan). Args: feature_extractor ([`CLIPFeatureExtractor`]): Feature extractor for image pre-processing before being encoded. """ def __init__( self, device: str, cfg: ExperimentConfig, system, ): self.device = device self.cfg = cfg self.system = system @classmethod def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): r""" A simpler version that instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. The pipeline is set in evaluation mode (`model.eval()`) by default. """ # 1. Download the checkpoints and configs # use snapshot download here to get it working from from_pretrained if not os.path.isdir(pretrained_model_name_or_path): ckpt_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="model.ckpt", repo_type="model") config_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="config.yaml", repo_type="model") else: ckpt_path = os.path.join(pretrained_model_name_or_path, "model.ckpt") config_path = os.path.join(pretrained_model_name_or_path, "config.yaml") # 2. Load the model device = kwargs.get("device", "cuda" if torch.cuda.is_available() else "cpu") cfg = load_config(config_path) system = craftsman.find(cfg.system_type)(cfg.system) print(f"Restoring states from the checkpoint path at {ckpt_path} with config {cfg}") ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) system.load_state_dict( ckpt["state_dict"] if "state_dict" in ckpt else ckpt, ) system = system.to(device).eval() return cls( device=device, cfg=cfg, system=system ) def check_inputs( self, image, ): r""" Check if the inputs are valid. Raise an error if not. """ if isinstance(image, str): assert os.path.isfile(image) or image.startswith("http"), "Input image must be a valid URL or a file path." elif isinstance(image, (torch.Tensor, PIL.Image.Image)): raise ValueError("Input image must be a `torch.Tensor` or `PIL.Image.Image`.") def preprocess_image( self, images_pil: List[PIL.Image.Image], force: bool = False, background_color: List[int] = [255, 255, 255], foreground_ratio: float = 1.0, ): r""" Crop and remote the background of the input image Args: image_pil (`List[PIL.Image.Image]`): List of `PIL.Image.Image` objects representing the input image. force (`bool`, *optional*, defaults to `False`): Whether to force remove the background even if the image has an alpha channel. Returns: `List[PIL.Image.Image]`: List of `PIL.Image.Image` objects representing the preprocessed image. """ preprocessed_images = [] for i in range(len(images_pil)): image = images_pil[i] do_remove = True if image.mode == "RGBA" and image.getextrema()[3][0] < 255: # explain why current do not rm bg print("alhpa channl not enpty, skip remove background, using alpha channel as mask") background = PIL.Image.new("RGBA", image.size, (*background_color, 0)) image = PIL.Image.alpha_composite(background, image) do_remove = False do_remove = do_remove or force if do_remove: image = rembg.remove(image) # calculate the min bbox of the image alpha = image.split()[-1] image = image.crop(alpha.getbbox()) # Calculate the new size after rescaling new_size = tuple(int(dim * foreground_ratio) for dim in image.size) # Resize the image while maintaining the aspect ratio resized_image = image.resize(new_size) # Create a new image with the original size and white background padded_image = PIL.Image.new("RGBA", image.size, (*background_color, 0)) paste_position = ((image.width - resized_image.width) // 2, (image.height - resized_image.height) // 2) padded_image.paste(resized_image, paste_position) # expand image to 1:1 width, height = padded_image.size if width == height: preprocessed_images.append(padded_image) continue new_size = (max(width, height), max(width, height)) new_image = PIL.Image.new("RGBA", new_size, (*background_color, 1)) paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2) new_image.paste(padded_image, paste_position) preprocessed_images.append(new_image) return preprocessed_images @torch.no_grad() def __call__( self, image: Union[torch.FloatTensor, PIL.Image.Image, str], num_inference_steps: int = 50, guidance_scale: float = 7.5, eta: float = 0.0, num_meshes_per_prompt: Optional[int] = 1, output_type: Optional[str] = "trimesh", return_dict: bool = True, seed: Optional[int] = None, force_remove_background: bool = False, background_color: List[int] = [255, 255, 255], foreground_ratio: float = 0.95, mc_depth: int = 8, ): r""" Function invoked when calling the pipeline for generation. Args: image (`torch.FloatTensor` or `PIL.Image.Image`): `Image`, or tensor representing an image batch. The image will be encoded to its CLIP/DINO-v2 embedding which the DiT will be conditioned on. num_inference_steps (`int`, *optional*, defaults to 20): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. guidance_scale (`float`, *optional*, defaults to 10.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. eta (`float`, *optional*, defaults to 0.0): The eta parameter as defined in [DDIM](https://arxiv.org/abs/2010.02502). `eta` is a parameter that controls the amount of noise added to the latent space. It is only used with the DDIM scheduler and will be ignored for other schedulers. `eta` should be between [0, 1]. num_meshes_per_prompt (`int`, *optional*, defaults to 1): The number of meshes to generate per prompt. output_type (`str`, *optional*, defaults to `"trimesh"`): The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image`, `latents` or `np.array of v and f`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. seed (`int`, *optional*, defaults to `None`): Seed for the random number generator. Setting a seed will ensure reproducibility. force_remove_background (`bool`, *optional*, defaults to `False`): Whether to force remove the background even if the image has an alpha channel. foreground_ratio (`float`, *optional*, defaults to 1.0): The ratio of the foreground in the image. The foreground is the part of the image that is not the background. The foreground is resized to the size of the background image while maintaining the aspect ratio. The background is filled with black color. The foreground ratio should be between [0, 1]. mc_depth (`int`, *optional*, defaults to 8): The resolution of the Marching Cubes algorithm. The resolution is the number of cubes in the x, y, and z. 8 means 2^8 = 256 cubes in each dimension. The higher the resolution, the more detailed the mesh will be. Examples: Returns: [`~MeshPipelineOutput`] or `tuple`: [`~MeshPipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated meshes. """ # 0. Check inputs. Raise error if not correct self.check_inputs( image=image, ) # 1. Define call parameters if isinstance(image, torch.Tensor): batch_size = image.shape[0] elif isinstance(image, PIL.Image.Image) or isinstance(image, str): batch_size = 1 do_classifier_free_guidance = guidance_scale != 1.0 # 2. Preprocess input image if isinstance(image, torch.Tensor): images_pil = [TF.to_pil_image(image[i]) for i in range(image.shape[0])] elif isinstance(image, PIL.Image.Image): images_pil = [image] elif isinstance(image, str): if image.startswith("http"): import requests images_pil = [PIL.Image.open(requests.get(image, stream=True).raw)] else: images_pil = [PIL.Image.open(image)] images_pil = self.preprocess_image( images_pil, force=force_remove_background, background_color=background_color, foreground_ratio=foreground_ratio ) # 3. Inference latents = self.system.sample( {'image': images_pil}, sample_times = num_meshes_per_prompt, steps = num_inference_steps, guidance_scale = guidance_scale, eta = eta, seed = seed ) # 4. Post-processing if not output_type == "latent": mesh = [] for i, cur_latents in enumerate(latents): print(f"Generating mesh {i+1}/{num_meshes_per_prompt}") mesh_v_f, has_surface = self.system.shape_model.extract_geometry( cur_latents, octree_depth=mc_depth, extract_mesh_func="mc" ) if output_type == "trimesh": import trimesh cur_mesh = trimesh.Trimesh(vertices=mesh_v_f[0][0], faces=mesh_v_f[0][1]) mesh.append(cur_mesh) elif output_type == "np": mesh.append(mesh_v_f[0]) else: mesh = latents if not return_dict: return tuple(mesh) return MeshPipelineOutput(meshes=mesh)