uuu

.idea/vcs.xml

@@ -2,7 +2,7 @@
 <project version="4">
   <component name="VcsDirectoryMappings">
     <mapping directory="" vcs="Git" />
-    <mapping directory="$PROJECT_DIR$/depth-fm" vcs="Git" />
+    <mapping directory="$PROJECT_DIR$/_image_gen_aux" vcs="Git" />
     <mapping directory="$PROJECT_DIR$/test_gradio" vcs="Git" />
   </component>
 </project>

app.py

@@ -11,6 +11,7 @@ from huggingface_hub import login
 import torch
 from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel
 from diffusers.utils import load_image
+from image_gen_aux import DepthPreprocessor
 
 # ----------------------------
 # Step 1: Download IP Adapter if not exists
@@ -70,9 +71,11 @@ pipe.init_ipadapter(
 @spaces.GPU
 def gui_generation(prompt,negative_prompt, ref_img, guidance_scale, ipadapter_scale):
     ref_img = load_image(ref_img.name).convert('RGB')
+    image = load_image(ref_img.name)
+
+    depth_preprocessor = DepthPreprocessor.from_pretrained("depth-anything/Depth-Anything-V2-Large-hf").to("cuda")
+    control_image = depth_preprocessor(image, invert=True)[0].convert("RGB")
 
-    control_image = load_image(
-        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/marigold/marigold_einstein_lcm_depth.png")
     generator = torch.Generator(device="cpu").manual_seed(0)
     image = pipe(
         width=1024,

image_gen_aux/__init__.py

@@ -0,0 +1,81 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from .utils import (
+    IMAGE_AUX_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+# Lazy Import based on
+# https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
+
+# When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names,
+# and is used to defer the actual importing for when the objects are requested.
+# This way `import image_gen_aux` provides the names in the namespace without actually importing anything (and especially none of the backends).
+
+_import_structure = {
+    "upscalers": [],
+    "preprocessors": ["LineArtStandardPreprocessor"],
+    "utils": [
+        "OptionalDependencyNotAvailable",
+        "is_torch_available",
+        "logging",
+    ],
+}
+
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["upscalers"].extend(["UpscaleWithModel"])
+
+    _import_structure["preprocessors"].extend(["DepthPreprocessor", "LineArtPreprocessor", "TeedPreprocessor"])
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    from .preprocessors import LineArtStandardPreprocessor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .preprocessors import LineArtPreprocessor, TeedPreprocessor
+        from .upscalers import UpscaleWithModel
+
+    try:
+        if not (is_torch_available() and is_transformers_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .preprocessors import DepthPreprocessor
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/image_processor.py

@@ -0,0 +1,233 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Union
+
+import cv2
+import numpy as np
+import PIL.Image
+import torch
+from PIL import Image
+
+
+class ImageMixin:
+    """
+    A mixin class for converting images between different formats: PIL, NumPy, and PyTorch tensors.
+
+    This class provides methods to:
+    - Convert a PIL image or a NumPy array to a PyTorch tensor.
+    - Post-process a PyTorch tensor image and convert it to the specified return type.
+    - Convert a PIL image or a list of PIL images to NumPy arrays.
+    - Convert a NumPy image to a PyTorch tensor.
+    - Convert a PyTorch tensor to a NumPy image.
+    - Convert a NumPy image or a batch of images to a PIL image.
+    """
+
+    def convert_image_to_tensor(
+        self, image: Union[PIL.Image.Image, np.ndarray, List[PIL.Image.Image]], normalize: bool = True
+    ) -> torch.Tensor:
+        """
+        Convert a PIL image or a NumPy array to a PyTorch tensor.
+
+        Args:
+            image (Union[PIL.Image.Image, np.ndarray]): The input image, either as a PIL image, a NumPy array or a list of
+                PIL images.
+
+        Returns:
+            torch.Tensor: The converted image as a PyTorch tensor.
+        """
+        if isinstance(image, (PIL.Image.Image, list)):
+            # We expect that if it is a list, it only should contain pillow images
+            if isinstance(image, list):
+                for single_image in image:
+                    if not isinstance(single_image, PIL.Image.Image):
+                        raise ValueError("All images in the list must be Pillow images.")
+
+            image = self.pil_to_numpy(image, normalize)
+
+        return self.numpy_to_pt(image)
+
+    def post_process_image(self, image: torch.Tensor, return_type: str):
+        """
+        Post-process a PyTorch tensor image and convert it to the specified return type.
+
+        Args:
+            image (torch.Tensor): The input image as a PyTorch tensor.
+            return_type (str): The desired return type, either "pt" for PyTorch tensor, "np" for NumPy array, or "pil" for PIL image.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray, List[PIL.Image.Image]]: The post-processed image in the specified return type.
+        """
+        if return_type == "pt":
+            return image
+
+        image = self.pt_to_numpy(image)
+        if return_type == "np":
+            return image
+
+        image = self.numpy_to_pil(image)
+        return image
+
+    @staticmethod
+    def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image], normalize: bool = True) -> np.ndarray:
+        """
+        Convert a PIL image or a list of PIL images to NumPy arrays.
+
+        Args:
+            images (Union[List[PIL.Image.Image], PIL.Image.Image]): The input image(s) as PIL image(s).
+
+        Returns:
+            np.ndarray: The converted image(s) as a NumPy array.
+        """
+        if not isinstance(images, list):
+            images = [images]
+
+        if normalize:
+            images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+        else:
+            images = [np.array(image).astype(np.float32) for image in images]
+
+        images = np.stack(images, axis=0)
+
+        return images
+
+    @staticmethod
+    def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
+        """
+        Convert a NumPy image to a PyTorch tensor.
+
+        Args:
+            images (np.ndarray): The input image(s) as a NumPy array.
+
+        Returns:
+            torch.Tensor: The converted image(s) as a PyTorch tensor.
+        """
+        if images.ndim == 3:
+            images = images[..., None]
+        images = torch.from_numpy(images.transpose(0, 3, 1, 2)).float()
+
+        return images
+
+    @staticmethod
+    def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
+        """
+        Convert a PyTorch tensor to a NumPy image.
+
+        Args:
+            images (torch.Tensor): The input image(s) as a PyTorch tensor.
+
+        Returns:
+            np.ndarray: The converted image(s) as a NumPy array.
+        """
+        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+        return images
+
+    @staticmethod
+    def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
+        """
+        Convert a NumPy image or a batch of images to a PIL image.
+
+        Args:
+            images (np.ndarray): The input image(s) as a NumPy array.
+
+        Returns:
+            List[PIL.Image.Image]: The converted image(s) as PIL images.
+        """
+        if images.ndim == 3:
+            images = images[None, ...]
+        images = (images * 255).round().astype("uint8")
+        if images.shape[-1] == 1:
+            # special case for grayscale (single channel) images
+            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+        else:
+            pil_images = [Image.fromarray(image) for image in images]
+
+        return pil_images
+
+    @staticmethod
+    def scale_image(image: torch.Tensor, scale: float, mutiple_factor: int = 8) -> torch.Tensor:
+        """
+        Scales an image while maintaining aspect ratio and ensuring dimensions are multiples of `multiple_factor`.
+
+        Args:
+            image (`torch.Tensor`): The input image tensor of shape (batch, channels, height, width).
+            scale (`float`): The scaling factor applied to the image dimensions.
+            multiple_factor (`int`, *optional*, defaults to 8): The factor by which the new dimensions should be divisible.
+
+        Returns:
+            `torch.Tensor`: The scaled image tensor.
+        """
+
+        if scale == 1.0:
+            return image, scale
+
+        _batch, _channels, height, width = image.shape
+
+        # Calculate new dimensions while maintaining aspect ratio
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+
+        # Ensure new dimensions are multiples of mutiple_factor
+        new_height = (new_height // mutiple_factor) * mutiple_factor
+        new_width = (new_width // mutiple_factor) * mutiple_factor
+
+        # if the final height and widht changed because of the multiple_factor, we need to set the scale too
+        scale = new_height / height
+
+        # Resize the image using the calculated dimensions
+        resized_image = torch.nn.functional.interpolate(
+            image, size=(new_height, new_width), mode="bilinear", align_corners=False
+        )
+
+        return resized_image, scale
+
+    @staticmethod
+    def resize_numpy_image(image: np.ndarray, scale: float, multiple_factor: int = 8) -> np.ndarray:
+        """
+        Resizes a NumPy image while maintaining aspect ratio and ensuring dimensions are multiples of `multiple_factor`.
+
+        Args:
+            image (`np.ndarray`): The input image array of shape (height, width, channels) or (height, width) for grayscale.
+            scale (`float`): The scaling factor applied to the image dimensions.
+            multiple_factor (`int`, *optional*, defaults to 8): The factor by which the new dimensions should be divisible.
+
+        Returns:
+            `np.ndarray`: The resized image array.
+        """
+        if len(image.shape) == 3:  # Single image without batch dimension
+            image = np.expand_dims(image, axis=0)
+
+        batch_size, height, width, channels = image.shape
+
+        # Calculate new dimensions while maintaining aspect ratio
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+
+        # Ensure new dimensions are multiples of multiple_factor
+        new_height = (new_height // multiple_factor) * multiple_factor
+        new_width = (new_width // multiple_factor) * multiple_factor
+
+        # if the final height and widht changed because of the multiple_factor, we need to set the scale too
+        scale = new_height / height
+
+        # Resize each image in the batch
+        resized_images = []
+        for i in range(batch_size):
+            resized_image = cv2.resize(image[i], (new_width, new_height), interpolation=cv2.INTER_LINEAR)
+            resized_images.append(resized_image)
+
+        # Stack resized images back into a single array
+        resized_images = np.stack(resized_images, axis=0)
+
+        return resized_images, scale

image_gen_aux/modeling_utils.py

@@ -0,0 +1,90 @@
+import itertools
+from typing import List, Tuple
+
+import torch
+from torch import Tensor
+
+
+def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
+    """
+    Gets the device of a PyTorch module's parameters or buffers.
+
+    Args:
+        parameter (`torch.nn.Module`): The PyTorch module from which to get the device.
+
+    Returns:
+        `torch.device`: The device of the module's parameters or buffers.
+    """
+    try:
+        parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
+        return next(parameters_and_buffers).device
+    except StopIteration:
+        # For torch.nn.DataParallel compatibility in PyTorch 1.5
+        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+            return tuples
+
+        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+        first_tuple = next(gen)
+        return first_tuple[1].device
+
+
+def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
+    """
+    Gets the data type of a PyTorch module's parameters or buffers.
+
+    Args:
+        parameter (`torch.nn.Module`): The PyTorch module from which to get the data type.
+
+    Returns:
+        `torch.dtype`: The data type of the module's parameters or buffers.
+    """
+    try:
+        params = tuple(parameter.parameters())
+        if len(params) > 0:
+            return params[0].dtype
+
+        buffers = tuple(parameter.buffers())
+        if len(buffers) > 0:
+            return buffers[0].dtype
+
+    except StopIteration:
+        # For torch.nn.DataParallel compatibility in PyTorch 1.5
+        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+            return tuples
+
+        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+        first_tuple = next(gen)
+        return first_tuple[1].dtype
+
+
+class ModelMixin(torch.nn.Module):
+    """
+    Provides convenient properties to access the device and data type
+    of a PyTorch module.
+
+    By inheriting from this class, your custom PyTorch modules can access these properties
+    without manual retrieval of device and data type information.
+
+    These properties assume that all module parameters and buffers reside
+    on the same device and have the same data type, respectively.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+        device).
+        """
+        return get_parameter_device(self)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
+        """
+        return get_parameter_dtype(self)

image_gen_aux/preprocessors/README.md

@@ -0,0 +1,54 @@
+# Preprocessors
+
+Preprocessors in this context refer to the application of machine learning models to images or groups of images as a preliminary step for other tasks. A common use case is to feed preprocessed images into a controlnet to guide the generation of a diffusion model. Examples of preprocessors include depth maps, normals, and edge detection.
+
+## Supported preprocessors
+
+This is a list of the currently supported preprocessors.
+
+* [DepthPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/depth/README.md)
+* [LineArtPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/lineart/README.md)
+* [LineArtStandardPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/lineart_standard/README.md)
+* [TeedPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/teed/README.md)
+
+## General preprocessor usage
+
+```python
+from image_gen_aux import LineArtPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/simple_upscale/hippowaffle.png"
+)
+
+lineart_preprocessor = LineArtPreprocessor.from_pretrained("OzzyGT/lineart").to("cuda")
+image = lineart_preprocessor(input_image)
+image.save("lineart.png")
+```
+
+Use the Hugging Face Hub model id, a URL, or a local path to load a model. You can also load pickle checkpoints, but it's not recommended due to the [vulnerabilities](https://docs.python.org/3/library/pickle.html) of pickled files. Loading [safetensor](https://hf.co/docs/safetensors/index) files is a more secure option.
+
+If the checkpoint has custom imports (high risk!), you'll need to set the `weights_only` argument to `False`.
+
+You can also specify a `filename` and a `subfolder` if needed.
+
+```python
+lineart_preprocessor = LineArtPreprocessor.from_pretrained("lllyasviel/Annotators", filename="sk_model.pth", weights_only=False).to("cuda")
+```
+
+### List of safetensors checkpoints
+
+This is the current list of safetensor checkpoints available on the Hub.
+
+|Preprocessor|Repository|Author|
+|---|---|---|
+|Depth Anything V2 Small|depth-anything/Depth-Anything-V2-Small-hf|<https://depth-anything-v2.github.io/>|
+|Depth Anything V2 Base|depth-anything/Depth-Anything-V2-Base-hf|<https://depth-anything-v2.github.io/>|
+|Depth Anything V2 Large|depth-anything/Depth-Anything-V2-Large-hf|<https://depth-anything-v2.github.io/>|
+|LineArt|OzzyGT/lineart|[Caroline Chan](https://github.com/carolineec)|
+|Teed|OzzyGT/teed|<https://github.com/xavysp/TEED>|
+|ZoeDepth NYU|Intel/zoedepth-nyu|<https://github.com/isl-org/ZoeDepth>|
+|ZoeDepth KITTI|Intel/zoedepth-kitti|<https://github.com/isl-org/ZoeDepth>|
+|ZoeDepth NYU and KITTI|Intel/zoedepth-nyu-kitti|<https://github.com/isl-org/ZoeDepth>|
+
+If you own the model and want us to change the repository to your name/organization please open an issue.

image_gen_aux/preprocessors/__init__.py

@@ -0,0 +1,66 @@
+from typing import TYPE_CHECKING
+
+from ..utils import (
+    IMAGE_AUX_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_import_structure = {
+    "lineart": [],
+    "depth": [],
+    "teed": [],
+    "lineart_standard": ["LineArtStandardPreprocessor"],
+}
+
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["lineart"] = ["LineArtPreprocessor"]
+    _import_structure["teed"] = ["TeedPreprocessor"]
+
+try:
+    if not (is_torch_available() and is_transformers_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["depth"] = [
+        "DepthPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    from .lineart_standard import LineArtStandardPreprocessor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .lineart import LineArtPreprocessor
+        from .teed import TeedPreprocessor
+
+    try:
+        if not (is_torch_available() and is_transformers_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .depth import DepthPreprocessor
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/preprocessors/depth/README.md

@@ -0,0 +1,45 @@
+# DEPTH
+
+Monocular depth estimation is a computer vision task that involves predicting the depth information of a scene from a single image. In other words, it is the process of estimating the distance of objects in a scene from a single camera viewpoint.
+
+Monocular depth estimation has various applications, including 3D reconstruction, augmented reality, autonomous driving, and robotics. It is a challenging task as it requires the model to understand the complex relationships between objects in the scene and the corresponding depth information, which can be affected by factors such as lighting conditions, occlusion, and texture.
+
+## Usage
+
+```python
+from image_gen_aux import DepthPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/depth/coffee_ship.png")
+
+depth_preprocessor = DepthPreprocessor.from_pretrained("depth-anything/Depth-Anything-V2-Large-hf").to("cuda")
+image = depth_preprocessor(input_image)[0]
+image.save("depth.png")
+```
+
+## Models
+
+The Depth Preprocessor supports any depth estimation model that the `transformers` library supports that doesn't have a fixed image size restriction, but we mainly recommend and ensure the correct functionality for these models:
+
+|Model|License|Project Page|
+|---|---|---|
+|Depth Anything V2|CC-BY-NC-4.0|<https://depth-anything-v2.github.io/>|
+|ZoeDepth|MIT|<https://github.com/isl-org/ZoeDepth>|
+
+Each model has different variations:
+
+### Depth Anything V2
+
+|Variation|Repo ID|
+|---|---|
+|Small|depth-anything/Depth-Anything-V2-Small-hf|
+|Base|depth-anything/Depth-Anything-V2-Base-hf|
+|Large|depth-anything/Depth-Anything-V2-Large-hf|
+
+### ZoeDepth
+
+|Variation|Repo ID|
+|---|---|
+|NYU|Intel/zoedepth-nyu|
+|KITTI|Intel/zoedepth-kitti|
+|NYU and KITTI|Intel/zoedepth-nyu-kitti|

image_gen_aux/preprocessors/depth/__init__.py

@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    IMAGE_AUX_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available() and is_transformers_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["depth_preprocessor"] = [
+        "DepthPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not (is_torch_available() and is_transformers_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .depth_preprocessor import (
+            DepthPreprocessor,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/preprocessors/depth/depth_preprocessor.py

@@ -0,0 +1,70 @@
+import os
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForDepthEstimation
+
+from ...image_processor import ImageMixin
+from ..preprocessor import Preprocessor
+
+
+class DepthPreprocessor(Preprocessor, ImageMixin):
+    """Preprocessor specifically designed for monocular depth estimation.
+
+    This class inherits from both `Preprocessor` and `ImageMixin`. Please refer to each
+    one to get more information.
+    """
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path: Union[str, os.PathLike]):
+        model = AutoModelForDepthEstimation.from_pretrained(pretrained_model_or_path)
+
+        return cls(model)
+
+    @torch.inference_mode
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]],
+        resolution_scale: float = 1.0,
+        invert: bool = False,
+        batch_size: int = 1,
+        return_type: str = "pil",
+    ):
+        if not isinstance(image, torch.Tensor):
+            image = self.convert_image_to_tensor(image).to(self.model.device)
+
+        image, resolution_scale = self.scale_image(image, resolution_scale)
+
+        processed_images = []
+
+        for i in range(0, len(image), batch_size):
+            batch = image[i : i + batch_size].to(self.model.device)
+
+            predicted_depth = self.model(batch).predicted_depth
+
+            # depth models returns only batch, height and width, so we add the channel
+            predicted_depth = predicted_depth.unsqueeze(1)
+
+            # models like depth anything can return a different size image
+            if batch.shape[2] != predicted_depth.shape[2] or batch.shape[3] != predicted_depth.shape[3]:
+                predicted_depth = F.interpolate(
+                    predicted_depth, size=(batch.shape[2], batch.shape[3]), mode="bilinear", align_corners=False
+                )
+
+            if invert:
+                predicted_depth = 255 - predicted_depth
+            processed_images.append(predicted_depth.cpu())
+        predicted_depth = torch.cat(processed_images, dim=0)
+
+        if resolution_scale != 1.0:
+            predicted_depth, _ = self.scale_image(predicted_depth, 1 / resolution_scale)
+
+        predicted_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
+        predicted_depth = predicted_depth.clamp(0, 1)
+
+        image = self.post_process_image(predicted_depth, return_type)
+
+        return image

image_gen_aux/preprocessors/lineart/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Caroline Chan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

image_gen_aux/preprocessors/lineart/README.md

@@ -0,0 +1,27 @@
+# Line Art
+
+Line Art is based on the original repository, [Informative Drawings: Learning to generate line drawings that convey geometry and semantics](https://github.com/carolineec/informative-drawings).
+
+From the project page [summary](https://carolineec.github.io/informative_drawings/):
+
+*Our method approaches line drawing generation as an unsupervised image translation problem which uses various losses to assess the information communicated in a line drawing. Our key idea is to view the problem as an encoding through a line drawing and to maximize the quality of this encoding through explicit geometry, semantic, and appearance decoding objectives. This evaluation is performed by deep learning methods which decode depth, semantics, and appearance from line drawings. The aim is for the extracted depth and semantic information to match the scene geometry and semantics of the input photographs.*
+
+## Usage
+
+```python
+from image_gen_aux import LineArtPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/simple_upscale/hippowaffle.png"
+)
+
+lineart_preprocessor = LineArtPreprocessor.from_pretrained("OzzyGT/lineart").to("cuda")
+image = lineart_preprocessor(input_image)[0]
+image.save("lineart.png")
+```
+
+## Additional resources
+
+* [Project page](https://carolineec.github.io/informative_drawings/)
+* [Paper](https://arxiv.org/abs/2203.12691)

image_gen_aux/preprocessors/lineart/__init__.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from ...utils import IMAGE_AUX_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["lineart_preprocessor"] = [
+        "LineArtPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .lineart_preprocessor import (
+            LineArtPreprocessor,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/preprocessors/lineart/lineart_preprocessor.py

@@ -0,0 +1,101 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from safetensors.torch import load_file
+
+from ...image_processor import ImageMixin
+from ...utils import SAFETENSORS_FILE_EXTENSION, get_model_path
+from ..preprocessor import Preprocessor
+from .model import Generator
+
+
+class LineArtPreprocessor(Preprocessor, ImageMixin):
+    """Preprocessor specifically designed for converting images to line art.
+
+    This class inherits from both `Preprocessor` and `ImageMixin`. Please refer to each
+    one to get more information.
+    """
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_or_path: Union[str, os.PathLike],
+        filename: str = None,
+        subfolder: str = None,
+        weights_only: bool = True,
+    ) -> Generator:
+        model_path = get_model_path(pretrained_model_or_path, filename, subfolder)
+
+        file_extension = os.path.basename(model_path).split(".")[-1]
+        if file_extension == SAFETENSORS_FILE_EXTENSION:
+            state_dict = load_file(model_path, device="cpu")
+        else:
+            state_dict = torch.load(model_path, map_location=torch.device("cpu"), weights_only=weights_only)
+
+        model = Generator()
+        model.load_state_dict(state_dict)
+
+        return cls(model)
+
+    @torch.inference_mode
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]],
+        resolution_scale: float = 1.0,
+        invert: bool = True,
+        batch_size: int = 1,
+        return_type: str = "pil",
+    ):
+        """Preprocesses an image and generates line art using the pre-trained model.
+
+        Args:
+            image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]]`): Input image as PIL Image,
+                NumPy array, PyTorch tensor format or a list of PIL Images
+            resolution_scale (`float`, optional, defaults to 1.0): Scale factor for image resolution during
+                preprocessing and post-processing. Defaults to 1.0 for no scaling.
+            invert (`bool`, *optional*, defaults to True): Inverts the generated image if True (white or black background).
+            batch_size (`int`, *optional*, defaults to 1): The number of images to process in each batch.
+            return_type (`str`, *optional*, defaults to "pil"): The desired return type, either "pt" for PyTorch tensor, "np" for NumPy array,
+                or "pil" for PIL image.
+
+        Returns:
+            `Union[PIL.Image.Image, np.ndarray, torch.Tensor]`: The generated line art in the
+                specified output format.
+        """
+        if not isinstance(image, torch.Tensor):
+            image = self.convert_image_to_tensor(image)
+
+        image, resolution_scale = self.scale_image(image, resolution_scale)
+
+        processed_images = []
+
+        for i in range(0, len(image), batch_size):
+            batch = image[i : i + batch_size].to(self.model.device)
+            lineart = self.model(batch)
+            if invert:
+                lineart = 255 - lineart
+            processed_images.append(lineart.cpu())
+        lineart = torch.cat(processed_images, dim=0)
+
+        if resolution_scale != 1.0:
+            lineart, _ = self.scale_image(lineart, 1 / resolution_scale)
+
+        image = self.post_process_image(lineart, return_type)
+
+        return image

image_gen_aux/preprocessors/lineart/model.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2022 Caroline Chan
+#
+# MIT License
+import torch.nn as nn
+
+from ...modeling_utils import ModelMixin
+
+
+norm_layer = nn.InstanceNorm2d
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features: int):
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            norm_layer(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            norm_layer(in_features),
+        ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+
+class Generator(ModelMixin):
+    def __init__(self, input_nc: int = 3, output_nc: int = 1, n_residual_blocks: int = 3, sigmoid: bool = True):
+        super(Generator, self).__init__()
+
+        # Initial convolution block
+        model0 = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, 7), norm_layer(64), nn.ReLU(inplace=True)]
+        self.model0 = nn.Sequential(*model0)
+
+        # Downsampling
+        model1 = []
+        in_features = 64
+        out_features = in_features * 2
+        for _ in range(2):
+            model1 += [
+                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                norm_layer(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features * 2
+        self.model1 = nn.Sequential(*model1)
+
+        model2 = []
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model2 += [ResidualBlock(in_features)]
+        self.model2 = nn.Sequential(*model2)
+
+        # Upsampling
+        model3 = []
+        out_features = in_features // 2
+        for _ in range(2):
+            model3 += [
+                nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+                norm_layer(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features // 2
+        self.model3 = nn.Sequential(*model3)
+
+        # Output layer
+        model4 = [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7)]
+        if sigmoid:
+            model4 += [nn.Sigmoid()]
+
+        self.model4 = nn.Sequential(*model4)
+
+    def forward(self, x):
+        out = self.model0(x)
+        out = self.model1(out)
+        out = self.model2(out)
+        out = self.model3(out)
+        out = self.model4(out)
+
+        return out

image_gen_aux/preprocessors/lineart_standard/README.md

@@ -0,0 +1,23 @@
+# Line Art Standard
+
+Line Art Standard was copied from the [comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux) repository.
+This preprocessor applies a Gaussian blur to the image to reduce noise and detail, then calculates the intensity difference between the blurred and original images to highlight edges.
+
+## Usage
+
+```python
+from image_gen_aux import LineArtStandardPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/simple_upscale/hippowaffle.png"
+)
+
+lineart_standard_preprocessor = LineArtStandardPreprocessor()
+image = lineart_standard_preprocessor(input_image)[0]
+image.save("lineart_standard.png")
+```
+
+## Additional resources
+
+* [Original implementation](https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/src/custom_controlnet_aux/lineart_standard/__init__.py)

image_gen_aux/preprocessors/lineart_standard/__init__.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from ...utils import IMAGE_AUX_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["lineart_standard_preprocessor"] = [
+        "LineArtStandardPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .lineart_standard_preprocessor import (
+            LineArtStandardPreprocessor,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/preprocessors/lineart_standard/lineart_standard_preprocessor.py

@@ -0,0 +1,82 @@
+from typing import List, Union
+
+import cv2
+import numpy as np
+import PIL.Image
+
+from ...image_processor import ImageMixin
+
+
+class LineArtStandardPreprocessor(ImageMixin):
+    """Preprocessor specifically designed for converting images to line art standard.
+
+    This class inherits from both `Preprocessor` and `ImageMixin`. Please refer to each
+    one to get more information.
+    """
+
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, List[PIL.Image.Image]],
+        resolution_scale: float = 1.0,
+        invert: bool = False,
+        gaussian_sigma=6.0,
+        intensity_threshold=8,
+        return_type: str = "pil",
+    ):
+        """Preprocesses an image and generates line art (standard) using the opencv.
+
+        Args:
+            image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]]`): Input image as PIL Image,
+                NumPy array, PyTorch tensor format or a list of PIL Images
+            resolution_scale (`float`, optional, defaults to 1.0): Scale factor for image resolution during
+                preprocessing and post-processing. Defaults to 1.0 for no scaling.
+            invert (`bool`, *optional*, defaults to True): Inverts the generated image if True (white or black background).
+            gaussian_sigma (float, optional): Sigma value for Gaussian blur. Defaults to 6.0.
+            intensity_threshold (int, optional): Threshold for intensity clipping. Defaults to 8.
+            return_type (`str`, *optional*, defaults to "pil"): The desired return type, either "pt" for PyTorch tensor, "np" for NumPy array,
+                or "pil" for PIL image.
+
+        Returns:
+            `Union[PIL.Image.Image, np.ndarray, torch.Tensor]`: The generated line art (standard) in the
+                specified output format.
+        """
+        if not isinstance(image, np.ndarray):
+            image = np.array(image).astype(np.float32)
+
+        # check if image has batch, if not, add it
+        if len(image.shape) == 3:
+            image = image[None, ...]
+
+        image, resolution_scale = (
+            self.resize_numpy_image(image, resolution_scale) if resolution_scale != 1.0 else image
+        )
+
+        batch_size, height, width, _channels = image.shape
+        processed_images = np.empty((batch_size, height, width), dtype=np.uint8)
+
+        # since we're using just cv2, we can't do batch processing
+        for i in range(batch_size):
+            gaussian = cv2.GaussianBlur(image[i], (0, 0), gaussian_sigma)
+            intensity = np.min(gaussian - image[i], axis=2).clip(0, 255)
+            intensity /= max(16, np.median(intensity[intensity > intensity_threshold]))
+            intensity *= 127
+            edges = intensity.clip(0, 255).astype(np.uint8)
+
+            processed_images[i] = edges
+
+        if invert:
+            processed_images = 255 - processed_images
+
+        processed_images = processed_images[..., None]
+
+        if resolution_scale != 1.0:
+            processed_images, _ = self.resize_numpy_image(processed_images, 1 / resolution_scale)
+            processed_images = processed_images[..., None]  # cv2 resize removes the channel dimension if grayscale
+
+        if return_type == "np":
+            return processed_images
+
+        processed_images = np.transpose(processed_images, (0, 3, 1, 2))
+        processed_images = [PIL.Image.fromarray(image.squeeze(), mode="L") for image in processed_images]
+
+        return processed_images

image_gen_aux/preprocessors/preprocessor.py

@@ -0,0 +1,75 @@
+from abc import ABC, abstractmethod
+from typing import Union
+
+import numpy as np
+import PIL.Image
+import torch
+
+
+class Preprocessor(ABC):
+    """
+    This abstract base class defines the interface for image preprocessors.
+
+    Subclasses should implement the abstract methods `from_pretrained` and
+    `__call__` to provide specific loading and preprocessing logic for their
+    respective models.
+
+    Args:
+        model (`nn.Module`): The torch model to use.
+    """
+
+    def __init__(self, model):
+        self.model = model
+
+    def to(self, device):
+        """
+        Moves the underlying model to the specified device
+        (e.g., CPU or GPU).
+
+        Args:
+            device (`torch.device`): The target device.
+
+        Returns:
+            `Preprocessor`: The preprocessor object itself (for method chaining).
+        """
+        self.model = self.model.to(device)
+        return self
+
+    @abstractmethod
+    def from_pretrained(self):
+        """
+        This abstract method defines how the preprocessor loads pre-trained
+        weights or configurations specific to the model it supports. Subclasses
+        must implement this method to handle model-specific loading logic.
+
+        This method might download pre-trained weights from a repository or
+        load them from a local file depending on the model's requirements.
+        """
+        pass
+
+    @abstractmethod
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        resolution_scale: float = 1.0,
+        invert: bool = True,
+        return_type: str = "pil",
+    ):
+        """
+        Preprocesses an image for use with the underlying model.
+
+        Args:
+            image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): Input image as PIL Image,
+                NumPy array, or PyTorch tensor format.
+            resolution_scale (`float`, optional, defaults to 1.0): Scale factor for image resolution during
+            resolution_scale (`float`, *optional*, defaults to 1.0): Scale factor for image resolution during
+                preprocessing and post-processing. Defaults to 1.0 for no scaling.
+            invert (`bool`, *optional*, defaults to True): Inverts the generated image if True.
+            return_type (`str`, *optional*, defaults to "pil"): The desired return type, either "pt" for PyTorch tensor,
+                "np" for NumPy array, or "pil" for PIL image.
+
+        Returns:
+            `Union[PIL.Image.Image, torch.Tensor]`: The preprocessed image in the
+                specified format.
+        """
+        pass

image_gen_aux/preprocessors/teed/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Xavier Soria Poma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

image_gen_aux/preprocessors/teed/README.md

@@ -0,0 +1,26 @@
+# TEED: Tiny and Efficient Model for the Edge Detection Generalization
+
+Tiny and Efficient Edge Detector (TEED) is a light convolutional neural network with only 58K parameters, less than 0.2% of the state-of-the-art models. Training on the [BIPED](https://www.kaggle.com/datasets/xavysp/biped) dataset takes less than 30 minutes,
+with each epoch requiring less than 5 minutes. Our proposed model is easy to train and it quickly converges within very first few
+epochs, while the predicted edge-maps are crisp and of high quality, see image above.
+[This paper has been accepted by ICCV 2023-Workshop RCV](https://arxiv.org/abs/2308.06468).
+
+## Usage
+
+```python
+from image_gen_aux import TeedPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/teed/20240922043215.png"
+)
+
+teed_preprocessor = TeedPreprocessor.from_pretrained("OzzyGT/teed").to("cuda")
+image = teed_preprocessor(input_image)[0]
+image.save("teed.png")
+```
+
+## Additional resources
+
+* [Project page](https://github.com/xavysp/TEED)
+* [Paper](https://arxiv.org/abs/2308.06468)

image_gen_aux/preprocessors/teed/__init__.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from ...utils import IMAGE_AUX_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["teed_preprocessor"] = [
+        "TeedPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .teed_preprocessor import (
+            TeedPreprocessor,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/preprocessors/teed/teed.py

@@ -0,0 +1,323 @@
+# Original from: https://github.com/xavysp/TEED
+# TEED: is a Tiny but Efficient Edge Detection, it comes from the LDC-B3
+# with a Slightly modification
+# LDC parameters:
+# 155665
+# TED > 58K
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...modeling_utils import ModelMixin
+
+
+"""
+smish_function and Smish script based on:
+Wang, Xueliang, Honge Ren, and Achuan Wang.
+ "Smish: A Novel Activation Function for Deep Learning Methods.
+ " Electronics 11.4 (2022): 540.
+smish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + sigmoid(x)))
+"""
+
+
+@torch.jit.script
+def smish_function(input):
+    """
+    Applies the mish function element-wise:
+    mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(sigmoid(x))))
+    See additional documentation for mish class.
+    """
+    return input * torch.tanh(torch.log(1 + torch.sigmoid(input)))
+
+
+class Smish(nn.Module):
+    """
+    Applies the mish function element-wise:
+    mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
+    Shape:
+        - Input: (N, *) where * means, any number of additional
+          dimensions
+        - Output: (N, *), same shape as the input
+    Examples:
+        >>> m = Mish()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.Mish.html
+    """
+
+    def __init__(self):
+        """
+        Init method.
+        """
+        super().__init__()
+
+    def forward(self, input):
+        """
+        Forward pass of the function.
+        """
+        return smish_function(input)
+
+
+def weight_init(m):
+    if isinstance(m, (nn.Conv2d,)):
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+    # for fusion layer
+    if isinstance(m, (nn.ConvTranspose2d,)):
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+
+class CoFusion(nn.Module):
+    # from LDC
+
+    def __init__(self, in_ch, out_ch):
+        super(CoFusion, self).__init__()
+        self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3, stride=1, padding=1)  # before 64
+        self.conv3 = nn.Conv2d(32, out_ch, kernel_size=3, stride=1, padding=1)  # before 64  instead of 32
+        self.relu = nn.ReLU()
+        self.norm_layer1 = nn.GroupNorm(4, 32)  # before 64
+
+    def forward(self, x):
+        # fusecat = torch.cat(x, dim=1)
+        attn = self.relu(self.norm_layer1(self.conv1(x)))
+        attn = F.softmax(self.conv3(attn), dim=1)
+        return ((x * attn).sum(1)).unsqueeze(1)
+
+
+class CoFusion2(nn.Module):
+    # TEDv14-3
+    def __init__(self, in_ch, out_ch):
+        super(CoFusion2, self).__init__()
+        self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3, stride=1, padding=1)  # before 64
+        self.conv3 = nn.Conv2d(32, out_ch, kernel_size=3, stride=1, padding=1)  # before 64  instead of 32
+        self.smish = Smish()  # nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        attn = self.conv1(self.smish(x))
+        attn = self.conv3(self.smish(attn))  # before , )dim=1)
+
+        return ((x * attn).sum(1)).unsqueeze(1)
+
+
+class DoubleFusion(nn.Module):
+    # TED fusion before the final edge map prediction
+    def __init__(self, in_ch, out_ch):
+        super(DoubleFusion, self).__init__()
+        self.DWconv1 = nn.Conv2d(in_ch, in_ch * 8, kernel_size=3, stride=1, padding=1, groups=in_ch)  # before 64
+        self.PSconv1 = nn.PixelShuffle(1)
+
+        self.DWconv2 = nn.Conv2d(24, 24 * 1, kernel_size=3, stride=1, padding=1, groups=24)  # before 64  instead of 32
+
+        self.AF = Smish()  # XAF() #nn.Tanh()# XAF() #   # Smish()#
+
+    def forward(self, x):
+        attn = self.PSconv1(self.DWconv1(self.AF(x)))  # #TEED best res TEDv14 [8, 32, 352, 352]
+
+        attn2 = self.PSconv1(self.DWconv2(self.AF(attn)))  # #TEED best res TEDv14[8, 3, 352, 352]
+
+        return smish_function(((attn2 + attn).sum(1)).unsqueeze(1))  # TED best res
+
+
+class _DenseLayer(nn.Sequential):
+    def __init__(self, input_features, out_features):
+        super(_DenseLayer, self).__init__()
+
+        (
+            self.add_module(
+                "conv1",
+                nn.Conv2d(
+                    input_features,
+                    out_features,
+                    kernel_size=3,
+                    stride=1,
+                    padding=2,
+                    bias=True,
+                ),
+            ),
+        )
+        (self.add_module("smish1", Smish()),)
+        self.add_module(
+            "conv2",
+            nn.Conv2d(out_features, out_features, kernel_size=3, stride=1, bias=True),
+        )
+
+    def forward(self, x):
+        x1, x2 = x
+
+        new_features = super(_DenseLayer, self).forward(smish_function(x1))  # F.relu()
+
+        return 0.5 * (new_features + x2), x2
+
+
+class _DenseBlock(nn.Sequential):
+    def __init__(self, num_layers, input_features, out_features):
+        super(_DenseBlock, self).__init__()
+        for i in range(num_layers):
+            layer = _DenseLayer(input_features, out_features)
+            self.add_module("denselayer%d" % (i + 1), layer)
+            input_features = out_features
+
+
+class UpConvBlock(nn.Module):
+    def __init__(self, in_features, up_scale):
+        super(UpConvBlock, self).__init__()
+        self.up_factor = 2
+        self.constant_features = 16
+
+        layers = self.make_deconv_layers(in_features, up_scale)
+        assert layers is not None, layers
+        self.features = nn.Sequential(*layers)
+
+    def make_deconv_layers(self, in_features, up_scale):
+        layers = []
+        all_pads = [0, 0, 1, 3, 7]
+        for i in range(up_scale):
+            kernel_size = 2**up_scale
+            pad = all_pads[up_scale]  # kernel_size-1
+            out_features = self.compute_out_features(i, up_scale)
+            layers.append(nn.Conv2d(in_features, out_features, 1))
+            layers.append(Smish())
+            layers.append(nn.ConvTranspose2d(out_features, out_features, kernel_size, stride=2, padding=pad))
+            in_features = out_features
+        return layers
+
+    def compute_out_features(self, idx, up_scale):
+        return 1 if idx == up_scale - 1 else self.constant_features
+
+    def forward(self, x):
+        return self.features(x)
+
+
+class SingleConvBlock(nn.Module):
+    def __init__(self, in_features, out_features, stride, use_ac=False):
+        super(SingleConvBlock, self).__init__()
+        self.use_ac = use_ac
+        self.conv = nn.Conv2d(in_features, out_features, 1, stride=stride, bias=True)
+        if self.use_ac:
+            self.smish = Smish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.use_ac:
+            return self.smish(x)
+        else:
+            return x
+
+
+class DoubleConvBlock(nn.Module):
+    def __init__(self, in_features, mid_features, out_features=None, stride=1, use_act=True):
+        super(DoubleConvBlock, self).__init__()
+
+        self.use_act = use_act
+        if out_features is None:
+            out_features = mid_features
+        self.conv1 = nn.Conv2d(in_features, mid_features, 3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(mid_features, out_features, 3, padding=1)
+        self.smish = Smish()  # nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.smish(x)
+        x = self.conv2(x)
+        if self.use_act:
+            x = self.smish(x)
+        return x
+
+
+class TEED(ModelMixin):
+    """Definition of  Tiny and Efficient Edge Detector
+    model
+    """
+
+    def __init__(self):
+        super(TEED, self).__init__()
+        self.block_1 = DoubleConvBlock(
+            3,
+            16,
+            16,
+            stride=2,
+        )
+        self.block_2 = DoubleConvBlock(16, 32, use_act=False)
+        self.dblock_3 = _DenseBlock(1, 32, 48)  # [32,48,100,100] before (2, 32, 64)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        # skip1 connection, see fig. 2
+        self.side_1 = SingleConvBlock(16, 32, 2)
+
+        # skip2 connection, see fig. 2
+        self.pre_dense_3 = SingleConvBlock(32, 48, 1)  # before (32, 64, 1)
+
+        # USNet
+        self.up_block_1 = UpConvBlock(16, 1)
+        self.up_block_2 = UpConvBlock(32, 1)
+        self.up_block_3 = UpConvBlock(48, 2)  # (32, 64, 1)
+
+        self.block_cat = DoubleFusion(3, 3)  # TEED: DoubleFusion
+
+        self.apply(weight_init)
+
+    def slice(self, tensor, slice_shape):
+        t_shape = tensor.shape
+        img_h, img_w = slice_shape
+        if img_w != t_shape[-1] or img_h != t_shape[2]:
+            new_tensor = F.interpolate(tensor, size=(img_h, img_w), mode="bicubic", align_corners=False)
+
+        else:
+            new_tensor = tensor
+        # tensor[..., :height, :width]
+        return new_tensor
+
+    def resize_input(self, tensor):
+        t_shape = tensor.shape
+        if t_shape[2] % 8 != 0 or t_shape[3] % 8 != 0:
+            img_w = ((t_shape[3] // 8) + 1) * 8
+            img_h = ((t_shape[2] // 8) + 1) * 8
+            new_tensor = F.interpolate(tensor, size=(img_h, img_w), mode="bicubic", align_corners=False)
+        else:
+            new_tensor = tensor
+        return new_tensor
+
+    def crop_bdcn(data1, h, w, crop_h, crop_w):
+        # Based on BDCN Implementation @ https://github.com/pkuCactus/BDCN
+        _, _, h1, w1 = data1.size()
+        assert h <= h1 and w <= w1
+        data = data1[:, :, crop_h : crop_h + h, crop_w : crop_w + w]
+        return data
+
+    def forward(self, x, single_test=False):
+        assert x.ndim == 4, x.shape
+        # supose the image size is 352x352
+
+        # Block 1
+        block_1 = self.block_1(x)  # [8,16,176,176]
+        block_1_side = self.side_1(block_1)  # 16 [8,32,88,88]
+
+        # Block 2
+        block_2 = self.block_2(block_1)  # 32 # [8,32,176,176]
+        block_2_down = self.maxpool(block_2)  # [8,32,88,88]
+        block_2_add = block_2_down + block_1_side  # [8,32,88,88]
+
+        # Block 3
+        block_3_pre_dense = self.pre_dense_3(block_2_down)  # [8,64,88,88] block 3 L connection
+        block_3, _ = self.dblock_3([block_2_add, block_3_pre_dense])  # [8,64,88,88]
+
+        # upsampling blocks
+        out_1 = self.up_block_1(block_1)
+        out_2 = self.up_block_2(block_2)
+        out_3 = self.up_block_3(block_3)
+
+        results = [out_1, out_2, out_3]
+
+        # concatenate multiscale outputs
+        block_cat = torch.cat(results, dim=1)  # Bx6xHxW
+        block_cat = self.block_cat(block_cat)  # Bx1xHxW DoubleFusion
+
+        results.append(block_cat)
+        return results

image_gen_aux/preprocessors/teed/teed_preprocessor.py

@@ -0,0 +1,121 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from safetensors.torch import load_file
+
+from ...image_processor import ImageMixin
+from ...utils import SAFETENSORS_FILE_EXTENSION, get_model_path
+from ..preprocessor import Preprocessor
+from .teed import TEED
+
+
+class TeedPreprocessor(Preprocessor, ImageMixin):
+    """Preprocessor specifically designed for detecting edges in images.
+
+    This class inherits from both `Preprocessor` and `ImageMixin`. Please refer to each
+    one to get more information.
+    """
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_or_path: Union[str, os.PathLike],
+        filename: str = None,
+        subfolder: str = None,
+        weights_only: bool = True,
+    ) -> TEED:
+        model_path = get_model_path(pretrained_model_or_path, filename, subfolder)
+
+        file_extension = os.path.basename(model_path).split(".")[-1]
+        if file_extension == SAFETENSORS_FILE_EXTENSION:
+            state_dict = load_file(model_path, device="cpu")
+        else:
+            state_dict = torch.load(model_path, map_location=torch.device("cpu"), weights_only=weights_only)
+
+        model = TEED()
+        model.load_state_dict(state_dict)
+
+        return cls(model)
+
+    @torch.inference_mode
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]],
+        resolution_scale: float = 1.0,
+        invert: bool = False,
+        safe_steps: int = 2,
+        batch_size: int = 1,
+        return_type: str = "pil",
+    ):
+        """Preprocesses an image and detects the edges using the pre-trained model.
+
+        Args:
+            image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image]]`): Input image as PIL Image,
+                NumPy array, PyTorch tensor format or a list of PIL Images
+            resolution_scale (`float`, optional, defaults to 1.0): Scale factor for image resolution during
+                preprocessing and post-processing. Defaults to 1.0 for no scaling.
+            invert (`bool`, *optional*, defaults to True): Inverts the generated image if True (white or black background).
+            safe_steps (int, optional):
+                Number of safe steps for the TEED model. Defaults to 2.
+            batch_size (`int`, *optional*, defaults to 1): The number of images to process in each batch.
+            return_type (`str`, *optional*, defaults to "pil"): The desired return type, either "pt" for PyTorch tensor, "np" for NumPy array,
+                or "pil" for PIL image.
+
+        Returns:
+            `Union[PIL.Image.Image, np.ndarray, torch.Tensor]`: The generated line art in the
+                specified output format.
+        """
+        if not isinstance(image, torch.Tensor):
+            image = self.convert_image_to_tensor(image, normalize=False)
+
+        image, resolution_scale = self.scale_image(image, resolution_scale)
+
+        processed_images = []
+
+        for i in range(0, len(image), batch_size):
+            batch = image[i : i + batch_size].to(self.model.device)
+
+            edges = self.model(batch)
+            edges = torch.stack([e[0, 0] for e in edges], dim=2)
+            mean_edges = torch.mean(edges, dim=2)
+            edge = torch.sigmoid(mean_edges)
+
+            if safe_steps != 0:
+                edge = self.safe_step(edge, safe_steps)
+
+            if invert:
+                edge = 1 - edge
+
+            processed_images.append(edge.unsqueeze(0).cpu())
+        teed = torch.cat(processed_images, dim=0)
+
+        # add missing channel
+        teed = teed.unsqueeze(1)
+
+        if resolution_scale != 1.0:
+            teed, _ = self.scale_image(teed, 1 / resolution_scale)
+
+        image = self.post_process_image(teed, return_type)
+
+        return image
+
+    def safe_step(self, x, step=2):
+        y = x.float() * float(step + 1)
+        y = y.int().float() / float(step)
+        return y

image_gen_aux/upscalers/README.md

@@ -0,0 +1,58 @@
+# UPSCALERS
+
+## Upscale with model
+
+Class to upscale images with a safetensor checkpoint. We use [spandrel](https://github.com/chaiNNer-org/spandrel) for loading, and you can see the list of supported models [here](https://github.com/chaiNNer-org/spandrel?tab=readme-ov-file#model-architecture-support).
+
+Most of the super-resolution models are provided as `pickle` checkpoints, which are considered unsafe. We promote the use of safetensor checkpoints, and for convenient use, we recommend using the Hugging Face Hub. You can still use a locally downloaded model.
+
+### Space
+
+You can test the current super resolution models you can use with this [Hugging Face Space](https://huggingface.co/spaces/OzzyGT/basic_upscaler).
+
+### How to use
+
+```python
+from image_gen_aux import UpscaleWithModel
+from image_gen_aux.utils import load_image
+
+original = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/simple_upscale/hippowaffle_small.png"
+)
+
+upscaler = UpscaleWithModel.from_pretrained("Kim2091/UltraSharp").to("cuda")
+image = upscaler(original)
+image.save("upscaled.png")
+```
+
+### Tiling
+
+Tiling can be enabled to use less resources.
+
+```python
+image = upscaler(original, tiling=True, tile_width=768, tile_height=768, overlap=8)
+```
+
+### Scale
+
+The scale is automatically obtained from the model but can be overridden with the `scale` argument:
+
+```python
+image = upscaler(original, scale=2)
+```
+
+### List of safetensors checkpoints
+
+This is the current list of safetensor checkpoints you can use from the hub.
+
+|Model|Scale|Repository|Owner|
+|---|---|---|---|
+|UltraSharp|4X|Kim2091/UltraSharp|[Kim2091](https://huggingface.co/Kim2091)|
+|DAT|2X|OzzyGT/DAT_X2|[zhengchen1999](https://github.com/zhengchen1999)|
+|DAT|3X|OzzyGT/DAT_X3|[zhengchen1999](https://github.com/zhengchen1999)|
+|DAT|4X|OzzyGT/DAT_X4|[zhengchen1999](https://github.com/zhengchen1999)|
+|RealPLKSR|4X|Phips/4xNomosWebPhoto_RealPLKSR|[Philip Hofmann](https://huggingface.co/Phips)|
+|DAT-2|4X|Phips/4xRealWebPhoto_v4_dat2|[Philip Hofmann](https://huggingface.co/Phips)|
+|BSRGAN|4X|OzzyGT/4xRemacri|[FoolhardyVEVO](https://openmodeldb.info/users/foolhardy)|
+
+If you own the model and want us to change the repository to your name/organization please open an issue.

image_gen_aux/upscalers/__init__.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from ..utils import IMAGE_AUX_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["upscale_with_model"] = [
+        "UpscaleWithModel",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .upscale_with_model import (
+            UpscaleWithModel,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

image_gen_aux/upscalers/upscale_with_model.py

@@ -0,0 +1,118 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from typing import Union
+
+import numpy as np
+import PIL.Image
+import torch
+from spandrel import ImageModelDescriptor, ModelLoader
+
+from ..image_processor import ImageMixin
+from ..utils import get_model_path, tiled_upscale
+
+
+class UpscaleWithModel(ImageMixin):
+    r"""
+    Upscaler class that uses a pytorch model.
+
+    Args:
+        model ([`ImageModelDescriptor`]):
+            Upscaler model, must be supported by spandrel.
+        scale (`int`, defaults to the scale of the model):
+            The number of times to scale the image, it is recommended to use the model default scale which
+            usually is what the model was trained for.
+    """
+
+    def __init__(self, model: ImageModelDescriptor, scale: int = None):
+        super().__init__()
+        self.model = model
+
+    def to(self, device):
+        self.model.to(device)
+        return self
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_or_path: Union[str, os.PathLike], filename: str = None, subfolder: str = None
+    ) -> ImageModelDescriptor:
+        r"""
+        Instantiate the Upscaler class from pretrained weights.
+
+        Parameters:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                Can be either:
+
+                    - A string, the *repo id* (for example `OzzyGT/UltraSharp`) of a pretrained model
+                      hosted on the Hub, must be saved in safetensors. If there's more than one checkpoint
+                      in the repository and the filename wasn't specified, the first one found will be loaded.
+                    - A path to a *directory* (for example `./upscaler_model/`) containing a pretrained
+                      upscaler checkpoint.
+            filename (`str`, *optional*):
+                The name of the file in the repo.
+            subfolder (`str`, *optional*):
+                An optional value corresponding to a folder inside the model repo.
+        """
+        model_path = get_model_path(pretrained_model_or_path, filename, subfolder)
+        model = ModelLoader().load_from_file(model_path)
+
+        # validate that it's the correct model
+        assert isinstance(model, ImageModelDescriptor)
+
+        return cls(model)
+
+    @torch.inference_mode
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        tiling: bool = False,
+        tile_width: int = 512,
+        tile_height: int = 512,
+        overlap: int = 8,
+        return_type: str = "pil",
+    ) -> Union[torch.Tensor, PIL.Image.Image, np.ndarray]:
+        r"""
+        Upscales the given image, optionally using tiling.
+
+        Args:
+            image (Union[PIL.Image.Image, np.ndarray, torch.Tensor]):
+                The image to be upscaled. Can be a PIL Image, NumPy array, or PyTorch tensor.
+            tiling (bool, optional):
+                Whether to use tiling for upscaling. Default is False.
+            tile_width (int, optional):
+                The width of each tile if tiling is used. Default is 512.
+            tile_height (int, optional):
+                The height of each tile if tiling is used. Default is 512.
+            overlap (int, optional):
+                The overlap between tiles if tiling is used. Default is 8.
+            return_type (str, optional):
+                The type of the returned image. Can be 'pil', 'numpy', or 'tensor'. Default is 'pil'.
+
+        Returns:
+            Union[torch.Tensor, PIL.Image.Image, np.ndarray]:
+                The upscaled image, in the format specified by `return_type`.
+        """
+        if not isinstance(image, torch.Tensor):
+            image = self.convert_image_to_tensor(image)
+
+        image = image.to(self.model.device)
+
+        if tiling:
+            upscaled_tensor = tiled_upscale(image, self.model, self.model.scale, tile_width, tile_height, overlap)
+        else:
+            upscaled_tensor = self.model(image)
+
+        image = self.post_process_image(upscaled_tensor, return_type)[0]
+
+        return image

image_gen_aux/utils/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .constants import SAFETENSORS_FILE_EXTENSION
+from .import_utils import (
+    IMAGE_AUX_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_torch_available,
+    is_transformers_available,
+)
+from .loading_utils import load_image
+from .model_utils import get_model_path
+from .tiling_utils import create_gradient_mask, tiled_upscale

image_gen_aux/utils/constants.py

@@ -0,0 +1 @@
+SAFETENSORS_FILE_EXTENSION = "safetensors"

image_gen_aux/utils/import_utils.py

@@ -0,0 +1,128 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Import utilities: Utilities related to imports and our lazy inits.
+"""
+import importlib.metadata as importlib_metadata
+import importlib.util
+import os
+from itertools import chain
+from types import ModuleType
+from typing import Any
+
+from . import logging
+
+
+logger = logging.get_logger(__name__)
+
+ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
+ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
+
+USE_TF = os.environ.get("USE_TF", "AUTO").upper()
+USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
+IMAGE_AUX_SLOW_IMPORT = os.environ.get("IMAGE_AUX_SLOW_IMPORT", "FALSE").upper()
+IMAGE_AUX_SLOW_IMPORT = IMAGE_AUX_SLOW_IMPORT in ENV_VARS_TRUE_VALUES
+
+
+_torch_version = "N/A"
+if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
+    _torch_available = importlib.util.find_spec("torch") is not None
+    if _torch_available:
+        try:
+            _torch_version = importlib_metadata.version("torch")
+            logger.info(f"PyTorch version {_torch_version} available.")
+        except importlib_metadata.PackageNotFoundError:
+            _torch_available = False
+else:
+    logger.info("Disabling PyTorch because USE_TORCH is set")
+    _torch_available = False
+
+_transformers_available = importlib.util.find_spec("transformers") is not None
+try:
+    _transformers_version = importlib_metadata.version("transformers")
+    logger.debug(f"Successfully imported transformers version {_transformers_version}")
+except importlib_metadata.PackageNotFoundError:
+    _transformers_available = False
+
+
+def is_torch_available():
+    return _torch_available
+
+
+def is_transformers_available():
+    return _transformers_available
+
+
+class OptionalDependencyNotAvailable(BaseException):
+    """An error indicating that an optional dependency of Diffusers was not found in the environment."""
+
+
+class _LazyModule(ModuleType):
+    """
+    Module class that surfaces all objects but only performs associated imports when the objects are requested.
+    """
+
+    # Very heavily inspired by optuna.integration._IntegrationModule
+    # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
+    def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):
+        super().__init__(name)
+        self._modules = set(import_structure.keys())
+        self._class_to_module = {}
+        for key, values in import_structure.items():
+            for value in values:
+                self._class_to_module[value] = key
+        # Needed for autocompletion in an IDE
+        self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values()))
+        self.__file__ = module_file
+        self.__spec__ = module_spec
+        self.__path__ = [os.path.dirname(module_file)]
+        self._objects = {} if extra_objects is None else extra_objects
+        self._name = name
+        self._import_structure = import_structure
+
+    # Needed for autocompletion in an IDE
+    def __dir__(self):
+        result = super().__dir__()
+        # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether
+        # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir.
+        for attr in self.__all__:
+            if attr not in result:
+                result.append(attr)
+        return result
+
+    def __getattr__(self, name: str) -> Any:
+        if name in self._objects:
+            return self._objects[name]
+        if name in self._modules:
+            value = self._get_module(name)
+        elif name in self._class_to_module.keys():
+            module = self._get_module(self._class_to_module[name])
+            value = getattr(module, name)
+        else:
+            raise AttributeError(f"module {self.__name__} has no attribute {name}")
+
+        setattr(self, name, value)
+        return value
+
+    def _get_module(self, module_name: str):
+        try:
+            return importlib.import_module("." + module_name, self.__name__)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
+                f" traceback):\n{e}"
+            ) from e
+
+    def __reduce__(self):
+        return (self.__class__, (self._name, self.__file__, self._import_structure))

image_gen_aux/utils/loading_utils.py

@@ -0,0 +1,63 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Callable, Union
+
+import PIL.Image
+import PIL.ImageOps
+import requests
+
+
+def load_image(
+    image: Union[str, PIL.Image.Image], convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] = None
+) -> PIL.Image.Image:
+    """
+    Loads `image` to a PIL Image.
+
+    Args:
+        image (`str` or `PIL.Image.Image`):
+            The image to convert to the PIL Image format.
+        convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], optional):
+            A conversion method to apply to the image after loading it. When set to `None` the image will be converted
+            "RGB".
+
+    Returns:
+        `PIL.Image.Image`:
+            A PIL Image.
+    """
+    if isinstance(image, str):
+        if image.startswith("http://") or image.startswith("https://"):
+            image = PIL.Image.open(requests.get(image, stream=True).raw)
+        elif os.path.isfile(image):
+            image = PIL.Image.open(image)
+        else:
+            raise ValueError(
+                f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path."
+            )
+    elif isinstance(image, PIL.Image.Image):
+        image = image
+    else:
+        raise ValueError(
+            "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image."
+        )
+
+    image = PIL.ImageOps.exif_transpose(image)
+
+    if convert_method is not None:
+        image = convert_method(image)
+    else:
+        image = image.convert("RGB")
+
+    return image

image_gen_aux/utils/logging.py

@@ -0,0 +1,341 @@
+# coding=utf-8
+# Copyright 2024 Optuna, Hugging Face
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Logging utilities."""
+
+import logging
+import os
+import sys
+import threading
+from logging import (
+    CRITICAL,  # NOQA
+    DEBUG,  # NOQA
+    ERROR,  # NOQA
+    FATAL,  # NOQA
+    INFO,  # NOQA
+    NOTSET,  # NOQA
+    WARN,  # NOQA
+    WARNING,  # NOQA
+)
+from typing import Dict, Optional
+
+from tqdm import auto as tqdm_lib
+
+
+_lock = threading.Lock()
+_default_handler: Optional[logging.Handler] = None
+
+log_levels = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.WARNING
+
+_tqdm_active = True
+
+
+def _get_default_logging_level() -> int:
+    """
+    If DIFFUSERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+    not - fall back to `_default_log_level`
+    """
+    env_level_str = os.getenv("DIFFUSERS_VERBOSITY", None)
+    if env_level_str:
+        if env_level_str in log_levels:
+            return log_levels[env_level_str]
+        else:
+            logging.getLogger().warning(
+                f"Unknown option DIFFUSERS_VERBOSITY={env_level_str}, "
+                f"has to be one of: { ', '.join(log_levels.keys()) }"
+            )
+    return _default_log_level
+
+
+def _get_library_name() -> str:
+    return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+    return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if _default_handler:
+            # This library has already configured the library root logger.
+            return
+        _default_handler = logging.StreamHandler()  # Set sys.stderr as stream.
+
+        if sys.stderr:  # only if sys.stderr exists, e.g. when not using pythonw in windows
+            _default_handler.flush = sys.stderr.flush
+
+        # Apply our default configuration to the library root logger.
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.addHandler(_default_handler)
+        library_root_logger.setLevel(_get_default_logging_level())
+        library_root_logger.propagate = False
+
+
+def _reset_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if not _default_handler:
+            return
+
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.removeHandler(_default_handler)
+        library_root_logger.setLevel(logging.NOTSET)
+        _default_handler = None
+
+
+def get_log_levels_dict() -> Dict[str, int]:
+    return log_levels
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Return a logger with the specified name.
+
+    This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
+    """
+
+    if name is None:
+        name = _get_library_name()
+
+    _configure_library_root_logger()
+    return logging.getLogger(name)
+
+
+def get_verbosity() -> int:
+    """
+    Return the current level for the 🤗 Diffusers' root logger as an `int`.
+
+    Returns:
+        `int`:
+            Logging level integers which can be one of:
+
+            - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+            - `40`: `diffusers.logging.ERROR`
+            - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+            - `20`: `diffusers.logging.INFO`
+            - `10`: `diffusers.logging.DEBUG`
+
+    """
+
+    _configure_library_root_logger()
+    return _get_library_root_logger().getEffectiveLevel()
+
+
+def set_verbosity(verbosity: int) -> None:
+    """
+    Set the verbosity level for the 🤗 Diffusers' root logger.
+
+    Args:
+        verbosity (`int`):
+            Logging level which can be one of:
+
+            - `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+            - `diffusers.logging.ERROR`
+            - `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+            - `diffusers.logging.INFO`
+            - `diffusers.logging.DEBUG`
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().setLevel(verbosity)
+
+
+def set_verbosity_info() -> None:
+    """Set the verbosity to the `INFO` level."""
+    return set_verbosity(INFO)
+
+
+def set_verbosity_warning() -> None:
+    """Set the verbosity to the `WARNING` level."""
+    return set_verbosity(WARNING)
+
+
+def set_verbosity_debug() -> None:
+    """Set the verbosity to the `DEBUG` level."""
+    return set_verbosity(DEBUG)
+
+
+def set_verbosity_error() -> None:
+    """Set the verbosity to the `ERROR` level."""
+    return set_verbosity(ERROR)
+
+
+def disable_default_handler() -> None:
+    """Disable the default handler of the 🤗 Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().removeHandler(_default_handler)
+
+
+def enable_default_handler() -> None:
+    """Enable the default handler of the 🤗 Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().addHandler(_default_handler)
+
+
+def add_handler(handler: logging.Handler) -> None:
+    """adds a handler to the HuggingFace Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None
+    _get_library_root_logger().addHandler(handler)
+
+
+def remove_handler(handler: logging.Handler) -> None:
+    """removes given handler from the HuggingFace Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None and handler in _get_library_root_logger().handlers
+    _get_library_root_logger().removeHandler(handler)
+
+
+def disable_propagation() -> None:
+    """
+    Disable propagation of the library log outputs. Note that log propagation is disabled by default.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = False
+
+
+def enable_propagation() -> None:
+    """
+    Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent
+    double logging if the root logger has been configured.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = True
+
+
+def enable_explicit_format() -> None:
+    """
+    Enable explicit formatting for every 🤗 Diffusers' logger. The explicit formatter is as follows:
+    ```
+    [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
+    ```
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s")
+        handler.setFormatter(formatter)
+
+
+def reset_format() -> None:
+    """
+    Resets the formatting for 🤗 Diffusers' loggers.
+
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        handler.setFormatter(None)
+
+
+def warning_advice(self, *args, **kwargs) -> None:
+    """
+    This method is identical to `logger.warning()`, but if env var DIFFUSERS_NO_ADVISORY_WARNINGS=1 is set, this
+    warning will not be printed
+    """
+    no_advisory_warnings = os.getenv("DIFFUSERS_NO_ADVISORY_WARNINGS", False)
+    if no_advisory_warnings:
+        return
+    self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_advice = warning_advice
+
+
+class EmptyTqdm:
+    """Dummy tqdm which doesn't do anything."""
+
+    def __init__(self, *args, **kwargs):  # pylint: disable=unused-argument
+        self._iterator = args[0] if args else None
+
+    def __iter__(self):
+        return iter(self._iterator)
+
+    def __getattr__(self, _):
+        """Return empty function."""
+
+        def empty_fn(*args, **kwargs):  # pylint: disable=unused-argument
+            return
+
+        return empty_fn
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        return
+
+
+class _tqdm_cls:
+    def __call__(self, *args, **kwargs):
+        if _tqdm_active:
+            return tqdm_lib.tqdm(*args, **kwargs)
+        else:
+            return EmptyTqdm(*args, **kwargs)
+
+    def set_lock(self, *args, **kwargs):
+        self._lock = None
+        if _tqdm_active:
+            return tqdm_lib.tqdm.set_lock(*args, **kwargs)
+
+    def get_lock(self):
+        if _tqdm_active:
+            return tqdm_lib.tqdm.get_lock()
+
+
+tqdm = _tqdm_cls()
+
+
+def is_progress_bar_enabled() -> bool:
+    """Return a boolean indicating whether tqdm progress bars are enabled."""
+    global _tqdm_active
+    return bool(_tqdm_active)
+
+
+def enable_progress_bar() -> None:
+    """Enable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = True
+
+
+def disable_progress_bar() -> None:
+    """Disable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = False

image_gen_aux/utils/model_utils.py

@@ -0,0 +1,37 @@
+import os
+
+from huggingface_hub import hf_hub_download, model_info
+
+
+def get_model_path(pretrained_model_or_path, filename=None, subfolder=None):
+    """
+    Retrieves the path to the model file.
+
+    If `pretrained_model_or_path` is a file, it returns the path directly.
+    Otherwise, it attempts to find a `.safetensors` file associated with the given model path.
+    If no `.safetensors` file is found, it raises a `FileNotFoundError`.
+
+    Parameters:
+    - pretrained_model_or_path (str): Path to the pretrained model or directory containing the model.
+    - filename (str, optional): Specific filename to load. If not provided, the function will search for a `.safetensors` file.
+    - subfolder (str, optional): Subfolder within the model directory to look for the file.
+
+    Returns:
+    - str: Path to the model file.
+
+    Raises:
+    - FileNotFoundError: If no `.safetensors` file is found when `filename` is not provided.
+    """
+    if os.path.isfile(pretrained_model_or_path):
+        return pretrained_model_or_path
+
+    if filename is None:
+        # If the filename is not passed, we only try to load a safetensor
+        info = model_info(pretrained_model_or_path)
+        filename = next(
+            (sibling.rfilename for sibling in info.siblings if sibling.rfilename.endswith(".safetensors")), None
+        )
+        if filename is None:
+            raise FileNotFoundError("No safetensors checkpoint found.")
+
+    return hf_hub_download(pretrained_model_or_path, filename, subfolder=subfolder)

image_gen_aux/utils/tiling_utils.py

@@ -0,0 +1,86 @@
+from typing import Callable, Tuple
+
+import torch
+
+
+def create_gradient_mask(shape: Tuple, feather: int, device="cpu") -> torch.Tensor:
+    """
+    Create a gradient mask for smooth blending of tiles.
+
+    Args:
+        shape (tuple): Shape of the mask (batch, channels, height, width)
+        feather (int): Width of the feathered edge
+
+    Returns:
+        torch.Tensor: Gradient mask
+    """
+    mask = torch.ones(shape).to(device)
+    _, _, h, w = shape
+    for feather_step in range(feather):
+        factor = (feather_step + 1) / feather
+        mask[:, :, feather_step, :] *= factor
+        mask[:, :, h - 1 - feather_step, :] *= factor
+        mask[:, :, :, feather_step] *= factor
+        mask[:, :, :, w - 1 - feather_step] *= factor
+    return mask
+
+
+def tiled_upscale(
+    samples: torch.Tensor,
+    function: Callable,
+    scale: int,
+    tile_width: int = 512,
+    tile_height: int = 512,
+    overlap: int = 8,
+) -> torch.Tensor:
+    """
+    Apply a scaling function to image samples in a tiled manner.
+
+    Args:
+        samples (torch.Tensor): Input tensor of shape (batch_size, channels, height, width)
+        function (Callable): The scaling function to apply to each tile
+        scale (int): Factor by which to upscale the image
+        tile_width (int): Width of each tile
+        tile_height (int): Height of each tile
+        overlap (int): Overlap between tiles
+
+    Returns:
+        torch.Tensor: Upscaled and processed output tensor
+    """
+    _batch, _channels, height, width = samples.shape
+    out_height, out_width = round(height * scale), round(width * scale)
+    output_device = samples.device
+
+    # Initialize output tensors
+    output = torch.empty((1, 3, out_height, out_width), device=output_device)
+    out = torch.zeros((1, 3, out_height, out_width), device=output_device)
+    out_div = torch.zeros_like(output)
+
+    # Process the image in tiles
+    for y in range(0, height, tile_height - overlap):
+        for x in range(0, width, tile_width - overlap):
+            # Ensure we don't go out of bounds
+            x_end = min(x + tile_width, width)
+            y_end = min(y + tile_height, height)
+            x = max(0, x_end - tile_width)
+            y = max(0, y_end - tile_height)
+
+            # Extract and process the tile
+            tile = samples[:, :, y:y_end, x:x_end]
+            processed_tile = function(tile).to(output_device)
+
+            # Calculate the position in the output tensor
+            out_y, out_x = round(y * scale), round(x * scale)
+            out_h, out_w = processed_tile.shape[2:]
+
+            # Create a feathered mask for smooth blending
+            mask = create_gradient_mask(processed_tile.shape, overlap * scale, device=output_device)
+
+            # Add the processed tile to the output
+            out[:, :, out_y : out_y + out_h, out_x : out_x + out_w] += processed_tile * mask
+            out_div[:, :, out_y : out_y + out_h, out_x : out_x + out_w] += mask
+
+    # Normalize the output
+    output = out / out_div
+
+    return output

requirements.txt

@@ -1,5 +1,4 @@
 einops
-diffusers==0.31.0
 transformers
 accelerate
 gradio
@@ -8,4 +7,5 @@ spaces
 pillow
 peft
 openai
-torch
+torch
+git+https://github.com/huggingface/diffusers