flowchart TD
    A["📝 User Prompt"] --> B["CLIP Encode\n(Qwen 3 4B)"]
    B --> C["Empty Latent\n(Image Format)"]
    C --> D["KSampler\n(20 steps, euler)"]
    D --> E["VAE Decode\n(ae.safetensors)"]
    E --> F["🖼️ Generated Image"]

    G["Z-Image Turbo FP8\n(Base Model)"] --> D
    H["ModelSamplingAuraFlow\n(Shift: 3.0)"] --> D

    style A fill:#0d1117,stroke:#58a6ff,color:#e6edf3,stroke-width:2px
    style F fill:#0d1117,stroke:#3fb950,color:#e6edf3,stroke-width:2px
    style G fill:#0d1117,stroke:#a371f7,color:#e6edf3,stroke-width:2px
    style H fill:#0d1117,stroke:#f97316,color:#e6edf3,stroke-width:2px

flowchart TD
    subgraph Text["📝 Text Encoding"]
        CLIP["CLIPLoader\nqwen_3_4b.safetensors\n(lumina2)"]
        POS["CLIPTextEncode\nPositive Prompt"]
        NEG["CLIPTextEncode\nNegative Prompt"]
        CLIP --> POS
        CLIP --> NEG
    end

    subgraph Latent["🖼️ Latent Space"]
        EMPTY["EmptyLatentImage\n1024×1024, batch=1"]
    end

    subgraph Model["🧠 Model Loading"]
        UNET["UNETLoader\nFP8 (.safetensors)"]
    end

    subgraph Sample["⚡ Sampling"]
        KS["KSampler\n20 steps · euler\nsimple scheduler · CFG 1.0"]
    end

    subgraph Output["💾 Output"]
        VAE_D["VAEDecode\nae.safetensors"]
        SAVE["SaveImage\nz-image-pro prefix"]
        VAE_D --> SAVE
    end

    VAE["VAELoader\nae.safetensors"]

    UNET --> KS
    POS --> KS
    NEG --> KS
    EMPTY --> KS
    KS --> VAE_D
    VAE --> VAE_D

    style Text fill:#0d1117,stroke:#58a6ff,color:#e6edf3
    style Latent fill:#0d1117,stroke:#a371f7,color:#e6edf3
    style Model fill:#0d1117,stroke:#f0c14b,color:#e6edf3
    style Sample fill:#0d1117,stroke:#f97316,color:#e6edf3
    style Output fill:#0d1117,stroke:#3fb950,color:#e6edf3

sequenceDiagram
    participant U as 👤 User
    participant NB as 📓 Notebook
    participant SRC as 📦 src/ Package
    participant DRV as 💾 Drive Cache
    participant CUI as ⚙️ ComfyUI
    participant GPU as 🖥️ T4 GPU

    U->>NB: Run Cell 1 (Initialize)
    NB->>NB: Clone ZImagePro + ComfyUI repos
    NB->>NB: Install deps + aria2c
    NB->>DRV: Mount Google Drive
    NB->>SRC: Import config, downloader modules
    SRC->>DRV: Check cache status (hit/miss per model)
    alt Cache Hit
        DRV-->>SRC: Models found in Drive (~30s copy)
    else Cache Miss
        SRC->>SRC: aria2c 16-parallel download (~5 min)
        SRC->>DRV: Save models to Drive cache
    end
    SRC-->>NB: ✅ Models ready

    U->>NB: Run Cell 2 (Load & Generate)
    NB->>SRC: load_models()
    SRC->>CUI: Load UNet + CLIP + VAE into VRAM
    CUI-->>SRC: Models ready + VRAM stats
    NB->>SRC: generate_image(params)
    SRC->>CUI: CLIP encode + KSampler + VAE decode
    GPU-->>CUI: Pixel image
    SRC-->>NB: PIL Image + save path + timing
    NB->>NB: Display inline image

    U->>NB: Run Cell 3 (Export)
    NB->>SRC: zip_outputs() + download_zip()
    SRC-->>U: 📥 Z_Image_Pro_Artworks.zip

    U->>NB: Run Cell 4 (Cleanup) [optional]
    NB->>DRV: Clear cache / Clean old outputs
    NB-->>U: ✅ Disk space freed