app_quant_latent.py

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

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):
    """
    Convert a latent tensor to a PIL image using pipe.vae
    """
    try:
        img_tensor = pipe.vae.decode(latent)
        img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
        pil_img = T.ToPILImage()(img_tensor[0].cpu())  # <--- single image
        return pil_img
    except Exception as e:
        log(f"⚠️ Failed to decode latent: {e}")
        # fallback blank image
        return Image.new("RGB", (latent.shape[-1]*8, latent.shape[-2]*8), color=(255,255,255))



# ============================================================
# SAFE TRANSFORMER INSPECTION
# ============================================================
def inspect_transformer(model, name):
    log(f"\n🔍🔍 FULL TRANSFORMER DEBUG DUMP: {name}")
    log("=" * 80)

    try:
        log(f"Model class          : {model.__class__.__name__}")
        log(f"DType                : {getattr(model, 'dtype', 'unknown')}")
        log(f"Device               : {next(model.parameters()).device}")
        log(f"Requires Grad?       : {any(p.requires_grad for p in model.parameters())}")

        # Check quantization
        if hasattr(model, "is_loaded_in_4bit"):
            log(f"4bit Quantization    : {model.is_loaded_in_4bit}")
        if hasattr(model, "is_loaded_in_8bit"):
            log(f"8bit Quantization    : {model.is_loaded_in_8bit}")

        # Find blocks
        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
        blocks = None
        chosen_attr = None

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

        log(f"Block container attr : {chosen_attr}")

        if blocks is None:
            log("⚠️ No valid block container found.")
            return

        if not hasattr(blocks, "__len__"):
            log("⚠️ Blocks exist but not iterable.")
            return

        total = len(blocks)
        log(f"Total Blocks         : {total}")
        log("-" * 80)

        # Inspect first N blocks
        N = min(20, total)
        for i in range(N):
            block = blocks[i]
            log(f"\n🧩 Block [{i}/{total-1}]")
            log(f"Class: {block.__class__.__name__}")

            # Print submodules
            for n, m in block.named_children():
                log(f"  ├─ {n}: {m.__class__.__name__}")

            # Print attention related
            if hasattr(block, "attn"):
                attn = block.attn
                log(f"  ├─ Attention: {attn.__class__.__name__}")
                log(f"  │   Heads     : {getattr(attn, 'num_heads', 'unknown')}")
                log(f"  │   Dim       : {getattr(attn, 'hidden_size', 'unknown')}")
                log(f"  │   Backend   : {getattr(attn, 'attention_backend', 'unknown')}")

            # Device + dtype info
            try:
                dev = next(block.parameters()).device
                log(f"  ├─ Device     : {dev}")
            except StopIteration:
                pass

            try:
                dt = next(block.parameters()).dtype
                log(f"  ├─ DType      : {dt}")
            except StopIteration:
                pass

        log("\n🔚 END TRANSFORMER DEBUG DUMP")
        log("=" * 80)

    except Exception as e:
        log(f"❌ ERROR IN INSPECTOR: {e}")
import torch
import time

# ---------- UTILITY ----------
def pretty_header(title):
    log("\n\n" + "=" * 80)
    log(f"🎛️  {title}")
    log("=" * 80 + "\n")


# ---------- MEMORY ----------
def get_vram(prefix=""):
    try:
        allocated = torch.cuda.memory_allocated() / 1024**2
        reserved = torch.cuda.memory_reserved() / 1024**2
        log(f"{prefix}Allocated VRAM : {allocated:.2f} MB")
        log(f"{prefix}Reserved VRAM  : {reserved:.2f} MB")
    except:
        log(f"{prefix}VRAM: CUDA not available")


# ---------- MODULE INSPECT ----------
def inspect_module(name, module):
    pretty_header(f"🔬 Inspecting {name}")

    try:
        log(f"📦 Class          : {module.__class__.__name__}")
        log(f"🔢 DType          : {getattr(module, 'dtype', 'unknown')}")
        log(f"💻 Device         : {next(module.parameters()).device}")
        log(f"🧮 Params         : {sum(p.numel() for p in module.parameters()):,}")

        # Quantization state
        if hasattr(module, "is_loaded_in_4bit"):
            log(f"⚙️  4-bit QLoRA     : {module.is_loaded_in_4bit}")
        if hasattr(module, "is_loaded_in_8bit"):
            log(f"⚙️  8-bit load      : {module.is_loaded_in_8bit}")

        # Attention backend (DiT)
        if hasattr(module, "set_attention_backend"):
            try:
                attn = getattr(module, "attention_backend", None)
                log(f"🚀 Attention Backend: {attn}")
            except:
                pass

        # Search for blocks
        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
        blocks = None
        chosen_attr = None

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

        log(f"\n📚 Block Container : {chosen_attr}")

        if blocks is None:
            log("⚠️ No block structure found")
            return

        if not hasattr(blocks, "__len__"):
            log("⚠️ Blocks exist but are not iterable")
            return

        total = len(blocks)
        log(f"🔢 Total Blocks   : {total}\n")

        # Inspect first 15 blocks
        N = min(15, total)

        for i in range(N):
            blk = blocks[i]
            log(f"\n🧩 Block [{i}/{total-1}] — {blk.__class__.__name__}")

            for n, m in blk.named_children():
                log(f"   ├─ {n:<15} {m.__class__.__name__}")

            # Attention details
            if hasattr(blk, "attn"):
                a = blk.attn
                log(f"   ├─ Attention")
                log(f"   │    Heads    : {getattr(a, 'num_heads', 'unknown')}")
                log(f"   │    Dim      : {getattr(a, 'hidden_size', 'unknown')}")
                log(f"   │    Backend  : {getattr(a, 'attention_backend', 'unknown')}")

            # Device / dtype
            try:
                log(f"   ├─ Device     : {next(blk.parameters()).device}")
                log(f"   ├─ DType      : {next(blk.parameters()).dtype}")
            except StopIteration:
                pass

        get_vram("   ▶ ")

    except Exception as e:
        log(f"❌ Module inspect error: {e}")


# ---------- LORA INSPECTION ----------
def inspect_loras(pipe):
    pretty_header("🧩 LoRA ADAPTERS")

    try:
        if not hasattr(pipe, "lora_state_dict") and not hasattr(pipe, "adapter_names"):
            log("⚠️ No LoRA system detected.")
            return

        if hasattr(pipe, "adapter_names"):
            names = pipe.adapter_names
            log(f"Available Adapters: {names}")

        if hasattr(pipe, "active_adapters"):
            log(f"Active Adapters   : {pipe.active_adapters}")

        if hasattr(pipe, "lora_scale"):
            log(f"LoRA Scale        : {pipe.lora_scale}")

        # LoRA modules
        if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "modules"):
            for name, module in pipe.transformer.named_modules():
                if "lora" in name.lower():
                    log(f"   🔧 LoRA Module: {name}  ({module.__class__.__name__})")

    except Exception as e:
        log(f"❌ LoRA inspect error: {e}")


# ---------- PIPELINE INSPECTOR ----------
def debug_pipeline(pipe):
    pretty_header("🚀 FULL PIPELINE DEBUGGING")

    try:
        log(f"Pipeline Class       : {pipe.__class__.__name__}")
        log(f"Attention Impl       : {getattr(pipe, 'attn_implementation', 'unknown')}")
        log(f"Device               : {pipe.device}")
    except:
        pass

    get_vram("▶ ")

    # Inspect TRANSFORMER
    if hasattr(pipe, "transformer"):
        inspect_module("Transformer", pipe.transformer)

    # Inspect TEXT ENCODER
    if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
        inspect_module("Text Encoder", pipe.text_encoder)

    # Inspect UNET (if ZImage pipeline has it)
    if hasattr(pipe, "unet"):
        inspect_module("UNet", pipe.unet)

    # LoRA adapters
    inspect_loras(pipe)

    pretty_header("🎉 END DEBUG REPORT")



# ============================================================
# 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,
       
    )
    # If transformer supports setting backend, prefer flash-3
    try:
        if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "set_attention_backend"):
            pipe.transformer.set_attention_backend("_flash_3")
            log("✅ transformer.set_attention_backend('_flash_3') called")
    except Exception as _e:
        log(f"⚠️ set_attention_backend failed: {_e}")

    # default LoRA load (keeps your existing behaviour)
    try:
        pipe.load_lora_weights("rahul7star/ZImageLora",
                               weight_name="NSFW/doggystyle_pov.safetensors", adapter_name="lora")
        pipe.set_adapters(["lora",], adapter_weights=[1.])
        pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
    except Exception as _e:
        log(f"⚠️ Default LoRA load failed: {_e}")

    debug_pipeline(pipe)
    #  pipe.unload_lora_weights()
    pipe.to("cuda")
    log("✅ Pipeline built successfully.")
    LOGS += log("Pipeline build completed.") + "\n"
except Exception as e:
    log(f"❌ Pipeline build failed: {e}")
    log(traceback.format_exc())
    pipe = None

log_system_stats("AFTER PIPELINE BUILD")


# -----------------------------
# Monkey-patch prepare_latents (safe)
# -----------------------------
if pipe is not None and hasattr(pipe, "prepare_latents"):
    original_prepare_latents = pipe.prepare_latents

    def logged_prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        try:
            result_latents = original_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents)
            log_msg = f"🔹 prepare_latents called | shape={result_latents.shape}, dtype={result_latents.dtype}, device={result_latents.device}"
            if hasattr(self, "_latents_log"):
                self._latents_log.append(log_msg)
            else:
                self._latents_log = [log_msg]
            return result_latents
        except Exception as e:
            log(f"⚠️ prepare_latents wrapper failed: {e}")
            raise

    # apply patch safely
    try:
        pipe.prepare_latents = logged_prepare_latents.__get__(pipe)
        log("✅ prepare_latents monkey-patched")
    except Exception as e:
        log(f"⚠️ Failed to attach prepare_latents patch: {e}")
else:
    log("❌ WARNING: Pipe not initialized or prepare_latents missing; skipping prepare_latents patch")


from PIL import Image
import torch

# --------------------------
# Helper: Safe latent extractor
# --------------------------
def safe_get_latents(pipe, height, width, generator, device, LOGS):
    """
    Safely prepare latents for any ZImagePipeline variant.
    Returns latents tensor, logs issues instead of failing.
    """
    try:
        # Determine number of channels
        num_channels = 4  # default fallback
        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
            num_channels = pipe.unet.in_channels
        elif hasattr(pipe, "vae") and hasattr(pipe.vae, "latent_channels"):
            num_channels = pipe.vae.latent_channels  # some pipelines define this
        LOGS.append(f"🔹 Using num_channels={num_channels} for latents")

        latents = pipe.prepare_latents(
            batch_size=1,
            num_channels_latents=num_channels,
            height=height,
            width=width,
            dtype=torch.float32,
            device=device,
            generator=generator,
        )

        LOGS.append(f"🔹 Latents shape: {latents.shape}, dtype: {latents.dtype}, device: {latents.device}")
        return latents
    except Exception as e:
        LOGS.append(f"⚠️ Latent extraction failed: {e}")
        # fallback: guess a safe shape
        fallback_channels = 16  # try standard default for ZImage pipelines
        latents = torch.randn((1, fallback_channels, height // 8, width // 8),
                              generator=generator, device=device)
        LOGS.append(f"🔹 Using fallback random latents shape: {latents.shape}")
        return latents

# --------------------------
# Main generation function (kept exactly as your logic)
# --------------------------
from huggingface_hub import HfApi, HfFolder
import torch
import os

HF_REPO_ID = "rahul7star/Zstudio-latent"  # Model repo
HF_TOKEN = HfFolder.get_token()  # Make sure you are logged in via `huggingface-cli login`

def upload_latents_to_hf(latent_dict, filename="latents.pt"):
    local_path = f"/tmp/{filename}"
    torch.save(latent_dict, local_path)
    try:
        api = HfApi()
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=filename,
            repo_id=HF_REPO_ID,
            token=HF_TOKEN,
            repo_type="model"  # since this is a model repo
        )
        os.remove(local_path)
        return f"https://huggingface.co/{HF_REPO_ID}/resolve/main/{filename}"
    except Exception as e:
        os.remove(local_path)
        raise e



import asyncio
import torch
from PIL import Image

async def async_upload_latents(latent_dict, filename, LOGS):
    try:
        hf_url = await upload_latents_to_hf(latent_dict, filename=filename)  # assume this can be async
        LOGS.append(f"🔹 All preview latents uploaded: {hf_url}")
    except Exception as e:
        LOGS.append(f"⚠️ Failed to upload all preview latents: {e}")


# this code genetae all frame for latest GPU expseinve bt decide fails sp use this later         
@spaces.GPU
def generate_image_all_latents(prompt, height, width, steps, seed, guidance_scale=0.0):
    LOGS = []
    device = "cpu"  # FORCE CPU
    generator = torch.Generator(device).manual_seed(int(seed))

    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
    latent_gallery = []
    final_gallery = []

    last_four_latents = []  # we only upload 4

    # --------------------------------------------------
    # LATENT PREVIEW GENERATION (CPU MODE)
    # --------------------------------------------------
    try:
        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
        latents = latents.to("cpu")  # keep EVERYTHING CPU

        timestep_count = len(pipe.scheduler.timesteps)
        preview_every = max(1, timestep_count // 10)

        for i, t in enumerate(pipe.scheduler.timesteps):

            # -------------- decode latent preview --------------
            try:
                with torch.no_grad():
                    latent_cpu = latents.to(pipe.vae.dtype)  # match VAE dtype
                    decoded = pipe.vae.decode(latent_cpu).sample  # [1,3,H,W]
                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
                    decoded = decoded[0].permute(1,2,0).cpu().numpy()
                    latent_img = Image.fromarray((decoded * 255).astype("uint8"))
            except Exception:
                latent_img = placeholder
                LOGS.append("⚠️ Latent preview decode failed.")

            latent_gallery.append(latent_img)

            # store last 4 latent states
            if len(last_four_latents) >= 4:
                last_four_latents.pop(0)
            last_four_latents.append(latents.cpu().clone())

            # UI preview yields
            if i % preview_every == 0:
                yield None, latent_gallery, LOGS

        # --------------------------------------------------
        # UPLOAD LAST 4 LATENTS (SYNC)
        # --------------------------------------------------
        try:
            upload_dict = {
                "last_4_latents": last_four_latents,
                "prompt": prompt,
                "seed": seed
            }

            hf_url = upload_latents_to_hf(
                upload_dict,
                filename=f"latents_last4_{seed}.pt"
            )

            LOGS.append(f"🔹 Uploaded last 4 latents: {hf_url}")

        except Exception as e:
            LOGS.append(f"⚠️ Failed to upload latents: {e}")

    except Exception as e:
        LOGS.append(f"⚠️ Latent generation failed: {e}")
        latent_gallery.append(placeholder)
        yield None, latent_gallery, LOGS

    # --------------------------------------------------
    # FINAL IMAGE - UNTOUCHED
    # --------------------------------------------------
    try:
        output = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        final_img = output.images[0]
        LOGS.append("✅ Standard pipeline succeeded.")

        yield final_img, latent_gallery, LOGS

    except Exception as e2:
        LOGS.append(f"❌ Standard pipeline failed: {e2}")
        yield placeholder, latent_gallery, LOGS

@spaces.GPU
def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
    LOGS = []
    device = "cuda"
    cpu_device = "cpu"
    generator = torch.Generator(device).manual_seed(int(seed))

    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
    latent_gallery = []
    final_gallery = []

    last_latents = []  # store last 5 preview latents on CPU

    try:
        # --- Initial latents ---
        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
        latents = latents.float().to(cpu_device)  # move to CPU

        num_previews = min(10, steps)
        preview_indices = torch.linspace(0, steps - 1, num_previews).long()

        for i, step_idx in enumerate(preview_indices):
            try:
                with torch.no_grad():
                    # --- Z-Image Turbo-style denoise simulation ---
                    t = 1.0 - (i / num_previews)  # linear decay [1.0 -> 0.0]
                    noise_scale = t ** 0.5  # reduce noise over steps (sqrt for smoother)
                    denoise_latent = latents * t + torch.randn_like(latents) * noise_scale

                    # Move to VAE device & dtype
                    denoise_latent = denoise_latent.to(pipe.vae.device).to(pipe.vae.dtype)

                    # Decode latent to image
                    decoded = pipe.vae.decode(denoise_latent, return_dict=False)[0]
                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
                    decoded = decoded.cpu().permute(0, 2, 3, 1).float().numpy()
                    decoded = (decoded * 255).round().astype("uint8")
                    latent_img = Image.fromarray(decoded[0])

            except Exception as e:
                LOGS.append(f"⚠️ Latent preview decode failed: {e}")
                latent_img = placeholder

            latent_gallery.append(latent_img)

            # Keep last 5 latents only
            last_latents.append(denoise_latent.cpu().clone())
            if len(last_latents) > 5:
                last_latents.pop(0)

            # Show only last 5 previews in UI
            yield None, latent_gallery[-5:], LOGS

        # Optionally: upload last 5 latents
        # latent_dict = {"latents": last_latents, "prompt": prompt, "seed": seed}
        # hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_last5_{seed}.pt")
        # LOGS.append(f"🔹 Last 5 latents uploaded: {hf_url}")

    except Exception as e:
        LOGS.append(f"⚠️ Latent generation failed: {e}")
        latent_gallery.append(placeholder)
        yield None, latent_gallery[-5:], LOGS

    # --- Final image on GPU ---
    try:
        output = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        final_img = output.images[0]
        final_gallery.append(final_img)
        latent_gallery.append(final_img)
        LOGS.append("✅ Standard pipeline succeeded.")
        yield final_img, latent_gallery[-5:] + [final_img], LOGS  # last 5 previews + final

    except Exception as e2:
        LOGS.append(f"❌ Standard pipeline failed: {e2}")
        final_gallery.append(placeholder)
        latent_gallery.append(placeholder)
        yield placeholder, latent_gallery[-5:] + [placeholder], LOGS



# this is astable vesopn tha can gen final and a noise to latent
@spaces.GPU
def generate_image_verygood_realnoise(prompt, height, width, steps, seed, guidance_scale=0.0):
    LOGS = []
    device = "cuda"
    generator = torch.Generator(device).manual_seed(int(seed))

    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
    latent_gallery = []
    final_gallery = []

    # --- Generate latent previews ---
    try:
        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
        latents = latents.float()  # keep float32 until decode

        num_previews = min(10, steps)
        preview_steps = torch.linspace(0, 1, num_previews)

        for alpha in preview_steps:
            try:
                with torch.no_grad():
                    # Simulate denoising progression like Z-Image Turbo
                    preview_latent = latents * alpha + latents * 0  # optional: simple progression

                    # Move to same device and dtype as VAE
                    preview_latent = preview_latent.to(pipe.vae.device).to(pipe.vae.dtype)

                    # Decode
                    decoded = pipe.vae.decode(preview_latent, return_dict=False)[0]

                    # Convert to PIL following same logic as final image
                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
                    decoded = decoded.cpu().permute(0, 2, 3, 1).float().numpy()
                    decoded = (decoded * 255).round().astype("uint8")
                    latent_img = Image.fromarray(decoded[0])

            except Exception as e:
                LOGS.append(f"⚠️ Latent preview decode failed: {e}")
                latent_img = placeholder

            latent_gallery.append(latent_img)
            yield None, latent_gallery, LOGS

    except Exception as e:
        LOGS.append(f"⚠️ Latent generation failed: {e}")
        latent_gallery.append(placeholder)
        yield None, latent_gallery, LOGS

    # --- Final image: untouched ---
    try:
        output = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        final_img = output.images[0]
        final_gallery.append(final_img)
        latent_gallery.append(final_img)  # fallback preview
        LOGS.append("✅ Standard pipeline succeeded.")
        yield final_img, latent_gallery, LOGS

    except Exception as e2:
        LOGS.append(f"❌ Standard pipeline failed: {e2}")
        final_gallery.append(placeholder)
        latent_gallery.append(placeholder)
        yield placeholder, latent_gallery, LOGS




# DO NOT TOUCH this is astable vesopn tha can gen final and a noise to latent with latent upload to repo
@spaces.GPU
def generate_image_safe(prompt, height, width, steps, seed, guidance_scale=0.0):
    LOGS = []
    device = "cuda"
    generator = torch.Generator(device).manual_seed(int(seed))

    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
    latent_gallery = []
    final_gallery = []

    # --- Generate latent previews in a loop ---
    try:
        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)

        # Convert latents to float32 if necessary
        if latents.dtype != torch.float32:
            latents = latents.float()

        # Loop for multiple previews before final image
        num_previews = min(10, steps)  # show ~10 previews
        preview_steps = torch.linspace(0, 1, num_previews)

        for i, alpha in enumerate(preview_steps):
            try:
                with torch.no_grad():
                    # Simple noise interpolation for preview (simulate denoising progress)
                    preview_latent = latents * alpha + torch.randn_like(latents) * (1 - alpha)
                    # Decode to PIL
                    latent_img_tensor = pipe.vae.decode(preview_latent).sample  # [1,3,H,W]
                    latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
                    latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
                    latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
            except Exception as e:
                LOGS.append(f"⚠️ Latent preview decode failed: {e}")
                latent_img = placeholder

            latent_gallery.append(latent_img)
            yield None, latent_gallery, LOGS  # update Gradio with intermediate preview

        # Save final latents to HF
        latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
        try:
            hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
            LOGS.append(f"🔹 Latents uploaded: {hf_url}")
        except Exception as e:
            LOGS.append(f"⚠️ Failed to upload latents: {e}")

    except Exception as e:
        LOGS.append(f"⚠️ Latent generation failed: {e}")
        latent_gallery.append(placeholder)
        yield None, latent_gallery, LOGS

    # --- Final image: untouched standard pipeline ---
    try:
        output = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        final_img = output.images[0]
        final_gallery.append(final_img)
        latent_gallery.append(final_img)  # fallback preview if needed
        LOGS.append("✅ Standard pipeline succeeded.")
        yield final_img, latent_gallery, LOGS

    except Exception as e2:
        LOGS.append(f"❌ Standard pipeline failed: {e2}")
        final_gallery.append(placeholder)
        latent_gallery.append(placeholder)
        yield placeholder, latent_gallery, LOGS







with gr.Blocks(title="Z-Image-Turbo") as demo:
    gr.Markdown("# 🎨 DO NOT RUN THIS ")
    with gr.Tabs():
        with gr.TabItem("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
                    )

        with gr.TabItem("Logs"):
            logs_box = gr.Textbox(label="All Logs", lines=25)

    # New UI: LoRA repo textbox, dropdown, refresh & rebuild
    with gr.Row():
        lora_repo = gr.Textbox(label="LoRA Repo (HF id)", value="rahul7star/ZImageLora", placeholder="e.g. rahul7star/ZImageLora")
        lora_dropdown = gr.Dropdown(choices=[], label="LoRA files (from local cache)")
        refresh_lora_btn = gr.Button("Refresh LoRA List")
        rebuild_pipe_btn = gr.Button("Rebuild pipeline (use selected LoRA)")

    # Refresh callback: repopulate dropdown from repo text
    def refresh_lora_list(repo_name):
        try:
            files = list_loras_from_repo(repo_name)
            if not files:
                return gr.update(choices=[], value=None)
            return gr.update(choices=files, value=files[0])
        except Exception as e:
            log(f"⚠️ refresh_lora_list failed: {e}")
            return gr.update(choices=[], value=None)

    refresh_lora_btn.click(refresh_lora_list, inputs=[lora_repo], outputs=[lora_dropdown])

    # Rebuild callback: build pipeline with selected lora file path (if any)
    def rebuild_pipeline_with_lora(lora_path, repo_name):
        global pipe, LOGS
        try:
            log(f"🔄 Rebuilding pipeline using LoRA repo={repo_name} file={lora_path}")
            # call existing logic to rebuild: attempt to create new pipeline then load lora file
            pipe = ZImagePipeline.from_pretrained(
                model_id,
                transformer=transformer,
                text_encoder=text_encoder,
                torch_dtype=torch_dtype,
            )
            # try set backend
            try:
                if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "set_attention_backend"):
                    pipe.transformer.set_attention_backend("_flash_3")
            except Exception as _e:
                log(f"⚠️ set_attention_backend failed during rebuild: {_e}")

            # load selected lora if provided
            if lora_path:
                weight_name_to_use = None

                # If dropdown provided a relative-style path (contains a slash or no leading /),
                # use it directly as weight_name (HF expects "path/inside/repo.safetensors")
                if ("/" in lora_path) and not os.path.isabs(lora_path):
                    weight_name_to_use = lora_path
                else:
                    # It might be an absolute path in cache; try to compute relative path to repo cache root
                    abs_path = lora_path if os.path.isabs(lora_path) else None
                    if abs_path and os.path.exists(abs_path):
                        # attempt to find repo-root-ish substring in abs_path
                        repo_variants = [
                            repo_name.replace("/", "--"),
                            repo_name.replace("/", "-"),
                            repo_name.replace("/", "_"),
                            repo_name.split("/")[-1],
                        ]
                        chosen_base = None
                        for v in repo_variants:
                            idx = abs_path.find(v)
                            if idx != -1:
                                chosen_base = abs_path[: idx + len(v)]
                                break
                        if chosen_base:
                            try:
                                rel = os.path.relpath(abs_path, chosen_base)
                                if rel and not rel.startswith(".."):
                                    weight_name_to_use = rel.replace(os.sep, "/")
                            except Exception:
                                weight_name_to_use = None

                    # fallback to basename
                    if weight_name_to_use is None:
                        weight_name_to_use = os.path.basename(lora_path)

                # Now attempt to load
                try:
                    pipe.load_lora_weights(repo_name or "rahul7star/ZImageLora",
                                           weight_name=weight_name_to_use,
                                           adapter_name="lora")
                    pipe.set_adapters(["lora"], adapter_weights=[1.])
                    pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
                    log(f"✅ Loaded LoRA weight: {weight_name_to_use} from repo {repo_name}")
                except Exception as _e:
                    log(f"⚠️ Failed to load selected LoRA during rebuild using weight_name='{weight_name_to_use}': {_e}")
                    # as last resort, try loading using basename
                    try:
                        fallback_name = os.path.basename(lora_path)
                        pipe.load_lora_weights(repo_name or "rahul7star/ZImageLora",
                                               weight_name=fallback_name,
                                               adapter_name="lora")
                        pipe.set_adapters(["lora"], adapter_weights=[1.])
                        pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
                        log(f"✅ Fallback loaded LoRA weight basename: {fallback_name}")
                    except Exception as _e2:
                        log(f"❌ Fallback LoRA load also failed: {_e2}")

            # finalize
            debug_pipeline(pipe)
            pipe.to("cuda")
            # re-attach monkey patch safely
            if pipe is not None and hasattr(pipe, "prepare_latents"):
                try:
                    original_prepare = pipe.prepare_latents
                    def logged_prepare(self, *args, **kwargs):
                        lat = original_prepare(*args, **kwargs)
                        msg = f"🔹 prepare_latents called | shape={lat.shape}, dtype={lat.dtype}"
                        if hasattr(self, "_latents_log"):
                            self._latents_log.append(msg)
                        else:
                            self._latents_log = [msg]
                        return lat
                    pipe.prepare_latents = logged_prepare.__get__(pipe)
                    log("✅ Re-applied prepare_latents monkey patch after rebuild")
                except Exception as _e:
                    log(f"⚠️ Could not re-apply prepare_latents patch: {_e}")
            return "\n".join([LOGS, "Rebuild complete."])
        except Exception as e:
            log(f"❌ Rebuild pipeline failed: {e}")
            log(traceback.format_exc())
            return "\n".join([LOGS, f"Rebuild failed: {e}"])

    rebuild_pipe_btn.click(rebuild_pipeline_with_lora, inputs=[lora_dropdown, lora_repo], outputs=[logs_box])

    # Wire the button AFTER all components exist
    
    run_btn.click(
         generate_image,
         inputs=[prompt, height, width, steps, seed],
         outputs=[final_image, latent_gallery, logs_box]
    )


demo.launch()