import gradio as gr
import torch
import random
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
from diffusers import DiffusionPipeline, LCMScheduler
from PIL import Image, ImageFilter

# -------------------------------
# SMALL PROMPT ENHANCER (CPU)
# -------------------------------
ENHANCER_MODEL = "HuggingFaceTB/SmolLM-135M-Instruct"
tokenizer_enhancer = AutoTokenizer.from_pretrained(ENHANCER_MODEL)
model_enhancer = AutoModelForCausalLM.from_pretrained(ENHANCER_MODEL)

def enhance_text(user_prompt, prefix):
    """
    Uses a tiny LLM to rewrite user input into a more detailed instruction.
    """
    instruction = f"Rewrite this for an image generator with detail: {user_prompt}"
    if prefix:
        instruction = f"{prefix} {user_prompt}"

    inputs = tokenizer_enhancer(instruction, return_tensors="pt")
    outputs = model_enhancer.generate(
        **inputs,
        max_new_tokens=50,
        temperature=0.7,
        do_sample=True
    )
    text = tokenizer_enhancer.decode(outputs[0], skip_special_tokens=True)
    return text.strip()

# -------------------------------
# IMAGE MODEL SETUP (CPU SAFE)
# -------------------------------
IMG_MODEL_ID = "runwayml/stable-diffusion-v1-5"
IMG_ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"

pipe = DiffusionPipeline.from_pretrained(
    IMG_MODEL_ID,
    torch_dtype=torch.float32,
    safety_checker=None
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(IMG_ADAPTER_ID)
pipe.to("cpu")
pipe.enable_attention_slicing()
pipe.enable_vae_slicing()
pipe.set_progress_bar_config(disable=True)

# -------------------------------
# ETA
# -------------------------------
def estimate_time(steps, resolution):
    steps = int(steps)
    resolution = int(resolution)
    per_step = {512:12, 768:25, 1024:45}[resolution]
    overhead = 2
    est = overhead + steps * per_step
    mins = est // 60
    secs = est % 60
    return f"⏱️ Estimated: ~{mins}m {secs}s"

# -------------------------------
# GENERATE WITH PROGRESSIVE BLUR
# -------------------------------
def generate(prompt, neg_prompt, resolution, steps):
    # 1️⃣ AI PROMPT ENHANCEMENT
    enhanced_prompt = enhance_text(prompt, "")
    enhanced_negative = enhance_text(neg_prompt, "Rewrite negative prompt:")

    # 2️⃣ White placeholder
    placeholder = Image.new("RGB", (int(resolution), int(resolution)), (255,255,255))
    yield [placeholder], "🟡 Generating..."

    # 3️⃣ CPU IMAGE GENERATION
    seed = random.randint(0, 10**9)
    gen = torch.Generator("cpu").manual_seed(seed)
    pipe.scheduler.set_timesteps(int(steps))

    img = pipe(
        prompt=enhanced_prompt,
        negative_prompt=enhanced_negative,
        num_inference_steps=int(steps),
        guidance_scale=1.2,
        width=int(resolution),
        height=int(resolution),
        generator=gen
    ).images[0]

    # 4️⃣ BLUR REVEAL
    max_blur = 20
    for i in range(10):
        blur_pct = 100 - i*10
        blurred = img.filter(ImageFilter.GaussianBlur(radius=max_blur * blur_pct/100))
        yield [blurred], "🟢 Revealing..."
        time.sleep(1)

    # 5️⃣ FINAL
    yield [img], f"✅ Done | Seed: {seed}"

# -------------------------------
# GRADIO UI
# -------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 👾 CREEPER AI — SMART IMAGE GENERATOR")
    gr.Markdown("1: The higher the resolution & steps, the longer the image takes to make.\n2: The more detailed the prompt and negative prompt, the better the result.")

    with gr.Row():
        with gr.Column():
            prompt_in = gr.Textbox(label="Prompt")
            neg_in = gr.Textbox(label="Negative Prompt")

            resolution = gr.Radio([512, 768, 1024], value=512, label="Resolution")
            steps = gr.Slider(6,8,value=6,step=1,label="Steps")

            eta = gr.Markdown("⏱️ Estimated: ~1m 0s")
            gen_btn = gr.Button("Generate")

            status = gr.Markdown("🟢 Ready")

        with gr.Column():
            gallery = gr.Gallery(columns=1)

    for ctrl in [steps, resolution]:
        ctrl.change(estimate_time, [steps, resolution], eta)

    gen_btn.click(
        generate,
        inputs=[prompt_in, neg_in, resolution, steps],
        outputs=[gallery, status]
    )

demo.launch()