Fix type for Img2Img.image_path

This commit is contained in:
Aarni Koskela 2025-04-25 14:01:13 +03:00
parent ca2f96d75e
commit 79ca146653

View File

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