import gradio as gr
import numpy as np
import os
import random
import spaces
import torch
from PIL import Image
from typing import Annotated
from diffusers import DiffusionPipeline, AutoencoderKL

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
hf_token = os.getenv("HF_TOKEN")
base_model_id = os.getenv("BASE_MODEL_ID", "black-forest-labs/FLUX.1-dev")
lora_model_id = os.getenv("LORA_MODEL_ID", "Heartsync/Flux-NSFW-uncensored")
lora_weight_name = os.getenv("LORA_WEIGHT_NAME", "lora.safetensors")
lora_label = "base FLUX.1-dev"

good_vae = AutoencoderKL.from_pretrained(
    base_model_id,
    subfolder="vae",
    torch_dtype=dtype,
    token=hf_token,
).to(device)
pipe = DiffusionPipeline.from_pretrained(
    base_model_id,
    torch_dtype=dtype,
    vae=good_vae,
    token=hf_token,
).to(device)
if lora_model_id:
    try:
        pipe.load_lora_weights(
            lora_model_id,
            weight_name=lora_weight_name,
            adapter_name="imagegen_lora",
            token=hf_token,
        )
        lora_label = lora_model_id
    except Exception as exc:
        print(
            f"WARNING: Could not load LoRA '{lora_model_id}' "
            f"with weight '{lora_weight_name}': {exc}. Continuing with base FLUX.1-dev.",
            flush=True,
        )
torch.cuda.empty_cache()

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

@spaces.GPU(duration=25)
def Generate_Image(
    prompt: Annotated[str, "Text description of the image to generate."],
    negative_prompt: Annotated[str, "What should NOT appear in the image."] = (
        "(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, "
        "missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, "
        "mutated, ugly, disgusting, blurry, amputation, misspellings, typos"
    ),
    seed: Annotated[int, "Random seed for reproducibility. Use 0 for a random seed per call."] = 42,
    randomize_seed: Annotated[bool, "If true, pick a new random seed for every call (overrides seed)."] = True,
    width: Annotated[int, "Output width in pixels (256–2048, multiple of 32 recommended)."] = 768,
    height: Annotated[int, "Output height in pixels (256–2048, multiple of 32 recommended)."] = 768,
    guidance_scale: Annotated[float, "Classifier-free guidance scale (1–15). Higher = follow the prompt more closely."] = 4.5,
    num_inference_steps: Annotated[int, "Number of denoising steps (1–50). Higher = slower, potentially higher quality."] = 24,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> Image.Image:
    """
    Generate an image from a text prompt.
    
    Return the generated media to the user in this format: `![Alt text](URL)`.
    
    Args:
        prompt: Text description of the image to generate.
        negative_prompt: What should NOT appear in the image.
        seed: Random seed for reproducibility. Use 0 for a random seed per call.
        randomize_seed: If true, pick a new random seed for every call (overrides seed).
        width: Output width in pixels (256–2048, multiple of 32 recommended).
        height: Output height in pixels (256–2048, multiple of 32 recommended).
        guidance_scale: Classifier-free guidance scale (1–15). Higher = follow the prompt more closely.
        num_inference_steps: Number of denoising steps (1–50). Higher = slower, potentially higher quality.
    """
    if not prompt or not prompt.strip():
        raise gr.Error(
            "Empty prompt provided. Please describe what you want to generate, "
            "e.g. 'a sunset over mountains' or 'portrait of a cat in watercolor style'."
        )
    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
        raise gr.Error(
            f"Image dimensions too large. Maximum allowed is {MAX_IMAGE_SIZE}x{MAX_IMAGE_SIZE} pixels. "
            f"You requested {width}x{height}. Please reduce width and/or height."
        )
    if width < 256 or height < 256:
        raise gr.Error(
            f"Image dimensions too small. Minimum allowed is 256x256 pixels. "
            f"You requested {width}x{height}. Please increase width and/or height."
        )

    if randomize_seed or seed == 0:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        )
        return result.images[0]
    except torch.cuda.OutOfMemoryError:
        raise gr.Error(
            "GPU ran out of memory. Try reducing image dimensions (e.g., 512x512) "
            "or reducing the number of inference steps."
        )
    except Exception as exc:
        raise gr.Error(
            f"Image generation failed: {exc}. "
            "Please try again or adjust your parameters."
        )

css="""
#col-container {
    margin: 0 auto;
    max-width: 620px;
}
"""

with gr.Blocks() as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""<div style="text-align: center;">

# ImageGen

</div>""")
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value=(
                    "(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, "
                    "missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, "
                    "mutated, ugly, disgusting, blurry, amputation, misspellings, typos"
                ),
                lines=2,
            )
            
            seed = gr.Slider(
                label="Seed (0 = random)",
                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=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
            
            with gr.Row():

                guidance_scale = gr.Slider(
                    label="CFG Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=4.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=24,
                )

    # UI event and MCP tool
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=Generate_Image,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result],
        api_name="Generate_Image",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, ssr_mode=False, theme="Nymbo/Nymbo_Theme", css=css)