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

# ===============================
# TEXT MODEL (PROMPT ENHANCER)
# ===============================
TEXT_MODEL_ID = "HuggingFaceTB/SmolLM-135M-Instruct"

tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_ID)

def enhance_prompt(user_prompt: str) -> str:
    instruction = (
        "Please enhance this prompt so it is suitable for an image generator "
        "that requires clear instructions. Analyse the prompt, and output as "
        "much visual detail as possible about it.\n\n"
        f"Prompt to enhance: {user_prompt}\n\n"
        "Enhanced prompt:"
    )

    inputs = tokenizer(instruction, return_tensors="pt")
    outputs = text_model.generate(
        **inputs,
        max_new_tokens=500,
        temperature=0.6,
        do_sample=True,
    )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "Enhanced prompt:" in decoded:
        decoded = decoded.split("Enhanced prompt:")[-1]

    return decoded.strip()

# ===============================
# IMAGE MODEL (CPU)
# ===============================
IMG_MODEL = "runwayml/stable-diffusion-v1-5"
LCM_LORA = "latent-consistency/lcm-lora-sdv1-5"

pipe = DiffusionPipeline.from_pretrained(
    IMG_MODEL,
    torch_dtype=torch.float32,
    safety_checker=None
)

pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(LCM_LORA)
pipe.to("cpu")

pipe.enable_attention_slicing()
pipe.enable_vae_slicing()
pipe.set_progress_bar_config(disable=True)

# ===============================
# TIME ESTIMATION
# ===============================
def estimate_time(steps, res):
    res = int(res)
    steps = int(steps)
    base = {512: 12, 768: 25, 1024: 45}[res]
    total = steps * base + 5
    return f"⏱️ Estimated: ~{total//60}m {total%60}s"

# ===============================
# GENERATION FUNCTION
# ===============================
def generate(prompt, negative, resolution, steps):
    size = int(resolution)

    # Placeholder while thinking
    yield (
        [Image.new("RGB", (size, size), "white")],
        "🧠 Enhancing prompt..."
    )

    enhanced = enhance_prompt(prompt)

    yield (
        [Image.new("RGB", (size, size), "white")],
        "🎨 Generating image..."
    )

    seed = random.randint(0, 1_000_000_000)
    generator = torch.Generator("cpu").manual_seed(seed)

    pipe.scheduler.set_timesteps(int(steps))

    start = time.time()
    image = pipe(
        prompt=enhanced,
        negative_prompt=negative,
        num_inference_steps=int(steps),
        guidance_scale=1.2,
        width=size,
        height=size,
        generator=generator
    ).images[0]
    elapsed = int(time.time() - start)

    # Blur reveal
    for i in range(10):
        blur = image.filter(ImageFilter.GaussianBlur(radius=(10 - i)))
        yield (
            [blur],
            f"👀 Revealing... ({i+1}/10)"
        )
        time.sleep(1)

    yield (
        [image],
        f"✅ Done in {elapsed}s | Seed {seed}"
    )

# ===============================
# UI
# ===============================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 👾 Creeper AI Image - CPU")
    gr.Markdown(
        "1️⃣ The higher the resolution & steps, the longer the image takes.\n\n"
        "2️⃣ The more detailed the prompt and negative prompt, the better the result."
    )

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            negative = 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")
            generate_btn = gr.Button("Generate")

            status = gr.Markdown("🟢 Ready")

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

    resolution.change(estimate_time, [steps, resolution], eta)
    steps.change(estimate_time, [steps, resolution], eta)

    generate_btn.click(
        generate,
        inputs=[prompt, negative, resolution, steps],
        outputs=[gallery, status]
    )

demo.launch(share=False)