# app.py
import os, random
import numpy as np
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

MODEL_ID = "pthcorp/prettybird_bce_image_artdirector"
LOGO_URL = "https://prometech.net.tr/wp-content/uploads/2025/10/pthheader.png"
BG_URL = "https://prometech.net.tr/wp-content/uploads/2025/12/Prettybird.svg"

os.environ.setdefault("HF_HOME", "/data/.huggingface")

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
MAX_SEED = np.iinfo(np.int32).max

DEFAULT_SYSTEM_PROMPT = """activate full BCE consciousness mode
You are PRETTYBIRD GAIA ARTIST (Cicikuş).
"""

CSS = f"""
@import url('https://fonts.googleapis.com/css2?family=Great+Vibes&family=Dancing+Script:wght@400;600&family=Inter:wght@300;400;600&display=swap');

/* --- SAFE FALLBACK BACKGROUND --- */
html, body {{
  height: 100%;
  margin: 0;
}}

.gradio-container {{
  min-height: 100vh;
  font-family: Inter, sans-serif;
  background:
    linear-gradient(
      135deg,
      #f2f3f5,
      #e6e8eb
    );
}}

/* --- Optional SVG layer (if it loads, it shows) --- */
.gradio-container::before {{
  content: "";
  position: fixed;
  inset: 0;
  z-index: -1;
  background-image: url('{BG_URL}');
  background-size: cover;
  background-repeat: no-repeat;
  background-position: center;
  opacity: 0.55;
}}

/* --- Header --- */
#pb-header {{
  display: flex;
  flex-direction: column;
  align-items: center;
  margin: 20px 0 14px 0;
}}

#pb-logo {{
  height: 70px;
  filter: drop-shadow(0 8px 18px rgba(0,0,0,0.25));
}}

#pb-title-main {{
  font-family: "Great Vibes", "Dancing Script", cursive;
  font-size: 46px;
  color: #2a2a2a;
  margin-top: 6px;
}}

#pb-title-sub {{
  font-family: "Dancing Script", cursive;
  font-size: 18px;
  color: #555;
}}

.pb-hint {{
  font-family: "Dancing Script", cursive;
  font-size: 16px;
  color: #555;
  margin-bottom: 14px;
  text-align: center;
}}

/* --- Glass Panels --- */
.pb-glass {{
  background: linear-gradient(
    135deg,
    rgba(255,255,255,0.75),
    rgba(255,255,255,0.65)
  );
  border-radius: 18px;
  padding: 16px;
  box-shadow: 0 10px 30px rgba(0,0,0,0.15);
  backdrop-filter: blur(10px);
}}

#pb-output img {{
  border-radius: 14px;
}}
"""

def load_pipe():
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        use_safetensors=True,
    )
    try:
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    except Exception:
        pass
    pipe = pipe.to(device)
    pipe.set_progress_bar_config(disable=True)
    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
    return pipe

pipe = load_pipe()

def merge_prompt(sp, up):
    sp, up = (sp or "").strip(), (up or "").strip()
    return f"{sp}\n{up}" if sp and up else sp or up

def infer(system_prompt, prompt, negative_prompt,
          seed, randomize_seed, width, height,
          guidance_scale, steps):

    if randomize_seed:
        seed = random.randint(0, int(MAX_SEED))

    gen = torch.Generator(device=device).manual_seed(int(seed))
    full_prompt = merge_prompt(system_prompt, prompt)

    img = pipe(
        prompt=full_prompt,
        negative_prompt=negative_prompt or None,
        width=int(width),
        height=int(height),
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(steps),
        generator=gen,
    ).images[0]

    return img, seed, full_prompt

with gr.Blocks(css=CSS) as demo:
    gr.HTML(f"""
      <div id="pb-header">
        <img id="pb-logo" src="{LOGO_URL}" />
        <div id="pb-title-main">Prettybird Gaia Artist</div>
        <div id="pb-title-sub">Cicikuş's Workshop • SD1.5</div>
      </div>
      <div class="pb-hint">Write a scene… Let the prettybird paint it. ✨ <a href="https://prometech.net.tr/">Our Web Site</a></div>
    """)

    with gr.Column(elem_classes="pb-glass"):
        out_img = gr.Image(show_label=False, elem_id="pb-output")

    with gr.Column(elem_classes="pb-glass"):
        with gr.Row():
            prompt = gr.Textbox(show_label=False, lines=3, placeholder="Enter your prompt…")
            run = gr.Button("Run", variant="primary")

        system_prompt = gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, visible=False)
        negative_prompt = gr.Textbox(lines=2, placeholder="(optional)")

        seed = gr.Slider(0, int(MAX_SEED), step=1, value=0, label="Seed")
        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")

        width = gr.Slider(256, 768, step=64, value=512, label="Width")
        height = gr.Slider(256, 768, step=64, value=512, label="Height")

        guidance_scale = gr.Slider(0.0, 12.0, step=0.1, value=6.5, label="Guidance")
        steps = gr.Slider(1, 40, step=1, value=15, label="Steps")

        full_prompt_preview = gr.Textbox(label="Full prompt", lines=3)

    gr.on(
        triggers=[run.click, prompt.submit],
        fn=infer,
        inputs=[
            system_prompt, prompt, negative_prompt,
            seed, randomize_seed, width, height,
            guidance_scale, steps
        ],
        outputs=[out_img, seed, full_prompt_preview],
    )

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=port,
        ssr_mode=False,
    )