import gradio as gr
import torch
from diffusers import (
    StableDiffusionXLPipeline,
    EulerDiscreteScheduler,
    AutoencoderKL,
    DPMSolverSinglestepScheduler,
)

# --- Configuration ---
# The base model your LoRA was trained on.
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"

# --- The file is local, so we just need its name ---
# The safetensors file is in the same directory as this script.
lora_file_path = "emilyh.safetensors"

# --- Load the Pipeline ---
# Use a recommended VAE for SDXL
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = StableDiffusionXLPipeline.from_pretrained(
    base_model_id,
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)

# --- Load the local LoRA file ---
# No download needed. We just load the local file directly.
pipe.load_lora_weights(lora_file_path)

# Move the pipeline to the GPU
pipe.to("cuda")

# --- Default Settings from your Recommendations ---
default_positive_prompt = "masterpiece, best quality, ultra-detailed, realistic skin, intricate details, highres"
default_negative_prompt = "low quality, worst quality, blurry, (deformed:1.3), extra fingers, cartoon, 3d, anime, bad anatomy"
default_sampler = "DPM++ 2M Karras"
default_cfg = 6.0
default_steps = 30
trigger_word = "emilyh"
lora_tag_main = "<lora:emilyh:0.9>"


# --- Define the Inference Function ---
def generate_image(prompt, negative_prompt, sampler, steps, cfg, width, height, seed):
    """
    Function to generate an image based on user inputs.
    """
    # Combine the user prompt with the trigger word and LoRA tag
    full_prompt = f"{lora_tag_main}, {trigger_word}, {prompt}"

    # Set the scheduler (sampler)
    if sampler == "DPM++ 2M Karras":
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
    elif sampler == "DPM++ SDE Karras":
        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
    else: # Default to DPM++ 2M Karras
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)

    # Set seed for reproducibility
    generator = torch.Generator("cuda").manual_seed(seed) if seed != -1 else None

    # Generate the image
    image = pipe(
        prompt=full_prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=cfg,
        num_inference_steps=steps,
        generator=generator,
        cross_attention_kwargs={"scale": 0.9} 
    ).images[0]

    return image

# --- Create the Gradio Interface ---
with gr.Blocks(css="style.css") as demo:
    gr.Markdown("# `emilyh` LoRA Image Generator")
    gr.Markdown(
        "A Gradio interface for the `emilyh` LoRA. "
        "Based on the recommendations provided."
    )

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Positive Prompt", value=default_positive_prompt, lines=3)
            negative_prompt = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
            
            with gr.Row():
                sampler = gr.Radio(
                    label="Sampler",
                    choices=["DPM++ 2M Karras", "DPM++ SDE Karras"],
                    value=default_sampler,
                )
                steps = gr.Slider(label="Steps", minimum=15, maximum=50, value=default_steps, step=1)
            
            cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=10.0, value=default_cfg, step=0.5)

            with gr.Row():
                width = gr.Slider(label="Width", minimum=512, maximum=1024, value=1024, step=64)
                height = gr.Slider(label="Height", minimum=512, maximum=1024, value=1024, step=64)

            seed = gr.Slider(label="Seed", minimum=-1, maximum=999999999, step=1, value=-1, info="Use -1 for a random seed.")
            
            generate_button = gr.Button("Generate Image", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil")
            gr.Markdown(
                """
                ### 🔧 Usage Guide
                * The trigger word `emilyh` and the LoRA tag `<lora:emilyh:0.9>` are automatically added to your prompt.
                * For best results, generate images in batches and choose the most consistent ones.
                * The LoRA captures the subject's appearance well across various poses and outfits.
                * A weight of 0.9 provides a good balance of likeness and flexibility. Using a weight closer to 1.0 can increase consistency but may cause stiffness.
                * This interface does not include ADetailer, which is recommended for final face refinement.
                """
            )

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, sampler, steps, cfg, width, height, seed],
        outputs=output_image
    )

demo.launch()