"""
Sentinel Tiny Image Space — CIFAR-10 diffusion with Sentinel noise schedule
"""
import gradio as gr
import torch
from diffusers import DDPMScheduler, UNet2DModel, DDPMPipeline
from PIL import Image
import numpy as np
import json

# ─── Load Model ──────────────────────────────────────────────────────────────
MODEL_ID = "5dimension/sentinel-tiny-image"

model_status = "⏳ Loading model..."

pipeline = None
try:
    pipeline = DDPMPipeline.from_pretrained(MODEL_ID)
    pipeline.unet.eval()
    model_status = f"✅ Model loaded — {sum(p.numel() for p in pipeline.unet.parameters()):,} params"
except Exception as e:
    # Fallback: create from config
    try:
        unet = UNet2DModel(
            sample_size=32, in_channels=3, out_channels=3, layers_per_block=1,
            block_out_channels=(32, 64, 64, 128),
            down_block_types=("DownBlock2D",)*4,
            up_block_types=("UpBlock2D",)*4,
            time_embedding_type="positional",
        )
        scheduler = DDPMScheduler(num_train_timesteps=1000)
        pipeline = DDPMPipeline(unet=unet, scheduler=scheduler)
        pipeline.unet.eval()
        model_status = f"⚠️ Using fresh UNet ({sum(p.numel() for p in unet.parameters()):,} params)"
    except:
        model_status = f"❌ Failed: {str(e)[:200]}"

# ─── Generate Images ─────────────────────────────────────────────────────────
def generate_images(batch_size=4, num_steps=50, seed=42):
    if pipeline is None:
        return [Image.new('RGB', (256, 256), color='gray') for _ in range(batch_size)]
    
    generator = torch.manual_seed(seed) if seed > 0 else None
    
    with torch.no_grad():
        results = pipeline(
            batch_size=batch_size,
            num_inference_steps=num_steps,
            generator=generator,
        ).images
    
    # Upscale from 32x32 to 256x256 for display
    upscaled = []
    for img in results:
        upscaled.append(img.resize((256, 256), Image.NEAREST))
    return upscaled

# ─── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="🎨 Sentinel Tiny Image", css="""
    .gradio-container { max-width: 1000px; margin: 0 auto; }
    .title { text-align: center; font-size: 2em; font-weight: bold; color: #6b4c9a; }
    .subtitle { text-align: center; color: #888; margin-bottom: 1em; }
""") as demo:
    gr.Markdown("""
    <div class="title">🎨 Sentinel Tiny Image</div>
    <div class="subtitle">2.9M parameter diffusion model with Sentinel super-exponential noise schedule</div>
    """)
    
    gr.Markdown(f"**Status**: {model_status}")
    
    with gr.Row():
        with gr.Column(scale=1):
            batch_size = gr.Slider(1, 8, value=4, step=1, label="Batch Size")
            num_steps = gr.Slider(10, 100, value=50, step=10, label="Inference Steps")
            seed = gr.Number(value=42, label="Random Seed (0=random)", precision=0)
            generate_btn = gr.Button("🎨 Generate", variant="primary")
        
        with gr.Column(scale=3):
            gallery = gr.Gallery(label="Generated Images", columns=4, height=300)
    
    with gr.Row():
        gr.Markdown("""
        ### About
        - **Architecture**: UNet2D, block channels [32,64,64,128], 1 layer/block
        - **Noise Schedule**: Sentinel super-exponential β(t)
        - **Dataset**: CIFAR-10 (2K samples demo)
        - **Output**: 32×32 images (upscaled 8× for display)
        - **Parameters**: 2.9M | **INT8**: [3 MB](https://huggingface.co/5dimension/sentinel-tiny-image-int8) | **INT4**: [1.4 MB](https://huggingface.co/5dimension/sentinel-tiny-image-int4)
        """)
    
    generate_btn.click(generate_images, [batch_size, num_steps, seed], gallery)

demo.launch()