import gradio as gr
import numpy as np
import random
import spaces  # [uncomment to use ZeroGPU inside HuggingFace Spaces]
import torch
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, AutoencoderKL

# --- НАСТРОЙКИ ---
ENABLE_REFINER = True
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

base_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
refiner_repo_id = "stabilityai/stable-diffusion-xl-refiner-1.0"
vae_repo_id = "madebyollin/sdxl-vae-fp16-fix"

print(f"Device: {device}, dtype: {torch_dtype}")

# 1. Загружаем VAE
vae = AutoencoderKL.from_pretrained(vae_repo_id, torch_dtype=torch_dtype)

# 2. Загружаем Базовую модель (Text-to-Image)
pipe = StableDiffusionXLPipeline.from_pretrained(
    base_repo_id,
    vae=vae,
    torch_dtype=torch_dtype,
    use_safetensors=True,
    variant="fp16"
).to(device)

# 👉 Подключаем LoRA здесь
#pipe.load_lora_weights("FaceNpenisV4XL.safetensors", adapter_name="my_lora")
#pipe.set_adapters(["my_lora"], adapter_weights=[1])





# 3. Загружаем Refiner как Image-to-Image
refiner_pipe = None
if ENABLE_REFINER:
    print("Loading Refiner...")
    refiner_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        refiner_repo_id,
        vae=vae,
        torch_dtype=torch_dtype,
        use_safetensors=True,
        variant="fp16"
    ).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024



@spaces.GPU
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator(device).manual_seed(seed)
    
    width = max(256, min(MAX_IMAGE_SIZE, width // 64 * 64))
    height = max(256, min(MAX_IMAGE_SIZE, height // 64 * 64))
    
    original_size = (height, width)
    target_size = (height, width)
    crop_coords_top_left = (0, 0)

    if ENABLE_REFINER and refiner_pipe is not None:
        denoising_end = 0.8
        base_out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
            denoising_end=denoising_end,
            output_type="latent",
            original_size=original_size,
            target_size=target_size,
            crop_coords_top_left=crop_coords_top_left,
        )
        latents = base_out.images

        image = refiner_pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
            denoising_start=denoising_end,
            image=latents,
            original_size=original_size,
            target_size=target_size,
            crop_coords_top_left=crop_coords_top_left,
        ).images[0]
        
    else:
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
            output_type="pil",
            original_size=original_size,
            target_size=target_size,
            crop_coords_top_left=crop_coords_top_left,
        ).images[0]

    return image, seed

# --- ГРАФИЧЕСКИЙ ИНТЕРФЕЙС ---
examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css = """
#big-prompt textarea {
    font-size: 20px;   /* крупный шрифт */
    height: 300px;     /* высота поля */
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(" # SDXL 1.0 High Quality (Corrected)")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                elem_id="big-prompt",   # связываем с CSS
                lines=10,               # больше строк
                placeholder="Enter your prompt",
            )
            run_button = gr.Button("Run", variant="primary")

            with gr.Accordion("Advanced Settings", open=True):
                negative_prompt = gr.Textbox(
                    label="Negative prompt",
                    lines=2,
                    value="blurry, low quality, bad anatomy, ugly, distortion",
                )
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    width = gr.Slider(label="Width", minimum=512, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
                    height = gr.Slider(label="Height", minimum=512, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
                with gr.Row():
                    guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
                    num_inference_steps = gr.Slider(label="Steps", minimum=10, maximum=100, step=1, value=40)

            gr.Examples(examples=examples, inputs=[prompt])

        with gr.Column(scale=2):
            result = gr.Image(label="Result", show_label=False)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()