monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #173 from akx/path-types

Browse Source 
Fix up some path-related types
This commit is contained in:
Filip Strand 2025-04-25 14:29:35 +02:00 committed by GitHub
commit 39f89026f0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 50 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/mflux/config/runtime_config.py
View File

@ -1,4 +1,5 @@
import logging
from pathlib import Path
import mlx.core as mx
import numpy as np
@ -48,7 +49,7 @@ class RuntimeConfig:
return self.model_config.num_train_steps
@property
def image_path(self) -> str:
def image_path(self) -> Path | None:
return self.config.image_path
@property
@ -56,15 +57,15 @@ class RuntimeConfig:
return self.config.image_strength
@property
def depth_image_path(self) -> str | None:
def depth_image_path(self) -> Path | None:
return self.config.depth_image_path
@property
def redux_image_paths(self) -> str | None:
def redux_image_paths(self) -> list[Path] | None:
return self.config.redux_image_paths
@property
def masked_image_path(self) -> str | None:
def masked_image_path(self) -> Path | None:
return self.config.masked_image_path
@property

4
src/mflux/dreambooth/state/zip_util.py
View File

@ -6,7 +6,9 @@ from zipfile import ZipFile
class ZipUtil:
@staticmethod
def unzip(zip_path: str | Path, filename: str, loader: callable):
def unzip(zip_path: str | Path | None, filename: str, loader: callable):
if not zip_path: # Would be nicer to do this in typing, but that's more effort on the callers' side
raise ValueError("zip_path cannot be None")
zip_path = Path(zip_path)
if not zip_path.exists():
raise FileNotFoundError(f"ZIP file not found at: {zip_path}")

11
src/mflux/flux_tools/depth/depth_util.py
View File

@ -1,5 +1,6 @@
import logging
import os
from pathlib import Path
import mlx.core as mx
import PIL.Image
@ -19,8 +20,8 @@ class DepthUtil:
vae: VAE,
depth_pro: DepthPro,
config: RuntimeConfig,
image_path: str | None = None,
depth_image_path: str | None = None,
image_path: str | Path | None = None,
depth_image_path: str | Path | None = None,
) -> (mx.array, PIL.Image.Image):
# 1. Create the depth map or use existing one
depth_image_path, depth_image = DepthUtil.get_or_create_depth_map(
@ -44,9 +45,9 @@ class DepthUtil:
@staticmethod
def get_or_create_depth_map(
depth_pro: DepthPro,
image_path: str | None = None,
depth_map_path: str | None = None,
) -> tuple[str, PIL.Image.Image | None]:
image_path: str | Path | None = None,
depth_map_path: str | Path | None = None,
) -> tuple[str | Path, PIL.Image.Image | None]:
# 1. If a depth map path is provided, use it directly
if depth_map_path:
if not os.path.exists(depth_map_path):

6
src/mflux/flux_tools/fill/mask_util.py
View File

@ -1,3 +1,5 @@
from pathlib import Path
import mlx.core as mx
from mflux.config.runtime_config import RuntimeConfig
@ -12,8 +14,8 @@ class MaskUtil:
vae: VAE,
config: RuntimeConfig,
latents: mx.array,
img_path: str,
mask_path: str | None,
img_path: str | Path,
mask_path: str | Path | None,
) -> mx.array:
if not img_path or not mask_path:
# Return empty latents if no image or mask is provided

4
src/mflux/flux_tools/redux/flux_redux.py
View File

@ -1,3 +1,5 @@
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
@ -158,7 +160,7 @@ class Flux1Redux(nn.Module):
clip_tokenizer: TokenizerCLIP,
t5_text_encoder: T5Encoder,
clip_text_encoder: CLIPEncoder,
image_paths: list[str],
image_paths: list[str] | list[Path],
image_encoder: SiglipVisionTransformer,
image_embedder: ReduxEncoder,
) -> (mx.array, mx.array):

6
src/mflux/flux_tools/redux/redux_util.py
View File

@ -1,3 +1,5 @@
from pathlib import Path
import mlx.core as mx
from mflux import ImageUtil
@ -8,7 +10,7 @@ from mflux.models.siglip_vision_transformer.siglip_vision_transformer import Sig
class ReduxUtil:
@staticmethod
def embed_images(
image_paths: list[str],
image_paths: list[str] | list[Path],
image_encoder: SiglipVisionTransformer,
image_embedder: ReduxEncoder,
) -> list[mx.array]: # fmt:off
@ -24,7 +26,7 @@ class ReduxUtil:
@staticmethod
def _embed_single_image(
image_path: str,
image_path: str | Path,
image_encoder: SiglipVisionTransformer,
image_embedder: ReduxEncoder,
) -> mx.array: # fmt:off

10
src/mflux/latent_creator/latent_creator.py
View File

@ -1,3 +1,5 @@
from pathlib import Path
import mlx.core as mx
from mflux.models.vae.vae import VAE
@ -11,7 +13,7 @@ class Img2Img:
vae: VAE,
sigmas: mx.array,
init_time_step: int,
image_path: int,
image_path: str | Path | None,
):
self.vae = vae
self.sigmas = sigmas
@ -39,9 +41,7 @@ class LatentCreator:
img2img: Img2Img,
) -> mx.array:
# 0. Determine type of image generation
is_text2img = img2img.image_path is None
if is_text2img:
if img2img.image_path is None:
# 1. Create the pure noise
return LatentCreator.create(
seed=seed,
@ -71,7 +71,7 @@ class LatentCreator:
)
@staticmethod
def encode_image(vae: VAE, image_path: str, height: int, width: int):
def encode_image(vae: VAE, image_path: str | Path, height: int, width: int):
scaled_user_image = ImageUtil.scale_to_dimensions(
image=ImageUtil.load_image(image_path).convert("RGB"),
target_width=width,

16
src/mflux/post_processing/generated_image.py
View File

@ -1,6 +1,6 @@
import importlib
import pathlib
import typing as t
from pathlib import Path
import mlx.core as mx
import PIL.Image
@ -23,13 +23,13 @@ class GeneratedImage:
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
controlnet_image_path: str | pathlib.Path | None = None,
controlnet_image_path: str | Path | None = None,
controlnet_strength: float | None = None,
image_path: str | pathlib.Path | None = None,
image_path: str | Path | None = None,
image_strength: float | None = None,
masked_image_path: str | pathlib.Path | None = None,
depth_image_path: str | pathlib.Path | None = None,
redux_image_paths: list[str] | list[pathlib.Path] | None = None,
masked_image_path: str | Path | None = None,
depth_image_path: str | Path | None = None,
redux_image_paths: list[str] | list[Path] | None = None,
):
self.image = image
self.model_config = model_config
@ -78,7 +78,7 @@ class GeneratedImage:
def save(
self,
path: t.Union[str, pathlib.Path],
path: t.Union[str, Path],
export_json_metadata: bool = False,
overwrite: bool = False,
) -> None:
@ -129,7 +129,7 @@ class GeneratedImage:
@staticmethod
def _get_version_from_toml() -> str | None:
# Search for pyproject.toml by traversing up from the current working directory
current_dir = pathlib.Path(__file__).resolve().parent
current_dir = Path(__file__).resolve().parent
for parent in current_dir.parents:
pyproject_path = parent / "pyproject.toml"
if pyproject_path.exists():

20
src/mflux/post_processing/image_util.py
View File

@ -1,7 +1,7 @@
import json
import logging
import pathlib
import typing as t
from pathlib import Path
import mlx.core as mx
import numpy as np
@ -27,12 +27,12 @@ class ImageUtil:
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
controlnet_image_path: str | None = None,
image_path: str | None = None,
redux_image_paths: list[str] | None = None,
controlnet_image_path: str | Path | None = None,
image_path: str | Path | None = None,
redux_image_paths: list[str] | list[Path] | None = None,
image_strength: float | None = None,
masked_image_path: str | None = None,
depth_image_path: str | None = None,
masked_image_path: str | Path | None = None,
depth_image_path: str | Path | None = None,
) -> GeneratedImage:
normalized = ImageUtil._denormalize(decoded_latents)
normalized_numpy = ImageUtil._to_numpy(normalized)
@ -113,7 +113,7 @@ class ImageUtil:
return array
@staticmethod
def load_image(path: str | pathlib.Path) -> PIL.Image.Image:
def load_image(path: str | Path) -> PIL.Image.Image:
return PIL.Image.open(path)
@staticmethod
@ -204,12 +204,12 @@ class ImageUtil:
@staticmethod
def save_image(
image: PIL.Image.Image,
path: t.Union[str, pathlib.Path],
path: t.Union[str, Path],
metadata: dict | None = None,
export_json_metadata: bool = False,
overwrite: bool = False,
) -> None:
file_path = pathlib.Path(path)
file_path = Path(path)
file_path.parent.mkdir(parents=True, exist_ok=True)
file_name = file_path.stem
file_extension = file_path.suffix
@ -240,7 +240,7 @@ class ImageUtil:
log.error(f"Error saving image: {e}")
@staticmethod
def _embed_metadata(metadata: dict, path: str) -> None:
def _embed_metadata(metadata: dict, path: str | Path) -> None:
try:
# Convert metadata dictionary to a string
metadata_str = json.dumps(metadata)

4
tools/create_outpaint_image_canvas_and_mask.py
View File

@ -1,5 +1,5 @@
import pathlib
import sys
from pathlib import Path
from mflux.post_processing.image_util import ImageUtil
from mflux.ui.box_values import AbsoluteBoxValues, BoxValues
@ -8,7 +8,7 @@ from mflux.ui.cli.parsers import CommandLineParser
def main():
parser = CommandLineParser(description="Create expanded canvas and mask for outpainting")
parser.add_argument("image_path", type=pathlib.Path, help="Path to the input image file")
parser.add_argument("image_path", type=Path, help="Path to the input image file")
parser.add_image_outpaint_arguments(required=True)
args = parser.parse_args()