"""
CPU-only LoRA image generator that asks the user for an HF token
Dataset: lambdalabs/naruto-blip-captions (gated)
"""

import os
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from datasets import load_dataset
from huggingface_hub import login

# ---------- global state ----------
pipe = None        # will hold the pipeline once the user logs in
DATASET_NAME = "lambdalabs/naruto-blip-captions"
MODEL_ID = "runwayml/stable-diffusion-v1-5"

# ---------- helpers ----------
def load_pipeline(hf_token: str):
    """Build pipeline after HF login."""
    global pipe
    if pipe is None:
        login(token=hf_token, add_to_git_credential=False)
        pipe = StableDiffusionPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float32,  # CPU safe
            safety_checker=None,
            requires_safety_checker=False,
        )
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe.load_lora_weights(DATASET_NAME, weight_name="pytorch_lora_weights.safetensors")
    return pipe

def generate(prompt, steps, guidance, hf_token):
    """Generate image (only if token supplied)."""
    if not hf_token.strip():
        return gr.Warning("Please enter your Hugging Face token first.")
    try:
        pl = load_pipeline(hf_token)
        img = pl(prompt, num_inference_steps=int(steps), guidance_scale=float(guidance)).images[0]
        return img
    except Exception as e:
        return gr.Warning(str(e))

# ---------- Gradio UI ----------
with gr.Blocks(title="LoRA Naruto Generator – token login") as demo:
    gr.Markdown("## 🍥 LoRA Naruto Generator (gated dataset)")
    gr.Markdown("Enter your **HF token** once, then generate images.")

    with gr.Row():
        token_box = gr.Textbox(
            label="Hugging Face Token",
            placeholder="hf_...",
            type="password",
            lines=1,
        )

    with gr.Row():
        prompt_box = gr.Textbox(
            label="Prompt",
            placeholder="Naruto riding a giant frog, anime style",
            lines=2,
        )

    with gr.Row():
        steps_slider = gr.Slider(10, 50, value=20, label="Inference steps")
        guidance_slider = gr.Slider(1.0, 20.0, value=7.5, label="Guidance scale")

    generate_btn = gr.Button("🚀 Generate")
    output_img = gr.Image(type="pil", label="Result")

    examples = gr.Examples(
        examples=[
            ["naruto wearing sunglasses, anime style", 20, 7.5],
            ["sakura in a pink dress, cherry blossoms", 25, 8.0],
            ["sasuke with lightning background, dark mood", 15, 7.0],
        ],
        inputs=[prompt_box, steps_slider, guidance_slider],
    )

    generate_btn.click(
        generate,
        inputs=[prompt_box, steps_slider, guidance_slider, token_box],
        outputs=output_img,
    )

if __name__ == "__main__":
    demo.queue(max_size=5).launch()