Rolled back to decoupled zimage and hunyuan 3d, removed redundant dependencies

inner_layer/models/zimage/zimage.py

@@ -1,153 +1,72 @@
 import io
 import base64
-import torch
-from diffusers import ZImagePipeline, FlowMatchEulerDiscreteScheduler
-from sdnq import SDNQConfig # import sdnq to register it into diffusers and transformers
-from sdnq.common import use_torch_compile as triton_is_available
-from sdnq.loader import apply_sdnq_options_to_model
-from pathlib import Path
 import uuid
-import trimesh
-import numpy as np
+from pathlib import Path
 from PIL import Image
+import numpy as np
+import trimesh
+from gradio_client import Client
 
-
+# --- Ensure output directories ---
 def ensure_dir(path: str):
     Path(path).mkdir(parents=True, exist_ok=True)
 
-# Global variable to store the loaded model
-_zimage_pipe = None
-
-
-def get_zimage_pipeline():
-    """
-    Lazily load and return the global Z-Image-Turbo pipeline.
-    This prevents reloading the model every time the API is called.
-    """
-    global _zimage_pipe
-    
-    if _zimage_pipe is None:
-        device = get_best_device()
-        print(f"[ZImage] Using device: {device}")
-        print(f"[ZImage] Loading Z-Image-Turbo on {device}...")
-        
-        # Use bfloat16 for better quality
-        dtype = torch.bfloat16 if device in ["mps", "cuda"] else torch.float32
-
-        _zimage_pipe = ZImagePipeline.from_pretrained(
-            "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-        )
-        
-        _zimage_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
-            _zimage_pipe.scheduler.config,
-            use_beta_sigmas=True,
-        )
-        
-        _zimage_pipe.to(device)
-        
-        # if triton_is_available and (torch.cuda.is_available() or torch.xpu.is_available()):
-        #     _zimage_pipe.transformer = apply_sdnq_options_to_model(_zimage_pipe.transformer, use_quantized_matmul=True)
-        #     _zimage_pipe.text_encoder = apply_sdnq_options_to_model(_zimage_pipe.text_encoder, use_quantized_matmul=True)
-        #     _zimage_pipe.transformer = torch.compile(_zimage_pipe.transformer) # optional for faster speeds
-
-        _zimage_pipe.enable_model_cpu_offload()
-
-        
-        _zimage_pipe.enable_attention_slicing()
-
-        # if hasattr(_zimage_pipe, "enable_vae_slicing"):
-        #     _zimage_pipe.enable_vae_slicing()
-        #     print("[ZImage] VAE slicing enabled")
-
-        # if hasattr(getattr(_zimage_pipe, "vae", None), "enable_tiling"):
-        #     _zimage_pipe.vae.enable_tiling()
-        #     print("[ZImage] VAE tiling enabled")
-        # print("[ZImage] Model loaded successfully.")
-    return _zimage_pipe
-
-import torch
-
-def get_best_device():
-    if torch.cuda.is_available():
-        return "cuda"
-    elif torch.version.hip is not None:  # ROCm (AMD)
-        return "hip"
-    elif torch.backends.mps.is_available():
-        return "mps"
-    elif torch.backends.mps.is_built():  # fallback in case of MPS build but not available
-        return "mps"
-    elif torch.backends.opencl.is_available():  # not always present, but check
-        return "opencl"
-    elif torch.has_mps:  # just in case for Apple
-        return "mps"
-    else:
-        return "cpu"
-
-
+# --- Convert PIL image to GLB ---
 def image_to_glb(image: Image.Image, thickness: float = 0.001) -> bytes:
     w, h = image.size
     aspect = w / h
 
-    # Thin box so viewers don’t cull it
-    mesh = trimesh.creation.box(
-        extents=(aspect, 1.0, thickness)
-    )
+    mesh = trimesh.creation.box(extents=(aspect, 1.0, thickness))
 
-    # Simple UV mapping
     uv = np.zeros((len(mesh.vertices), 2))
     uv[:, 0] = (mesh.vertices[:, 0] / aspect + 0.5)
     uv[:, 1] = (mesh.vertices[:, 1] + 0.5)
 
-    mesh.visual = trimesh.visual.TextureVisuals(
-        uv=uv,
-        image=image
-    )
-
+    mesh.visual = trimesh.visual.TextureVisuals(uv=uv, image=image)
     return mesh.export(file_type="glb")
 
 
+# --- Global Gradio client ---
+_client = None
+def get_gradio_client():
+    global _client
+    if _client is None:
+        _client = Client("mrfakename/Z-Image-Turbo")
+    return _client
+
+
+# --- Main image generation function ---
 def generate_image_base64(
     prompt: str,
     height: int = 768,
     width: int = 768,
-    steps: int = 50,
+    steps: int = 8,
     seed: int = 5,
     convert_to_glb: bool = False,
+    randomize_seed: bool = True,
 ) -> str:
     """
-    Generate an image from text prompt using Z-Image-Turbo.
+    Generate an image from text prompt using Gradio Z-Image-Turbo.
     Returns:
       - base64 PNG if convert_to_glb=False
       - base64 GLB if convert_to_glb=True
     """
-    pipe = get_zimage_pipeline()
-
-    if seed is None:
-        seed = torch.randint(0, 2**32, (1,)).item()
-    print(f"Generating with seed {seed}...")
-
-    device = get_best_device()
-
-    if device == "cuda":
-        generator = torch.Generator("cuda").manual_seed(seed)
-    elif device == "mps":
-        generator = torch.Generator("mps").manual_seed(seed)
-    else:
-        generator = torch.Generator().manual_seed(seed)
+    client = get_gradio_client()
+
+    # Predict using the Gradio API
+    result, used_seed = client.predict(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        seed=seed,
+        randomize_seed=randomize_seed,
+        api_name="/generate_image"
+    )
 
-    with torch.inference_mode():
-        result = pipe(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            guidance_scale=0.0,
-            generator=generator,
-        )
+   
 
-    image = result.images[0]
+    image = Image.open(result)
 
     ensure_dir("output/images")
     img_name = f"zimage_{uuid.uuid4().hex}.png"
@@ -155,19 +74,17 @@ def generate_image_base64(
     image.save(img_path)
     print(f"[ZImage] Saved image to {img_path}")
 
-    # --- If GLB is NOT requested, return image ---
+    # --- Return PNG base64 ---
     if not convert_to_glb:
         
         return img_path
 
-    # --- Convert to GLB only if requested ---
+    # --- Convert to GLB ---
     glb_bytes = image_to_glb(image)
-
     glb_name = f"zimage_{uuid.uuid4().hex}.glb"
     glb_path = Path("output/images") / glb_name
     with open(glb_path, "wb") as f:
         f.write(glb_bytes)
+    print(f"[ZImage] Saved GLB to {glb_path}")
 
-    print(f"[ZImage] Saved glb to {glb_path}")
-
-    return base64.b64encode(glb_bytes).decode("utf-8")
+    return base64.b64encode(glb_bytes).decode("utf-8")

requirements.txt

@@ -1,139 +1,102 @@
-accelerate==1.12.0
 aiohappyeyeballs==2.6.1
-aiohttp==3.13.2
+aiohttp==3.13.3
 aiosignal==1.4.0
 annotated-doc==0.0.4
 annotated-types==0.7.0
-anyio==4.12.0
-async-timeout==5.0.1
+anyio==4.12.1
 attrs==25.4.0
-Authlib==1.6.6
-backports.tarfile==1.2.0
+Authlib==1.6.7
 beartype==0.22.9
-cachetools==6.2.3
-certifi==2025.11.12
+cachetools==7.0.0
+certifi==2026.1.4
 cffi==2.0.0
 charset-normalizer==3.4.4
 click==8.3.1
 cloudpickle==3.1.2
-coloredlogs==15.0.1
-cryptography==46.0.3
-cyclopts==4.3.0
-diffusers @ git+https://github.com/huggingface/diffusers@17c0e79dbdf53fb6705e9c09cc1a854b84c39249
+croniter==6.0.0
+cryptography==46.0.4
+cyclopts==4.5.1
 diskcache==5.6.3
 dnspython==2.8.0
 docstring_parser==0.17.0
-docutils==0.22.3
+docutils==0.22.4
 email-validator==2.3.0
 exceptiongroup==1.3.1
-fakeredis==2.32.1
-fastapi==0.124.4
-fastapi-mcp==0.4.0
-fastmcp==2.14.0
-filelock==3.20.0
-flatbuffers==25.9.23
+fakeredis==2.33.0
+fastapi==0.128.4
+fastmcp==2.14.5
+filelock==3.20.3
 frozenlist==1.8.0
-fsspec==2025.12.0
+fsspec==2026.2.0
+gradio_client==2.0.3
 h11==0.16.0
 hf-xet==1.2.0
 httpcore==1.0.9
 httpx==0.28.1
 httpx-sse==0.4.3
-huggingface-hub==0.36.0
-humanfriendly==10.0
+huggingface_hub==1.4.1
 idna==3.11
 ImageIO==2.37.2
-importlib_metadata==8.7.0
+importlib_metadata==8.7.1
 jaraco.classes==3.4.0
-jaraco.context==6.0.1
-jaraco.functools==4.3.0
-Jinja2==3.1.6
-jsonschema==4.25.1
+jaraco.context==6.1.0
+jaraco.functools==4.4.0
+jsonref==1.1.0
+jsonschema==4.26.0
 jsonschema-path==0.3.4
 jsonschema-specifications==2025.9.1
 keyring==25.7.0
-lazy_loader==0.4
-llvmlite==0.46.0
 lupa==2.6
 markdown-it-py==4.0.0
-MarkupSafe==3.0.3
-mcp==1.24.0
+mcp==1.26.0
 mdurl==0.1.2
 more-itertools==10.8.0
-mpmath==1.3.0
-multidict==6.7.0
-networkx==3.4.2
-numba==0.63.1
-numpy==2.2.6
-onnxruntime==1.23.2
+multidict==6.7.1
+numpy==2.4.2
 openapi-pydantic==0.5.1
-opencv-python-headless==4.12.0.88
 opentelemetry-api==1.39.1
-opentelemetry-exporter-prometheus==0.60b1
-opentelemetry-instrumentation==0.60b1
-opentelemetry-sdk==1.39.1
-opentelemetry-semantic-conventions==0.60b1
-packaging==25.0
+packaging==26.0
 pathable==0.4.4
 pathvalidate==3.3.1
-pillow==12.0.0
+pillow==12.1.0
 platformdirs==4.5.1
-pooch==1.8.2
-prometheus_client==0.23.1
+prometheus_client==0.24.1
 propcache==0.4.1
-protobuf==6.33.2
-psutil==7.1.3
 py-key-value-aio==0.3.0
 py-key-value-shared==0.3.0
-pybind11==3.0.1
-pycparser==2.23
+pycparser==3.0
 pydantic==2.12.5
 pydantic-settings==2.12.0
 pydantic_core==2.41.5
-pydocket==0.15.5
+pydocket==0.17.5
 Pygments==2.19.2
-PyJWT==2.10.1
-PyMatting==1.1.14
+PyJWT==2.11.0
 pyperclip==1.11.0
+python-dateutil==2.9.0.post0
 python-dotenv==1.2.1
 python-json-logger==4.0.0
-python-multipart==0.0.20
+python-multipart==0.0.22
+pytz==2025.2
 PyYAML==6.0.3
 redis==7.1.0
 referencing==0.36.2
-regex==2025.11.3
-rembg==2.0.69
 requests==2.32.5
-rich==14.2.0
+rich==14.3.2
 rich-rst==1.3.2
 rpds-py==0.30.0
-safetensors==0.7.0
-scikit-image==0.25.2
-scipy==1.15.3
-sdnq==0.1.2
 shellingham==1.5.4
+six==1.17.0
 sortedcontainers==2.4.0
-sse-starlette==3.0.3
-starlette==0.50.0
-sympy==1.14.0
-tifffile==2025.5.10
-tokenizers==0.22.1
-tomli==2.3.0
-torch==2.9.1
-torchvision==0.24.1
-tqdm==4.67.1
-transformers==4.57.3
-typer==0.20.0
-typer-slim==0.20.0
+sse-starlette==3.2.0
+starlette==0.52.1
+tqdm==4.67.3
+trimesh==4.11.1
+typer==0.21.1
+typer-slim==0.21.1
 typing-inspection==0.4.2
 typing_extensions==4.15.0
-urllib3==2.6.2
-uvicorn==0.38.0
-websockets==15.0.1
-wrapt==1.17.3
+urllib3==2.6.3
+uvicorn==0.40.0
+websockets==16.0
 yarl==1.22.0
 zipp==3.23.0
-
-
-trimesh
-gradio_client