# kaggle_template.py
# =====================================================
# KAGGLE IMAGE WORKER (FLUX.1-SCHNELL)
# =====================================================
import os, torch, gc, subprocess, sys

# Install bitsandbytes if missing (Critical for T4 GPU)
try:
    import bitsandbytes
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "bitsandbytes"])

from diffusers import FluxPipeline, FluxTransformer2DModel, BitsAndBytesConfig
from pathlib import Path

# --- CONFIG ---
# The automation script will inject the prompts here automatically
PROMPTS = [
    # {{PROMPTS_PLACEHOLDER}}
]

OUTPUT_DIR = Path("/kaggle/working/images")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# --- MODEL LOADING ---
print("📦 Loading Quantized Flux...")

# 4-bit config to fit Flux on Kaggle T4s
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

# Load Transformer
transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    subfolder="transformer",
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True
)

# Load Pipeline
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)

pipe.enable_model_cpu_offload()

# --- GENERATION ---
print(f"🎨 Generating {len(PROMPTS)} images...")

for i, prompt in enumerate(PROMPTS, 1):
    print(f"   Frame {i}/{len(PROMPTS)}...")
    gc.collect()
    torch.cuda.empty_cache()
    
    # Generate
    image = pipe(
        prompt=prompt,
        width=1072, 
        height=1920,
        num_inference_steps=4, # Schnell is fast
        guidance_scale=1.0,
        max_sequence_length=512,
    ).images[0]
    
    # Save
    save_path = OUTPUT_DIR / f"{i:02d}.png"
    image.save(save_path)
    print(f"   ✅ Saved: {save_path.name}")

print("🏁 Job Complete")