app.py

import io, json
from typing import List, Dict, Optional, Tuple

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import torch
from diffusers import (
    StableDiffusionXLPipeline,
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionXLInpaintPipeline,
    StableDiffusionXLControlNetPipeline,
    ControlNetModel,
    StableDiffusionUpscalePipeline,
    DPMSolverMultistepScheduler, EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler, HeunDiscreteScheduler,
)

# ---------------- Optional deps (safe imports: ไม่มีก็ข้าม) ----------------
try:
    from rembg import remove as rembg_remove
except Exception:
    rembg_remove = None

_HAS_GFP = False
GFPGANer = None
GFP = None
try:
    import gfpgan  # type: ignore
    if hasattr(gfpgan, "GFPGANer"):
        GFPGANer = gfpgan.GFPGANer  # type: ignore
        _HAS_GFP = True
except Exception as e:
    print("[WARN] GFPGAN not available:", e)

_HAS_REALESRGAN = False
RealESRGAN = None
REALSR = None
try:
    from realesrgan import RealESRGAN  # type: ignore
    _HAS_REALESRGAN = True
except Exception as e:
    print("[WARN] RealESRGAN not available:", e)

# ---------------- Runtime setup ----------------
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype  = torch.float16 if device == "cuda" else torch.float32

# ---------------- Registries ----------------
MODELS: List[Tuple[str,str,str]] = [
    ("stabilityai/stable-diffusion-xl-base-1.0",   "SDXL Base 1.0", "เอนกประสงค์"),
    ("stabilityai/stable-diffusion-xl-refiner-1.0","SDXL Refiner",  "เสริมรายละเอียด (pass 2)"),
    ("SG161222/RealVisXL_V4.0",                    "RealVis XL v4", "โฟโต้เรียล คน/สินค้า"),
    ("Lykon/dreamshaper-xl-v2",                    "DreamShaper XL","แฟนตาซี-เรียลลิสติก"),
    ("RunDiffusion/Juggernaut-XL",                 "Juggernaut XL", "คอนทราสต์แรง"),
    ("emilianJR/epiCRealismXL",                    "EpicRealism XL","แฟชั่น/พอร์เทรต"),
    ("black-forest-labs/FLUX.1-dev",               "FLUX.1-dev",    "แนวสมัยใหม่ (ไม่ใช่ SDXL)"),
    ("stabilityai/sd-turbo",                       "SD-Turbo",      "เร็วมากสำหรับร่างไอเดีย"),
    ("stabilityai/stable-diffusion-2-1",           "SD 2.1",        "แลนด์สเคปกว้าง"),
    ("runwayml/stable-diffusion-v1-5",             "SD 1.5",        "คลาสสิก"),
    ("timbrooks/instruct-pix2pix",                 "Instruct-Pix2Pix","แก้ภาพตามคำสั่ง"),
]

LORAS: List[Tuple[str,str,str]] = [
    ("ByteDance/SDXL-Lightning",           "SDXL-Lightning",   "สปีด"),
    ("ostris/epicrealism-xl-lora",         "EpicRealism XL",   "โทนจริง"),
    ("alpha-diffusion/sdxl-anime-lora",    "Anime-Style XL",   "อนิเม"),
    ("alpha-diffusion/sdxl-cinematic-lora","Cinematic-Drama",  "แสงหนัง"),
    ("alpha-diffusion/sdxl-watercolor-lora","Watercolor",      "สีน้ำ"),
    ("alpha-diffusion/sdxl-fashion-lora",  "Fashion",          "แฟชั่น"),
    ("alpha-diffusion/sdxl-product-lora",  "Product-Studio",   "สินค้า"),
    ("alpha-diffusion/sdxl-interior-lora", "Interior-Archi",   "สถาปัตย์"),
    ("alpha-diffusion/sdxl-food-lora",     "Food-Tasty",       "อาหาร"),
    ("alpha-diffusion/sdxl-logo-lora",     "Logo-Clean",       "โลโก้"),
]

# ใช้ 5 ชนิดหลักเพื่อ UI กระชับและเสถียร
CONTROLNETS: List[Tuple[str,str,str,str]] = [
    ("diffusers/controlnet-canny-sdxl-1.0",    "Canny",     "เส้นขอบ",    "canny"),
    ("diffusers/controlnet-openpose-sdxl-1.0", "OpenPose",  "ท่าทางคน",   "pose"),
    ("diffusers/controlnet-depth-sdxl-1.0",    "Depth",     "ระยะลึก",    "depth"),
    ("diffusers/controlnet-softedge-sdxl-1.0", "SoftEdge",  "เส้นนุ่ม",    "softedge"),
    ("diffusers/controlnet-lineart-sdxl-1.0",  "Lineart",   "เส้นร่าง",    "lineart"),
]

PRESETS = {
    "Cinematic": ", cinematic lighting, 50mm, bokeh, film grain, high dynamic range",
    "Studio":    ", studio photo, softbox lighting, sharp focus, high detail",
    "Anime":     ", anime style, clean lineart, vibrant colors, high quality",
    "Product":   ", product photography, seamless background, diffused light, reflections",
}
NEG_DEFAULT = "lowres, blurry, bad anatomy, extra fingers, watermark, jpeg artifacts, text"

SCHEDULERS = {
    "DPM-Solver (Karras)": DPMSolverMultistepScheduler,
    "Euler": EulerDiscreteScheduler,
    "Euler a": EulerAncestralDiscreteScheduler,
    "Heun": HeunDiscreteScheduler,
}

# ---------------- Caches ----------------
PIPE_CACHE: Dict[str, object] = {}
CONTROL_CACHE: Dict[str, ControlNetModel] = {}
UPSCALE_PIPE: Optional[StableDiffusionUpscalePipeline] = None

# ---------------- Helpers ----------------
def set_sched(pipe, name: str):
    cls = SCHEDULERS.get(name, DPMSolverMultistepScheduler)
    pipe.scheduler = cls.from_config(pipe.scheduler.config)

def seed_gen(sd: int):
    if sd is None or sd < 0: return None
    g = torch.Generator(device=("cuda" if device=="cuda" else "cpu"))
    g.manual_seed(int(sd)); return g

def prep_pipe(model_id: str, control_ids: List[str]):
    key = f"{model_id}|{'-'.join(control_ids) if control_ids else 'none'}"
    if key in PIPE_CACHE: return PIPE_CACHE[key]

    if control_ids:
        cn_models = []
        for cid in control_ids:
            if cid not in CONTROL_CACHE:
                CONTROL_CACHE[cid] = ControlNetModel.from_pretrained(cid, torch_dtype=dtype, use_safetensors=True)
            cn_models.append(CONTROL_CACHE[cid])
        pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model_id, controlnet=cn_models, torch_dtype=dtype, use_safetensors=True)
    else:
        pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=dtype, use_safetensors=True)

    pipe.to(device)
    try:
        if device == "cuda":
            pipe.enable_vae_tiling(); pipe.enable_vae_slicing()
            pipe.enable_xformers_memory_efficient_attention()
        else:
            pipe.enable_attention_slicing()
    except Exception:
        pass

    PIPE_CACHE[key] = pipe
    return pipe

def apply_loras(pipe, lora_ids: List[str]):
    for rid in [x for x in lora_ids if x]:
        try:
            pipe.load_lora_weights(rid)
        except Exception as e:
            print(f"[LoRA] load failed {rid}: {e}")

def to_info(meta: dict) -> str:
    return json.dumps(meta, ensure_ascii=False, indent=2)

# ---------------- Post-process ----------------
def ensure_upscalers():
    global UPSCALE_PIPE, GFP, REALSR
    if UPSCALE_PIPE is None:
        try:
            UPSCALE_PIPE = StableDiffusionUpscalePipeline.from_pretrained(
                "stabilityai/stable-diffusion-x4-upscaler",
                torch_dtype=dtype, use_safetensors=True
            ).to(device)
        except Exception as e:
            print("[Upscaler] SD x4 not available:", e)

    if _HAS_GFP and GFP is None and GFPGANer is not None:
        try:
            GFP = GFPGANer(model_path=None, upscale=1, arch="clean", channel_multiplier=2)
        except Exception as e:
            print("[GFPGAN] init failed:", e)

    if _HAS_REALESRGAN and REALSR is None and device == "cuda":
        try:
            REALSR = RealESRGAN(torch.device("cuda"), scale=4)  # ต้องมี weights เองจึงจะทำงานจริง
        except Exception as e:
            REALSR = None
            print("[RealESRGAN] init failed:", e)

def post_process(img: Image.Image, do_up: bool, do_face: bool, do_bg: bool):
    ensure_upscalers()
    out = img

    # Upscale: RealESRGAN (ถ้ามี) > SD x4 > skip
    if do_up:
        try:
            if REALSR is not None:
                out = Image.fromarray(REALSR.predict(np.array(out)))
            elif UPSCALE_PIPE is not None:
                if device == "cuda":
                    with torch.autocast("cuda"):
                        out = UPSCALE_PIPE(prompt="", image=out).images[0]
                else:
                    out = UPSCALE_PIPE(prompt="", image=out).images[0]
        except Exception as e:
            print("[Upscale] skipped:", e)

    if do_face and _HAS_GFP and GFP is not None:
        try:
            _, _, restored = GFP.enhance(np.array(out), has_aligned=False, only_center_face=False, paste_back=True)
            out = Image.fromarray(restored)
        except Exception as e:
            print("[GFPGAN] skipped:", e)

    if do_bg and rembg_remove is not None:
        try:
            out = Image.open(io.BytesIO(rembg_remove(np.array(out))))
        except Exception as e:
            print("[rembg] skipped:", e)

    return out

# ---------------- Generators ----------------
def run_txt2img(
    model_id, model_custom, prompt, preset, negative,
    steps, cfg, width, height, scheduler_name, seed,
    lora_selected, lora_custom,
    ctrl_selected, img_canny, img_pose, img_depth, img_softedge, img_lineart,
    do_up, do_face, do_bg
):
    if not prompt or not str(prompt).strip():
        raise gr.Error("กรุณากรอก prompt")

    model = (model_custom.strip() or model_id).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not negative.strip(): negative = NEG_DEFAULT

    # ControlNet mapping (เฉพาะภาพที่อัปโหลดจริง)
    label_to_img = {
        "Canny": img_canny, "OpenPose": img_pose, "Depth": img_depth,
        "SoftEdge": img_softedge, "Lineart": img_lineart
    }
    control_ids, cond_images = [], []
    for cid, label, note, key in CONTROLNETS:
        if label in ctrl_selected and label_to_img.get(label) is not None:
            control_ids.append(cid); cond_images.append(label_to_img[label])

    pipe = prep_pipe(model, control_ids)
    set_sched(pipe, scheduler_name)

    # LoRA
    lora_ids = [s.split(" — ")[0].strip() for s in (lora_selected or [])]
    if lora_custom and lora_custom.strip():
        lora_ids += [x.strip() for x in lora_custom.split(",") if x.strip()]
    apply_loras(pipe, lora_ids)

    width, height = int(width), int(height)
    gen = seed_gen(seed)

    if device == "cuda":
        with torch.autocast("cuda"):
            if control_ids:
                img = pipe(
                    prompt=prompt, negative_prompt=negative,
                    width=width, height=height,
                    num_inference_steps=int(steps), guidance_scale=float(cfg),
                    controlnet_conditioning_image=cond_images if len(cond_images)>1 else cond_images[0],
                    generator=gen
                ).images[0]
            else:
                img = pipe(
                    prompt=prompt, negative_prompt=negative,
                    width=width, height=height,
                    num_inference_steps=int(steps), guidance_scale=float(cfg),
                    generator=gen
                ).images[0]
    else:
        if control_ids:
            img = pipe(
                prompt=prompt, negative_prompt=negative,
                width=width, height=height,
                num_inference_steps=int(steps), guidance_scale=float(cfg),
                controlnet_conditioning_image=cond_images if len(cond_images)>1 else cond_images[0],
                generator=gen
            ).images[0]
        else:
            img = pipe(
                prompt=prompt, negative_prompt=negative,
                width=width, height=height,
                num_inference_steps=int(steps), guidance_scale=float(cfg),
                generator=gen
            ).images[0]

    img = post_process(img, do_up, do_face, do_bg)
    meta = {
        "mode":"txt2img","model":model,"loras":lora_ids,"controlnets":ctrl_selected,
        "prompt":prompt,"negative":negative,"size":f"{width}x{height}",
        "steps":steps,"cfg":cfg,"scheduler":scheduler_name,"seed":seed,
        "post":{"upscale":do_up,"face_restore":do_face,"remove_bg":do_bg}
    }
    return img, to_info(meta)

def run_img2img(
    model_id, model_custom, init_image, strength,
    prompt, preset, negative, steps, cfg, width, height, scheduler_name, seed,
    do_up, do_face, do_bg
):
    if init_image is None: raise gr.Error("โปรดอัปโหลดภาพเริ่มต้น")
    model = (model_custom.strip() or model_id).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not negative.strip(): negative = NEG_DEFAULT

    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(model, torch_dtype=dtype, use_safetensors=True).to(device)
    try:
        if device=="cuda": pipe.enable_xformers_memory_efficient_attention()
    except Exception: pass
    set_sched(pipe, scheduler_name); gen = seed_gen(seed)

    if device=="cuda":
        with torch.autocast("cuda"):
            img = pipe(prompt=prompt, negative_prompt=negative,
                       image=init_image, strength=float(strength),
                       num_inference_steps=int(steps), guidance_scale=float(cfg),
                       generator=gen).images[0]
    else:
        img = pipe(prompt=prompt, negative_prompt=negative,
                   image=init_image, strength=float(strength),
                   num_inference_steps=int(steps), guidance_scale=float(cfg),
                   generator=gen).images[0]

    img = post_process(img, do_up, do_face, do_bg)
    meta = {"mode":"img2img","model":model,"prompt":prompt,"neg":negative,
            "steps":steps,"cfg":cfg,"seed":seed,"strength":strength}
    return img, to_info(meta)

def expand_canvas_for_outpaint(img: Image.Image, expand_px: int, direction: str) -> Tuple[Image.Image, Image.Image]:
    w, h = img.size
    if direction == "left":
        new = Image.new("RGBA",(w+expand_px,h),(0,0,0,0)); new.paste(img,(expand_px,0))
        mask = Image.new("L",(w+expand_px,h),0); d=ImageDraw.Draw(mask); d.rectangle([0,0,expand_px,h], fill=255)
    elif direction == "right":
        new = Image.new("RGBA",(w+expand_px,h),(0,0,0,0)); new.paste(img,(0,0))
        mask = Image.new("L",(w+expand_px,h),0); d=ImageDraw.Draw(mask); d.rectangle([w,0,w+expand_px,h], fill=255)
    elif direction == "top":
        new = Image.new("RGBA",(w,h+expand_px),(0,0,0,0)); new.paste(img,(0,expand_px))
        mask = Image.new("L",(w,h+expand_px),0); d=ImageDraw.Draw(mask); d.rectangle([0,0,w,expand_px], fill=255)
    else:
        new = Image.new("RGBA",(w,h+expand_px),(0,0,0,0)); new.paste(img,(0,0))
        mask = Image.new("L",(w,h+expand_px),0); d=ImageDraw.Draw(mask); d.rectangle([0,h,w,h+expand_px], fill=255)
    return new.convert("RGB"), mask

def run_inpaint_outpaint(
    model_id, model_custom, base_image, mask_image, mode, expand_px, expand_dir,
    prompt, preset, negative, steps, cfg, width, height, scheduler_name, seed,
    strength, do_up, do_face, do_bg
):
    if base_image is None: raise gr.Error("โปรดอัปโหลดภาพฐาน")
    model = (model_custom.strip() or model_id).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not negative.strip(): negative = NEG_DEFAULT

    pipe = StableDiffusionXLInpaintPipeline.from_pretrained(model, torch_dtype=dtype, use_safetensors=True).to(device)
    try:
        if device=="cuda": pipe.enable_xformers_memory_efficient_attention()
    except Exception: pass
    set_sched(pipe, scheduler_name); gen = seed_gen(seed)

    if mode == "Outpaint":
        base_image, mask_image = expand_canvas_for_outpaint(base_image, int(expand_px), expand_dir)

    if device=="cuda":
        with torch.autocast("cuda"):
            img = pipe(prompt=prompt, negative_prompt=negative,
                       image=base_image, mask_image=mask_image,
                       strength=float(strength),
                       num_inference_steps=int(steps), guidance_scale=float(cfg),
                       generator=gen).images[0]
    else:
        img = pipe(prompt=prompt, negative_prompt=negative,
                   image=base_image, mask_image=mask_image,
                   strength=float(strength),
                   num_inference_steps=int(steps), guidance_scale=float(cfg),
                   generator=gen).images[0]

    img = post_process(img, do_up, do_face, do_bg)
    meta = {"mode":mode,"model":model,"prompt":prompt,"steps":steps,"cfg":cfg,"seed":seed}
    return img, to_info(meta)

# ---------------- UI ----------------
def build_ui():
    with gr.Blocks(theme=gr.themes.Soft(), title="Masterpiece SDXL Studio Pro") as demo:
        gr.Markdown("# 🖼️ Masterpiece SDXL Studio Pro")
        gr.Markdown("Text2Img • Img2Img • Inpaint/Outpaint • Multi-LoRA • ControlNet • Upscale/FaceRestore/RemoveBG (optional)")

        # Common controls
        model_dd = gr.Dropdown(choices=[m[0] for m in MODELS], value=MODELS[0][0], label="Model")
        model_custom = gr.Textbox(label="Custom Model ID", placeholder="(ถ้าอยากใช้โมเดลของคุณเอง กรอกที่นี่)")

        preset = gr.Dropdown(choices=list(PRESETS.keys()), value=None, label="Style Preset (optional)")
        negative = gr.Textbox(value=NEG_DEFAULT, label="Negative Prompt")

        steps = gr.Slider(10, 60, 30, step=1, label="Steps")
        cfg = gr.Slider(1.0, 12.0, 7.0, step=0.1, label="CFG")
        width = gr.Slider(512, 1024, 832, step=64, label="Width")
        height= gr.Slider(512, 1024, 832, step=64, label="Height")
        scheduler = gr.Dropdown(list(SCHEDULERS.keys()), value="DPM-Solver (Karras)", label="Scheduler")
        seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")

        # LoRA & ControlNet
        lora_sel = gr.CheckboxGroup(choices=[f"{rid} — {lbl} ({note})" for rid,lbl,note in LORAS], label="LoRA (เลือกได้หลายตัว)")
        lora_custom = gr.Textbox(label="Custom LoRA IDs (comma separated)")

        ctrl_sel = gr.CheckboxGroup(choices=[c[1] for c in CONTROLNETS], label="ControlNet ชนิดที่ใช้")
        img_canny = gr.Image(type="pil", label="Canny")
        img_pose = gr.Image(type="pil", label="OpenPose")
        img_depth = gr.Image(type="pil", label="Depth")
        img_softedge = gr.Image(type="pil", label="SoftEdge")
        img_lineart = gr.Image(type="pil", label="Lineart")

        with gr.Row():
            do_up = gr.Checkbox(False, label="Upscale x4 (ถ้ามี)")
            do_face = gr.Checkbox(False, label="Face Restore (ถ้ามี)")
            do_bg = gr.Checkbox(False, label="Remove BG (ถ้ามี)")

        with gr.Tab("Text → Image"):
            prompt_txt = gr.Textbox(lines=3, label="Prompt")
            btn_txt = gr.Button("🚀 Generate")
            out_img_txt = gr.Image(type="pil", label="Result")
            out_meta_txt = gr.Textbox(label="Metadata", lines=10)

        with gr.Tab("Image → Image"):
            init_img = gr.Image(type="pil", label="Init Image")
            strength = gr.Slider(0.1, 1.0, 0.7, 0.05, label="Strength")
            prompt_i2i = gr.Textbox(lines=3, label="Prompt")
            btn_i2i = gr.Button("🚀 Img2Img")
            out_img_i2i = gr.Image(type="pil", label="Result")
            out_meta_i2i = gr.Textbox(label="Metadata", lines=10)

        with gr.Tab("Inpaint / Outpaint"):
            base_img = gr.Image(type="pil", label="Base Image")
            mask_img = gr.Image(type="pil", label="Mask (ขาว=แก้, ดำ=คงเดิม)")
            mode_io = gr.Radio(["Inpaint","Outpaint"], value="Inpaint", label="Mode")
            expand_px = gr.Slider(64, 1024, 256, 64, label="Outpaint pixels")
            expand_dir = gr.Radio(["left","right","top","bottom"], value="right", label="Outpaint direction")
            prompt_io = gr.Textbox(lines=3, label="Prompt")
            btn_io = gr.Button("🚀 Inpaint/Outpaint")
            out_img_io = gr.Image(type="pil", label="Result")
            out_meta_io = gr.Textbox(label="Metadata", lines=10)

        # Bindings
        btn_txt.click(
            fn=run_txt2img,
            inputs=[
                model_dd, model_custom, prompt_txt, preset, negative,
                steps, cfg, width, height, scheduler, seed,
                lora_sel, lora_custom,
                ctrl_sel, img_canny, img_pose, img_depth, img_softedge, img_lineart,
                do_up, do_face, do_bg
            ],
            outputs=[out_img_txt, out_meta_txt],
            api_name="txt2img"
        )

        btn_i2i.click(
            fn=run_img2img,
            inputs=[
                model_dd, model_custom, init_img, strength,
                prompt_i2i, preset, negative, steps, cfg, width, height, scheduler, seed,
                do_up, do_face, do_bg
            ],
            outputs=[out_img_i2i, out_meta_i2i],
            api_name="img2img"
        )

        btn_io.click(
            fn=run_inpaint_outpaint,
            inputs=[
                model_dd, model_custom, base_img, mask_img, mode_io, expand_px, expand_dir,
                prompt_io, preset, negative, steps, cfg, width, height, scheduler, seed,
                strength, do_up, do_face, do_bg
            ],
            outputs=[out_img_io, out_meta_io],
            api_name="inpaint_outpaint"
        )

        gr.Markdown("ℹ️ ถ้าโมดูลเสริมหรือบางโมเดลไม่พร้อมใช้งาน ระบบจะข้ามอย่างปลอดภัยและแจ้งเตือนใน Console")

    return demo

demo = build_ui()
demo.queue(max_size=8).launch()