nexaai 1.0.16rc8__cp310-cp310-macosx_13_0_x86_64.whl → 1.0.16rc10__cp310-cp310-macosx_13_0_x86_64.whl

nexaai/_version.py

@@ -1,4 +1,4 @@
 # This file is generated by CMake from _version.py.in
 # Do not modify this file manually - it will be overwritten
 
-__version__ = "1.0.16-rc8" 
+__version__ = "1.0.16-rc10" 

nexaai/mlx_backend/ml.py

@@ -1,9 +1,6 @@
 # This file defines the python interface that c-lib expects from a python backend
 
 from __future__ import annotations
-from typing import Optional
-from pathlib import Path
-from dataclasses import dataclass
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -104,12 +101,9 @@ class ModelConfig:
     n_threads_batch: int = 0  # number of threads to use for batch processing
     n_batch: int = 0  # logical maximum batch size that can be submitted to llama_decode
     n_ubatch: int = 0  # physical maximum batch size
-    # max number of sequences (i.e. distinct states for recurrent models)
-    n_seq_max: int = 0
-    # path to chat template file, optional
-    chat_template_path: Optional[Path] = None
-    # content of chat template file, optional
-    chat_template_content: Optional[str] = None
+    n_seq_max: int = 0  # max number of sequences (i.e. distinct states for recurrent models)
+    chat_template_path: Optional[Path] = None  # path to chat template file, optional
+    chat_template_content: Optional[str] = None  # content of chat template file, optional
 
 
 @dataclass
@@ -124,8 +118,7 @@ class SamplerConfig:
     frequency_penalty: float = 0.0
     seed: int = -1  # –1 for random
     grammar_path: Optional[Path] = None
-    # Optional grammar string (BNF-like format)
-    grammar_string: Optional[str] = None
+    grammar_string: Optional[str] = None  # Optional grammar string (BNF-like format)
 
 
 @dataclass
@@ -135,10 +128,8 @@ class GenerationConfig:
     stop: Sequence[str] = field(default_factory=tuple)
     n_past: int = 0
     sampler_config: Optional[SamplerConfig] = None
-    # Array of image paths for VLM (None if none)
-    image_paths: Optional[Sequence[Path]] = None
-    # Array of audio paths for VLM (None if none)
-    audio_paths: Optional[Sequence[Path]] = None
+    image_paths: Optional[Sequence[Path]] = None  # Array of image paths for VLM (None if none)
+    audio_paths: Optional[Sequence[Path]] = None  # Array of audio paths for VLM (None if none)
 
 
 @dataclass
@@ -179,32 +170,6 @@ class RerankConfig:
     normalize_method: str = "softmax"  # "softmax" | "min-max" | "none"
 
 
-# image-gen
-
-
-@dataclass
-class ImageGenTxt2ImgInput:
-    """Input structure for text-to-image generation."""
-    prompt: str
-    config: ImageGenerationConfig
-    output_path: Optional[Path] = None
-
-
-@dataclass
-class ImageGenImg2ImgInput:
-    """Input structure for image-to-image generation."""
-    init_image_path: Path
-    prompt: str
-    config: ImageGenerationConfig
-    output_path: Optional[Path] = None
-
-
-@dataclass
-class ImageGenOutput:
-    """Output structure for image generation."""
-    output_image_path: Path
-
-
 @dataclass
 class ImageSamplerConfig:
     """Configuration for image sampling."""
@@ -215,27 +180,17 @@ class ImageSamplerConfig:
     seed: int = -1  # –1 for random
 
 
-@dataclass
-class ImageGenCreateInput:
-    """Configuration for image generation."""
-    model_name: str
-    model_path: Path
-    config: ModelConfig
-    scheduler_config_path: Path
-    plugin_id: str
-    device_id: Optional[str] = None
-
-
 @dataclass
 class ImageGenerationConfig:
     """Configuration for image generation."""
-    prompts: List[str]
-    sampler_config: ImageSamplerConfig
-    scheduler_config: SchedulerConfig
-    strength: float
-    negative_prompts: Optional[List[str]] = None
+    prompts: str | List[str]
+    negative_prompts: str | List[str] | None = None
     height: int = 512
     width: int = 512
+    sampler_config: Optional[ImageSamplerConfig] = None
+    lora_id: int = -1  # –1 for none
+    init_image: Optional[Image] = None
+    strength: float = 1.0
 
 
 @dataclass
@@ -306,7 +261,7 @@ class TTSResult:
 class BoundingBox:
     """Generic bounding box structure."""
     x: float  # X coordinate (normalized or pixel, depends on model)
-    y: float  # Y coordinate (normalized or pixel, depends on model)
+    y: float  # Y coordinate (normalized or pixel, depends on model) 
     width: float  # Width
     height: float  # Height
 
@@ -320,8 +275,7 @@ class CVResult:
     confidence: float = 0.0  # Confidence score [0.0-1.0]
     bbox: Optional[BoundingBox] = None  # Bounding box (example: YOLO)
     text: Optional[str] = None  # Text result (example: OCR)
-    # Feature embedding (example: CLIP embedding)
-    embedding: Optional[List[float]] = None
+    embedding: Optional[List[float]] = None  # Feature embedding (example: CLIP embedding)
     embedding_dim: int = 0  # Embedding dimension
 
 

nexaai/mlx_backend/sd/modeling/model_io.py

@@ -1,7 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import json
-import os
 from typing import Optional
 
 import mlx.core as mx
@@ -177,37 +176,19 @@ def _load_safetensor_weights(mapper, model, weight_file, float16: bool = False):
 
 
 def _check_key(key: str, part: str):
-    # Check if it's a local path
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # For local paths, we'll use a default model structure
-        return
     if key not in _MODELS:
         raise ValueError(
             f"[{part}] '{key}' model not found, choose one of {{{','.join(_MODELS.keys())}}}"
         )
 
-def _get_model_path(key: str, file_path: str):
-    """Get the full path for a model file, supporting both local and HuggingFace paths"""
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path
-        return os.path.join(key, file_path)
-    else:
-        # HuggingFace path
-        return hf_hub_download(key, file_path)
-
 
 def load_unet(key: str = _DEFAULT_MODEL, float16: bool = False):
     """Load the stable diffusion UNet from Hugging Face Hub."""
     _check_key(key, "load_unet")
 
-    # Get the config path
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        unet_config = "unet/config.json"
-    else:
-        unet_config = _MODELS[key]["unet_config"]
-    
-    with open(_get_model_path(key, unet_config)) as f:
+    # Download the config and create the model
+    unet_config = _MODELS[key]["unet_config"]
+    with open(hf_hub_download(key, unet_config)) as f:
         config = json.load(f)
 
     n_blocks = len(config["block_out_channels"])
@@ -238,13 +219,8 @@ def load_unet(key: str = _DEFAULT_MODEL, float16: bool = False):
     )
 
     # Download the weights and map them into the model
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        unet_weights = "unet/diffusion_pytorch_model.safetensors"
-    else:
-        unet_weights = _MODELS[key]["unet"]
-    
-    weight_file = _get_model_path(key, unet_weights)
+    unet_weights = _MODELS[key]["unet"]
+    weight_file = hf_hub_download(key, unet_weights)
     _load_safetensor_weights(map_unet_weights, model, weight_file, float16)
 
     return model
@@ -262,13 +238,8 @@ def load_text_encoder(
     config_key = config_key or (model_key + "_config")
 
     # Download the config and create the model
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        text_encoder_config = f"{model_key}/config.json"
-    else:
-        text_encoder_config = _MODELS[key][config_key]
-    
-    with open(_get_model_path(key, text_encoder_config)) as f:
+    text_encoder_config = _MODELS[key][config_key]
+    with open(hf_hub_download(key, text_encoder_config)) as f:
         config = json.load(f)
 
     with_projection = "WithProjection" in config["architectures"][0]
@@ -286,13 +257,8 @@ def load_text_encoder(
     )
 
     # Download the weights and map them into the model
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        text_encoder_weights = f"{model_key}/model.safetensors"
-    else:
-        text_encoder_weights = _MODELS[key][model_key]
-    
-    weight_file = _get_model_path(key, text_encoder_weights)
+    text_encoder_weights = _MODELS[key][model_key]
+    weight_file = hf_hub_download(key, text_encoder_weights)
     _load_safetensor_weights(map_clip_text_encoder_weights, model, weight_file, float16)
 
     return model
@@ -303,13 +269,8 @@ def load_autoencoder(key: str = _DEFAULT_MODEL, float16: bool = False):
     _check_key(key, "load_autoencoder")
 
     # Download the config and create the model
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        vae_config = "vae/config.json"
-    else:
-        vae_config = _MODELS[key]["vae_config"]
-    
-    with open(_get_model_path(key, vae_config)) as f:
+    vae_config = _MODELS[key]["vae_config"]
+    with open(hf_hub_download(key, vae_config)) as f:
         config = json.load(f)
 
     model = Autoencoder(
@@ -326,13 +287,8 @@ def load_autoencoder(key: str = _DEFAULT_MODEL, float16: bool = False):
     )
 
     # Download the weights and map them into the model
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        vae_weights = "vae/diffusion_pytorch_model.safetensors"
-    else:
-        vae_weights = _MODELS[key]["vae"]
-    
-    weight_file = _get_model_path(key, vae_weights)
+    vae_weights = _MODELS[key]["vae"]
+    weight_file = hf_hub_download(key, vae_weights)
     _load_safetensor_weights(map_vae_weights, model, weight_file, float16)
 
     return model
@@ -342,13 +298,8 @@ def load_diffusion_config(key: str = _DEFAULT_MODEL):
     """Load the stable diffusion config from Hugging Face Hub."""
     _check_key(key, "load_diffusion_config")
 
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        diffusion_config = "scheduler/scheduler_config.json"
-    else:
-        diffusion_config = _MODELS[key]["diffusion_config"]
-    
-    with open(_get_model_path(key, diffusion_config)) as f:
+    diffusion_config = _MODELS[key]["diffusion_config"]
+    with open(hf_hub_download(key, diffusion_config)) as f:
         config = json.load(f)
 
     return DiffusionConfig(
@@ -366,17 +317,11 @@ def load_tokenizer(
 ):
     _check_key(key, "load_tokenizer")
 
-    if os.path.exists(key) or '/' in key or '\\' in key:
-        # Local path - use SDXL Turbo structure
-        vocab_file = _get_model_path(key, f"tokenizer/{vocab_key.split('_')[1]}.json")
-        merges_file = _get_model_path(key, f"tokenizer/{merges_key.split('_')[1]}.txt")
-    else:
-        vocab_file = _get_model_path(key, _MODELS[key][vocab_key])
-        merges_file = _get_model_path(key, _MODELS[key][merges_key])
-    
+    vocab_file = hf_hub_download(key, _MODELS[key][vocab_key])
     with open(vocab_file, encoding="utf-8") as f:
         vocab = json.load(f)
 
+    merges_file = hf_hub_download(key, _MODELS[key][merges_key])
     with open(merges_file, encoding="utf-8") as f:
         bpe_merges = f.read().strip().split("\n")[1 : 49152 - 256 - 2 + 1]
     bpe_merges = [tuple(m.split()) for m in bpe_merges]

{nexaai-1.0.16rc8.dist-info → nexaai-1.0.16rc10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nexaai
-Version: 1.0.16rc8
+Version: 1.0.16rc10
 Summary: Python bindings for NexaSDK C-lib backend
 Author-email: "Nexa AI, Inc." <dev@nexa.ai>
 Project-URL: Homepage, https://github.com/NexaAI/nexasdk-bridge

{nexaai-1.0.16rc8.dist-info → nexaai-1.0.16rc10.dist-info}/RECORD

@@ -1,6 +1,6 @@
 nexaai/__init__.py,sha256=jXdC4vv6DBK1fVewYTYSUhOOYfvf_Mk81UIeMGGIKUg,2029
-nexaai/_stub.cpython-310-darwin.so,sha256=zvVd8ZOu9tG4yYAHPt-3aGryeGsJ7-5c1KTZdh2U99k,49832
-nexaai/_version.py,sha256=HYA8rSkVin3FWg68b5gJ7nJTXAuo0Wk_FZ6Xr8Fv-hQ,143
+nexaai/_stub.cpython-310-darwin.so,sha256=bjv9gcR4eVwpi9riqlrWzsFcawoQgULf9ZcwFX0em3g,49832
+nexaai/_version.py,sha256=rM7mysA3I6mJ-jXt3DwjydINMplIqs78L_Oid1wMpI8,144
 nexaai/asr.py,sha256=NljMXDErwPNMOPaRkJZMEDka9Nk8xyur7L8i924TStY,2054
 nexaai/base.py,sha256=N8PRgDFA-XPku2vWnQIofQ7ipz3pPlO6f8YZGnuhquE,982
 nexaai/common.py,sha256=yBnIbqYaQYnfrl7IczOBh6MDibYZVxwaRJEglYcKgGs,3422
@@ -18,7 +18,7 @@ nexaai/asr_impl/pybind_asr_impl.py,sha256=pE9Hb_hMi5yAc4MF83bLVOb8zDtreCkB3_u7XE
 nexaai/binds/__init__.py,sha256=eYuay_8DDXeOUWz2_R9HFSabohxs6hvZn391t2L0Po0,104
 nexaai/binds/common_bind.cpython-310-darwin.so,sha256=KX_nfaQQPzxGv0GS0efcO-ByWai7y7RskMZvSAAaJWI,233960
 nexaai/binds/embedder_bind.cpython-310-darwin.so,sha256=b2NoXFAJvPLi_P1X7lXLKmAUU0v2HJI3Zwa10gfqHdw,202032
-nexaai/binds/libnexa_bridge.dylib,sha256=yaKtbQlBTvXqO14Ge5Pz3XFJpAu-wbOh5Fjjs9410xY,250376
+nexaai/binds/libnexa_bridge.dylib,sha256=haJ62A2Xnv3ZQ76GIT-HTqXh9maiMpNTSaFCCTNIJko,250376
 nexaai/binds/llm_bind.cpython-310-darwin.so,sha256=aKcT2kW1PL1xPFX7vsT6Gs79ZydcVfg8bKtnEthRpI4,183008
 nexaai/binds/vlm_bind.cpython-310-darwin.so,sha256=LGd-tykePnQFfGca25HnPIBfXsfrMzbwyx6d5Ld3xps,183000
 nexaai/binds/nexa_llama_cpp/libggml-base.dylib,sha256=GyOkHOM-5uHp7NUZ4Sr9BWak6BYpcc9aqI9A-zPnQp4,629528
@@ -40,7 +40,7 @@ nexaai/image_gen_impl/pybind_image_gen_impl.py,sha256=ms34VYoD5AxZFG6cIG0QAJDjCt
 nexaai/llm_impl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nexaai/llm_impl/mlx_llm_impl.py,sha256=4v7jUFzHfE7zw2uViekGQDaTROz8A6oaW31Z3iVe6tg,11157
 nexaai/llm_impl/pybind_llm_impl.py,sha256=aooqkcXZWhCo07wbSafGgBrA3WnijtnUADShjjgFsBQ,8051
-nexaai/mlx_backend/ml.py,sha256=DKXVOAfh8cg7KTKljh7jpcPwfQFNigc6uv_ZXF6lse8,23977
+nexaai/mlx_backend/ml.py,sha256=LafDM_TeXmuQkld2tdQxUBGgooT0JPMXngLam2TADqU,23179
 nexaai/mlx_backend/profiling.py,sha256=Dc-mybFwBdCIKFWL7CbSHjkOJGAoYHG7r_e_XPhzwBU,9361
 nexaai/mlx_backend/asr/__init__.py,sha256=fuT_9_xpYJ28m4yjly5L2jChUrzlSQz-b_S7nujxkSM,451
 nexaai/mlx_backend/asr/interface.py,sha256=pE5expr8sP7O9bEgWaUaJ4ITX0MsCxFovG9iVWqVDVU,4246
@@ -57,18 +57,6 @@ nexaai/mlx_backend/embedding/interface.py,sha256=M7AGiq_UVLNIi2Ie6H08ySnMxIjIhUl
 nexaai/mlx_backend/embedding/main.py,sha256=xKRebBcooKuf8DzWKwCicftes3MAcYAd1QvcT9_AAPQ,6003
 nexaai/mlx_backend/embedding/modeling/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nexaai/mlx_backend/embedding/modeling/nexa_jina_v2.py,sha256=F9Z_9r-Dh0wNThiMp5W5hqE2dt5bf4ps5_c6h4BuWGw,15218
-nexaai/mlx_backend/image_gen/__init__.py,sha256=8eFAF00-yWdEs0LJiszjHEsgPDAF8tSruBYvujCNgE0,42
-nexaai/mlx_backend/image_gen/generate_sd.py,sha256=8DuPsJlxVf1LlFARVThiW807G-0cWPReYONj-x8qKRo,8616
-nexaai/mlx_backend/image_gen/interface.py,sha256=7ElIiLm5gZXAfKPs497wLkn8qqLFLL5T-oV-hXVmEdw,3285
-nexaai/mlx_backend/image_gen/main.py,sha256=jKkGDfqmGIsQwhHJaL_j_CTtr04xplD6dnnGe3AcilU,8123
-nexaai/mlx_backend/image_gen/stable_diffusion/__init__.py,sha256=wriLb0wA5vCBlCoQMtfKrVVWMJw8fhXCCk6R9_Nrb9c,9524
-nexaai/mlx_backend/image_gen/stable_diffusion/clip.py,sha256=feHQXi1NiGa01AMo7nK8M-sgBoZBHI95xAGMfAv64kE,3733
-nexaai/mlx_backend/image_gen/stable_diffusion/config.py,sha256=lq2sWXevpnCk8KXOywEVOzgFT6WBNnG4xr7NLWgidII,1773
-nexaai/mlx_backend/image_gen/stable_diffusion/model_io.py,sha256=ei-pdrqFxjeHLL6Ov5NZaHT5lceqZbHDyNFfUFUh_04,13755
-nexaai/mlx_backend/image_gen/stable_diffusion/sampler.py,sha256=nJxiSmF79E2aDGzLqygOat4zpgaTzfR7Kp1PD0bmhWA,3397
-nexaai/mlx_backend/image_gen/stable_diffusion/tokenizer.py,sha256=WMcmyQmb5MmhCp-NNU72Bqjq9vQvDzgR8KAfGKA2Qso,2991
-nexaai/mlx_backend/image_gen/stable_diffusion/unet.py,sha256=rYj1-baAugL7CNyYYvLch891WUgpunhAJCVBNKxfZSA,14762
-nexaai/mlx_backend/image_gen/stable_diffusion/vae.py,sha256=sQ1943x5SKyTmeL2xe3t84FL9DGO72Eab2m88r-BSe8,8044
 nexaai/mlx_backend/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nexaai/mlx_backend/llm/generate.py,sha256=Phes0tzxbbEWA2hDylQvD0LjorMaPwvcfZq9RKCAOt0,4399
 nexaai/mlx_backend/llm/interface.py,sha256=SZFkuAUi2vxj_dSqj8RXf9vPTGMtpks_pZxxrF7iIe8,29330
@@ -236,7 +224,7 @@ nexaai/mlx_backend/sd/main.py,sha256=cHlZhG8KHDFXbYRH-diKA7B1Qacq1euuGw0pKeXJGAI
 nexaai/mlx_backend/sd/modeling/__init__.py,sha256=wriLb0wA5vCBlCoQMtfKrVVWMJw8fhXCCk6R9_Nrb9c,9524
 nexaai/mlx_backend/sd/modeling/clip.py,sha256=feHQXi1NiGa01AMo7nK8M-sgBoZBHI95xAGMfAv64kE,3733
 nexaai/mlx_backend/sd/modeling/config.py,sha256=lq2sWXevpnCk8KXOywEVOzgFT6WBNnG4xr7NLWgidII,1773
-nexaai/mlx_backend/sd/modeling/model_io.py,sha256=xUjF5XNUfRLHQz-LtGT_D3XGQ1MI7ZQWknmdUDMpi_s,13732
+nexaai/mlx_backend/sd/modeling/model_io.py,sha256=0jpMHJdMoTk0Jo4Uz1ZGiXekAWzPuTTD5g7ms4m4idY,11668
 nexaai/mlx_backend/sd/modeling/sampler.py,sha256=nJxiSmF79E2aDGzLqygOat4zpgaTzfR7Kp1PD0bmhWA,3397
 nexaai/mlx_backend/sd/modeling/tokenizer.py,sha256=WMcmyQmb5MmhCp-NNU72Bqjq9vQvDzgR8KAfGKA2Qso,2991
 nexaai/mlx_backend/sd/modeling/unet.py,sha256=rYj1-baAugL7CNyYYvLch891WUgpunhAJCVBNKxfZSA,14762
@@ -387,7 +375,7 @@ nexaai/utils/quantization_utils.py,sha256=4gvp6UQfSO9G1FYBwnFtQspTzH9sDbi1PBXw2t
 nexaai/vlm_impl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nexaai/vlm_impl/mlx_vlm_impl.py,sha256=pLtWm_ckz8a0U-AtAOMVseFDO4OVPvHyYO2KlfBaGYk,10833
 nexaai/vlm_impl/pybind_vlm_impl.py,sha256=FAbhpRJzHgI78r0mUvKybO97R1szvNhH0aTn_I52oT4,8597
-nexaai-1.0.16rc8.dist-info/METADATA,sha256=mMeCi8dxuYWGVRLrai250qYihewe89ROiq72He4RPSo,1201
-nexaai-1.0.16rc8.dist-info/WHEEL,sha256=0KYp5feZ1CMUhsfFXKpSQTbSmQbXy4mv6yPPVBXg2EM,110
-nexaai-1.0.16rc8.dist-info/top_level.txt,sha256=LRE2YERlrZk2vfuygnSzsEeqSknnZbz3Z1MHyNmBU4w,7
-nexaai-1.0.16rc8.dist-info/RECORD,,
+nexaai-1.0.16rc10.dist-info/METADATA,sha256=_p5OXiFfM3Nb6BN_grqE9--aJ4A9bz_B84UPqWYxp9I,1202
+nexaai-1.0.16rc10.dist-info/WHEEL,sha256=0KYp5feZ1CMUhsfFXKpSQTbSmQbXy4mv6yPPVBXg2EM,110
+nexaai-1.0.16rc10.dist-info/top_level.txt,sha256=LRE2YERlrZk2vfuygnSzsEeqSknnZbz3Z1MHyNmBU4w,7
+nexaai-1.0.16rc10.dist-info/RECORD,,

nexaai/mlx_backend/image_gen/__init__.py

@@ -1 +0,0 @@
-# Image generation module for MLX backend

nexaai/mlx_backend/image_gen/generate_sd.py

@@ -1,244 +0,0 @@
-from __future__ import annotations
-
-from typing import (
-    List,
-    Optional,
-)
-
-import mlx.core as mx
-import numpy as np
-from PIL import Image as PILImage
-import mlx.nn as nn
-import os
-
-from .stable_diffusion import StableDiffusion, StableDiffusionXL
-
-
-class Image:
-    def __init__(self, data: List[float], width: int, height: int, channels: int) -> None:
-        """Initialize an image with pixel data"""
-        self.data = data
-        self.width = width
-        self.height = height
-        self.channels = channels
-
-    @classmethod
-    def from_numpy(cls, array: np.ndarray) -> 'Image':
-        """Create Image from numpy array (H, W, C)"""
-        height, width, channels = array.shape
-        data = array.flatten().tolist()
-        return cls(data, width, height, channels)
-
-    @classmethod
-    def from_pil(cls, pil_image: PILImage.Image) -> 'Image':
-        """Create Image from PIL Image"""
-        array = np.array(pil_image).astype(np.float32) / 255.0
-        return cls.from_numpy(array)
-
-    def to_numpy(self) -> np.ndarray:
-        """Convert to numpy array (H, W, C)"""
-        return np.array(self.data).reshape(self.height, self.width, self.channels)
-
-    def to_pil(self) -> PILImage.Image:
-        """Convert to PIL Image"""
-        array = (self.to_numpy() * 255).astype(np.uint8)
-        return PILImage.fromarray(array)
-
-
-class ImageSamplerConfig:
-    def __init__(
-        self,
-        method: str = "ddim",
-        steps: int = 4,  # SDXL Turbo typically uses fewer steps
-        guidance_scale: float = 0.0,  # SDXL Turbo works well with no guidance
-        eta: float = 0.0,
-        seed: int = -1,
-    ) -> None:
-        """Initialize sampler configuration optimized for SDXL Turbo"""
-        self.method = method
-        self.steps = steps
-        self.guidance_scale = guidance_scale
-        self.eta = eta
-        self.seed = seed
-
-
-class ImageGenerationConfig:
-    def __init__(
-        self,
-        prompts: str | List[str],
-        negative_prompts: str | List[str] | None = None,
-        height: int = 512,
-        width: int = 512,
-        sampler_config: Optional[ImageSamplerConfig] = None,
-        lora_id: int = -1,  # Not used but kept for compatibility
-        init_image: Optional[Image] = None,
-        strength: float = 1.0,
-        n_images: int = 1,
-        n_rows: int = 1,
-        decoding_batch_size: int = 1,
-    ) -> None:
-        """Initialize image generation configuration"""
-        self.prompts = prompts
-        self.negative_prompts = negative_prompts or ""
-        self.height = height
-        self.width = width
-        self.sampler_config = sampler_config or ImageSamplerConfig()
-        self.lora_id = lora_id
-        self.init_image = init_image
-        self.strength = strength
-        self.n_images = n_images
-        self.n_rows = n_rows
-        self.decoding_batch_size = decoding_batch_size
-
-
-class ImageGen:
-    def __init__(
-        self,
-        model_path: str,
-        scheduler_config_path: Optional[str] = None,
-        device: Optional[str] = None,
-        float16: bool = True,
-        quantize: bool = False,
-    ) -> None:
-        """Initialize the image generation model for SDXL Turbo"""
-        self.model_path = model_path
-        self.scheduler_config_path = scheduler_config_path
-        self.float16 = float16
-        self.quantize = quantize
-        self.model = None
-
-    @staticmethod
-    def load_model(model_path: str, float16: bool = True, quantize: bool = False) -> StableDiffusion:
-        """Load a model from the given path - following txt2img.py pattern"""
-
-        # Check if it's a local path or HuggingFace repo
-        # If it contains path separators or exists as a file/directory, treat as local
-        is_local_path = (
-            '/' in model_path or '\\' in model_path or os.path.exists(model_path))
-
-        if is_local_path:
-            # For local paths, determine model type from the path or model files
-            if "xl" in model_path.lower() or "turbo" in model_path.lower():
-                model = StableDiffusionXL(model_path, float16=float16)
-            else:
-                model = StableDiffusion(model_path, float16=float16)
-        else:
-            # For HuggingFace repo names, use the original logic
-            if "xl" in model_path.lower() or "turbo" in model_path.lower():
-                model = StableDiffusionXL(model_path, float16=float16)
-            else:
-                model = StableDiffusion(model_path, float16=float16)
-
-        # Apply quantization if requested - same as txt2img.py
-        if quantize:
-            if "xl" in model_path.lower() or "turbo" in model_path.lower():
-                nn.quantize(
-                    model.text_encoder_1, class_predicate=lambda _, m: isinstance(
-                        m, nn.Linear)
-                )
-                nn.quantize(
-                    model.text_encoder_2, class_predicate=lambda _, m: isinstance(
-                        m, nn.Linear)
-                )
-            else:
-                nn.quantize(
-                    model.text_encoder, class_predicate=lambda _, m: isinstance(
-                        m, nn.Linear)
-                )
-            nn.quantize(model.unet, group_size=32, bits=8)
-        return model
-
-    def txt2img(self, prompt: str, config: ImageGenerationConfig, clear_cache: bool = True) -> Image:
-        """Generate an image from a text prompt - following txt2img.py pattern"""
-        if not self.model:
-            self.model = self.load_model(self.model_path)
-            if not self.model:
-                raise RuntimeError("Model not loaded")
-
-        sampler_config = config.sampler_config
-
-        negative_prompt = ""
-        if config.negative_prompts:
-            negative_prompt = config.negative_prompts if isinstance(
-                config.negative_prompts, str) else config.negative_prompts[0]
-
-        # Generate latents - following txt2img.py approach
-        latents_generator = self.model.generate_latents(
-            prompt,
-            n_images=1,
-            num_steps=sampler_config.steps,
-            cfg_weight=sampler_config.guidance_scale,
-            negative_text=negative_prompt,
-            seed=sampler_config.seed if sampler_config.seed >= 0 else None
-        )
-
-        # Get final latents - following txt2img.py pattern
-        final_latents = None
-        for latents in latents_generator:
-            final_latents = latents
-            mx.eval(final_latents)
-
-        if final_latents is None:
-            raise RuntimeError("No latents generated")
-
-        # Decode to image - following txt2img.py pattern
-        decoded_image = self.model.decode(final_latents)
-        mx.eval(decoded_image)
-
-        # Convert to numpy array
-        image_array = np.array(decoded_image.squeeze(0))
-
-        if clear_cache:
-            mx.clear_cache()
-
-        return Image.from_numpy(image_array)
-
-    def img2img(self, init_image: Image, prompt: str, config: ImageGenerationConfig, clear_cache: bool = True) -> Image:
-        """Generate an image from an initial image and a text prompt using SDXL Turbo"""
-        if not self.model:
-            self.model = self.load_model(self.model_path)
-            if not self.model:
-                raise RuntimeError("Model not loaded")
-
-        sampler_config = config.sampler_config
-
-        negative_prompt = ""
-        if config.negative_prompts:
-            negative_prompt = config.negative_prompts if isinstance(
-                config.negative_prompts, str) else config.negative_prompts[0]
-
-        img_tensor = _prepare_image_for_sd(
-            init_image, config.width, config.height)
-
-        # Generate latents from image
-        latents_generator = self.model.generate_latents_from_image(
-            img_tensor,
-            prompt,
-            n_images=1,
-            strength=config.strength,
-            num_steps=sampler_config.steps,
-            cfg_weight=sampler_config.guidance_scale,
-            negative_text=negative_prompt,
-            seed=sampler_config.seed if sampler_config.seed >= 0 else None
-        )
-
-        # Get final latents
-        final_latents = None
-        for latents in latents_generator:
-            final_latents = latents
-            mx.eval(final_latents)
-
-        if final_latents is None:
-            raise RuntimeError("No latents generated")
-
-        # Decode to image
-        decoded_image = self.model.decode(final_latents)
-        mx.eval(decoded_image)
-
-        # Convert to numpy array
-        image_array = np.array(decoded_image.squeeze(0))
-
-        if clear_cache:
-            mx.clear_cache()
-
-        return Image.from_numpy(image_array)