app_quant_latent1.py

import torch
import spaces
import gradio as gr
import sys
import platform
import diffusers
import transformers
import psutil
import os
import time

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers import ZImagePipeline, AutoModel
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
latent_history = []

# ============================================================
# LOGGING BUFFER
# ============================================================
LOGS = ""
def log(msg):
    global LOGS
    print(msg)
    LOGS += msg + "\n"
    return msg


# ============================================================
# SYSTEM METRICS — LIVE GPU + CPU MONITORING
# ============================================================
def log_system_stats(tag=""):
    try:
        log(f"\n===== 🔥 SYSTEM STATS {tag} =====")

        # ============= GPU STATS =============
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated(0) / 1e9
            reserved = torch.cuda.memory_reserved(0) / 1e9
            total = torch.cuda.get_device_properties(0).total_memory / 1e9
            free = total - allocated

            log(f"💠 GPU Total     : {total:.2f} GB")
            log(f"💠 GPU Allocated : {allocated:.2f} GB")
            log(f"💠 GPU Reserved  : {reserved:.2f} GB")
            log(f"💠 GPU Free      : {free:.2f} GB")

        # ============= CPU STATS ============
        cpu = psutil.cpu_percent()
        ram_used = psutil.virtual_memory().used / 1e9
        ram_total = psutil.virtual_memory().total / 1e9

        log(f"🧠 CPU Usage     : {cpu}%")
        log(f"🧠 RAM Used      : {ram_used:.2f} GB / {ram_total:.2f} GB")

    except Exception as e:
        log(f"⚠️ Failed to log system stats: {e}")


# ============================================================
# ENVIRONMENT INFO
# ============================================================
log("===================================================")
log("🔍 Z-IMAGE-TURBO DEBUGGING + LIVE METRIC LOGGER")
log("===================================================\n")

log(f"📌 PYTHON VERSION       : {sys.version.replace(chr(10),' ')}")
log(f"📌 PLATFORM             : {platform.platform()}")
log(f"📌 TORCH VERSION        : {torch.__version__}")
log(f"📌 TRANSFORMERS VERSION : {transformers.__version__}")
log(f"📌 DIFFUSERS VERSION    : {diffusers.__version__}")
log(f"📌 CUDA AVAILABLE       : {torch.cuda.is_available()}")

log_system_stats("AT STARTUP")

if not torch.cuda.is_available():
    raise RuntimeError("❌ CUDA Required")

device = "cuda"
gpu_id = 0

# ============================================================
# MODEL SETTINGS
# ============================================================
model_cache = "./weights/"
model_id = "Tongyi-MAI/Z-Image-Turbo"
torch_dtype = torch.bfloat16
USE_CPU_OFFLOAD = False

log("\n===================================================")
log("🧠 MODEL CONFIGURATION")
log("===================================================")
log(f"Model ID              : {model_id}")
log(f"Model Cache Directory : {model_cache}")
log(f"torch_dtype           : {torch_dtype}")
log(f"USE_CPU_OFFLOAD       : {USE_CPU_OFFLOAD}")

log_system_stats("BEFORE TRANSFORMER LOAD")


# ============================================================
# FUNCTION TO CONVERT LATENTS TO IMAGE
# ============================================================
def latent_to_image(latent):
    try:
        img_tensor = pipe.vae.decode(latent)
        img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
        pil_img = T.ToPILImage()(img_tensor[0])
        return pil_img
    except Exception as e:
        log(f"⚠️ Failed to decode latent: {e}")
        return None



# ============================================================
# SAFE TRANSFORMER INSPECTION
# ============================================================
def inspect_transformer(model, name):
    log(f"\n🔍 Inspecting {name}")
    try:
        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
        blocks = None

        for attr in candidates:
            if hasattr(model, attr):
                blocks = getattr(model, attr)
                break

        if blocks is None:
            log(f"⚠️ No block structure found in {name}")
            return

        if hasattr(blocks, "__len__"):
            log(f"Total Blocks = {len(blocks)}")
        else:
            log("⚠️ Blocks exist but are not iterable")

        for i in range(min(10, len(blocks) if hasattr(blocks, "__len__") else 0)):
            log(f"Block {i} = {blocks[i].__class__.__name__}")

    except Exception as e:
        log(f"⚠️ Transformer inspect error: {e}")


# ============================================================
# LOAD TRANSFORMER — WITH LIVE STATS
# ============================================================
log("\n===================================================")
log("🔧 LOADING TRANSFORMER BLOCK")
log("===================================================")

log("📌 Logging memory before load:")
log_system_stats("START TRANSFORMER LOAD")

try:
    quant_cfg = DiffusersBitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch_dtype,
        bnb_4bit_use_double_quant=True,
    )

    transformer = AutoModel.from_pretrained(
        model_id,
        cache_dir=model_cache,
        subfolder="transformer",
        quantization_config=quant_cfg,
        torch_dtype=torch_dtype,
        device_map=device,
    )
    log("✅ Transformer loaded successfully.")

except Exception as e:
    log(f"❌ Transformer load failed: {e}")
    transformer = None

log_system_stats("AFTER TRANSFORMER LOAD")

if transformer:
    inspect_transformer(transformer, "Transformer")


# ============================================================
# LOAD TEXT ENCODER
# ============================================================
log("\n===================================================")
log("🔧 LOADING TEXT ENCODER")
log("===================================================")

log_system_stats("START TEXT ENCODER LOAD")

try:
    quant_cfg2 = TransformersBitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch_dtype,
        bnb_4bit_use_double_quant=True,
    )

    text_encoder = AutoModel.from_pretrained(
        model_id,
        cache_dir=model_cache,
        subfolder="text_encoder",
        quantization_config=quant_cfg2,
        torch_dtype=torch_dtype,
        device_map=device,
    )
    log("✅ Text encoder loaded successfully.")

except Exception as e:
    log(f"❌ Text encoder load failed: {e}")
    text_encoder = None

log_system_stats("AFTER TEXT ENCODER LOAD")

if text_encoder:
    inspect_transformer(text_encoder, "Text Encoder")


# ============================================================
# BUILD PIPELINE
# ============================================================
log("\n===================================================")
log("🔧 BUILDING PIPELINE")
log("===================================================")

log_system_stats("START PIPELINE BUILD")

try:
    pipe = ZImagePipeline.from_pretrained(
        model_id,
        transformer=transformer,
        text_encoder=text_encoder,
        torch_dtype=torch_dtype,
        attn_implementation="kernels-community/vllm-flash-attn3",
    )
    pipe.to(device)
    log("✅ Pipeline built successfully.")

except Exception as e:
    log(f"❌ Pipeline build failed: {e}")
    pipe = None

log_system_stats("AFTER PIPELINE BUILD")





from PIL import Image
import torch




def safe_generate_with_latents(
 transformer,
 vae,
 text_encoder,
 tokenizer,
 scheduler,
 pipe,
 prompt,
 height,
 width,
 steps,
 guidance_scale,
 negative_prompt,
 num_images_per_prompt,
 generator,
 cfg_normalization,
 cfg_truncation,
 max_sequence_length,
 ):

 try:

  latents_or_images = generate(
  transformer=transformer,
  vae=vae,
  text_encoder=text_encoder,
  tokenizer=tokenizer,
  scheduler=scheduler,
  prompt=prompt,
  height=height,
  width=width,
  num_inference_steps=steps,
  guidance_scale=guidance_scale,
  negative_prompt=negative_prompt,
  num_images_per_prompt=num_images_per_prompt,
  generator=generator,
  cfg_normalization=cfg_normalization,
  cfg_truncation=cfg_truncation,
  max_sequence_length=max_sequence_length,
  output_type="latent", # IMPORTANT
  )

  return latents_or_images, None

 except Exception as e:
    return None, e











def safe_get_latents(pipe, height, width, generator, device, LOGS):
    """
    Attempts multiple ways to get latents.
    Returns a valid tensor even if pipeline hides UNet.
    """
    # Try official prepare_latents
    try:
        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
            num_channels = pipe.unet.in_channels
            latents = pipe.prepare_latents(
                batch_size=1,
                num_channels=num_channels,
                height=height,
                width=width,
                dtype=torch.float32,
                device=device,
                generator=generator
            )
            LOGS.append("✅ Latents extracted using official prepare_latents.")
            return latents
    except Exception as e:
        LOGS.append(f"⚠️ Official latent extraction failed: {e}")

    # Try hidden internal attribute
    try:
        if hasattr(pipe, "_default_latents"):
            LOGS.append("⚠️ Using hidden _default_latents.")
            return pipe._default_latents
    except:
        pass

    # Fallback: raw Gaussian tensor
    try:
        LOGS.append("⚠️ Using raw Gaussian latents fallback.")
        return torch.randn(
            (1, 4, height // 8, width // 8),
            generator=generator,
            device=device,
            dtype=torch.float32
        )
    except Exception as e:
        LOGS.append(f"⚠️ Gaussian fallback failed: {e}")

    LOGS.append("❗ Using CPU hard fallback latents.")
    return torch.randn((1, 4, height // 8, width // 8))


# --------------------------
# Main generation function
# --------------------------
@spaces.GPU
def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
 LOGS = []
 def log(msg):
  LOGS.append(msg)
  print(msg)


 generator = torch.Generator(device).manual_seed(int(seed))
 log("🎨 START IMAGE GENERATION")
# ==========================================================
# 🧪 TRY ADVANCED LATENT GENERATOR (Your original generate())
# ==========================================================
 latents, latent_err = safe_generate_with_latents(
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    scheduler=scheduler,
    pipe=pipe,
    prompt=prompt,
    height=height,
    width=width,
    steps=steps,
    guidance_scale=guidance_scale,
    negative_prompt="",
    num_images_per_prompt=1,
    generator=generator,
    cfg_normalization=False,
    cfg_truncation=1.0,
    max_sequence_length=4096,
   )

 if latent_err is None:
    log("✅ Latent generator succeeded.")
    try:
        # Decode latents to image
        shift_factor = getattr(vae.config, "shift_factor", 0.0) or 0.0
        dec = (latents.to(vae.dtype) / vae.config.scaling_factor) + shift_factor
        image = vae.decode(dec, return_dict=False)[0]

        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
        image = (image * 255).round().astype("uint8")
        from PIL import Image
        image = Image.fromarray(image[0])

        log("🟢 Final image decoded from latent generator.")
        return image, latents, LOGS

    except Exception as decode_error:
        log(f"⚠️ Latent decode failed: {decode_error}")
        log("🔁 Falling back to standard pipeline...")

 else:
    log(f"⚠️ Latent generator failed: {latent_err}")
    log("🔁 Switching to standard pipeline...")

# ==========================================================
# 🟩 STANDARD PIPELINE FALLBACK (Never fails)
# ==========================================================
 try:
    output = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        generator=generator,
    )
    image = output.images[0]
    log("🟢 Standard pipeline succeeded.")
    return image, None, LOGS

 except Exception as e:
    log(f"❌ Standard pipeline failed: {e}")
    return None, None, LOGS

# --------------------------
# Helper: Safe latent extractor
# --------------------------
def safe_get_latents0(pipe, height, width, generator, device, LOGS):
    """
    Attempts multiple ways to get latents.
    Returns a valid tensor even if pipeline hides UNet.
    """
    # Try official prepare_latents
    try:
        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
            num_channels = pipe.unet.in_channels
            latents = pipe.prepare_latents(
                batch_size=1,
                num_channels=num_channels,
                height=height,
                width=width,
                dtype=torch.float32,
                device=device,
                generator=generator
            )
            LOGS.append("✅ Latents extracted using official prepare_latents.")
            return latents
    except Exception as e:
        LOGS.append(f"⚠️ Official latent extraction failed: {e}")

    # Try hidden internal attribute
    try:
        if hasattr(pipe, "_default_latents"):
            LOGS.append("⚠️ Using hidden _default_latents.")
            return pipe._default_latents
    except:
        pass

    # Fallback: raw Gaussian tensor
    try:
        LOGS.append("⚠️ Using raw Gaussian latents fallback.")
        return torch.randn(
            (1, 4, height // 8, width // 8),
            generator=generator,
            device=device,
            dtype=torch.float32
        )
    except Exception as e:
        LOGS.append(f"⚠️ Gaussian fallback failed: {e}")

    LOGS.append("❗ Using CPU hard fallback latents.")
    return torch.randn((1, 4, height // 8, width // 8))


# --------------------------
# Main generation function
# --------------------------
@spaces.GPU
def generate_image0(prompt, height, width, steps, seed, guidance_scale=0.0):
    LOGS = []
    latents = None
    image = None
    gallery = []

    # placeholder image if all fails
    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
    print(prompt)

    try:
        generator = torch.Generator(device).manual_seed(int(seed))

        # -------------------------------
        # Try advanced latent extraction
        # -------------------------------
        try:
            latents = safe_get_latents(pipe, height, width, generator, device, LOGS)

            output = pipe(
                prompt=prompt,
                height=height,
                width=width,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
                latents=latents
            )

            image = output.images[0]
            gallery = [image]
            LOGS.append("✅ Advanced latent pipeline succeeded.")

        except Exception as e:
            LOGS.append(f"⚠️ Latent mode failed: {e}")
            LOGS.append("🔁 Switching to standard pipeline...")

            try:
                output = pipe(
                    prompt=prompt,
                    height=height,
                    width=width,
                    num_inference_steps=steps,
                    guidance_scale=guidance_scale,
                    generator=generator,
                )
                image = output.images[0]
                gallery = [image]
                LOGS.append("✅ Standard pipeline succeeded.")

            except Exception as e2:
                LOGS.append(f"❌ Standard pipeline failed: {e2}")
                image = placeholder
                gallery = [image]

        return image, gallery, LOGS

    except Exception as e:
        LOGS.append(f"❌ Total failure: {e}")
        return placeholder, [placeholder], LOGS

# ============================================================
# UI
# ============================================================

with gr.Blocks(title="Z-Image- experiment - dont run")as demo:
  gr.Markdown("# **🚀 do not run Z-Image-Turbo — Final Image & Latents**")


  with gr.Row():
    with gr.Column(scale=1):
        prompt = gr.Textbox(label="Prompt", value="boat in Ocean")
        height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
        width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
        steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
        seed = gr.Number(value=42, label="Seed")
        run_btn = gr.Button("Generate Image")

    with gr.Column(scale=1):
        final_image = gr.Image(label="Final Image")
        latent_gallery = gr.Gallery(
           label="Latent Steps",
                columns=4,
              height=256,
             preview=True
              )

        logs_box = gr.Textbox(label="Logs", lines=15)

    run_btn.click(
      generate_image,
      inputs=[prompt, height, width, steps, seed],
      outputs=[final_image, latent_gallery, logs_box]
     )



demo.launch()