Added support for finetuning OmniGen2.

extensions_built_in/diffusion_models/omnigen2/src/utils/img_util.py

@@ -0,0 +1,31 @@
+from typing import List
+
+from PIL import Image
+
+import torch
+from torchvision.transforms.functional import to_pil_image
+
+def resize_image(image, max_pixels, img_scale_num):
+    width, height = image.size
+    cur_pixels = height * width
+    ratio = (max_pixels / cur_pixels) ** 0.5
+    ratio = min(ratio, 1.0) # do not upscale input image
+
+    new_height, new_width = int(height * ratio) // img_scale_num * img_scale_num, int(width * ratio) // img_scale_num * img_scale_num
+
+    image = image.resize((new_width, new_height), resample=Image.BICUBIC)
+    return image
+
+def create_collage(images: List[torch.Tensor]) -> Image.Image:
+    """Create a horizontal collage from a list of images."""
+    max_height = max(img.shape[-2] for img in images)
+    total_width = sum(img.shape[-1] for img in images)
+    canvas = torch.zeros((3, max_height, total_width), device=images[0].device)
+    
+    current_x = 0
+    for img in images:
+        h, w = img.shape[-2:]
+        canvas[:, :h, current_x:current_x+w] = img * 0.5 + 0.5
+        current_x += w
+    
+    return to_pil_image(canvas)

extensions_built_in/diffusion_models/omnigen2/src/utils/import_utils.py

@@ -0,0 +1,46 @@
+# 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.util
+import sys
+
+# The package importlib_metadata is in a different place, depending on the python version.
+if sys.version_info < (3, 8):
+    import importlib_metadata
+else:
+    import importlib.metadata as importlib_metadata
+
+def _is_package_available(pkg_name: str):
+    pkg_exists = importlib.util.find_spec(pkg_name) is not None
+    pkg_version = "N/A"
+
+    if pkg_exists:
+        try:
+            pkg_version = importlib_metadata.version(pkg_name)
+        except (ImportError, importlib_metadata.PackageNotFoundError):
+            pkg_exists = False
+
+    return pkg_exists, pkg_version
+
+_triton_available, _triton_version = _is_package_available("triton")
+_flash_attn_available, _flash_attn_version = _is_package_available("flash_attn")
+
+def is_triton_available():
+    return _triton_available
+
+def is_flash_attn_available():
+    return _flash_attn_available