app.py

# ================= ZeroGPU-Optimized =================
# - Cache ไปที่ /tmp (ล้างทุก Restart)
# - Lazy load + LRU (เก็บ pipeline ล่าสุดแค่ 1-2 ตัว)
# - ใช้โมเดลเบาเป็นค่าเริ่มต้น (SD 1.5 / SD-Turbo)
# - ControlNet เฉพาะ Canny (เล็กและเร็ว)
# - ปุ่ม Clear cache ใน UI
# - Auto-retry ลดขนาด/steps เมื่อ OOM หรือค้างนาน
# =====================================================
import os, io, json, time, gc, shutil
from typing import Dict, List, Optional, Tuple
from collections import OrderedDict

# 1) ส่ง cache ไป /tmp เพื่อลดการสะสมพื้นที่
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["HF_HUB_CACHE"] = "/tmp/hf/hub"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf/transformers"
os.environ["DIFFUSERS_CACHE"] = "/tmp/hf/diffusers"

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import torch
from diffusers import (
    StableDiffusionPipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipelineLegacy,
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    DPMSolverMultistepScheduler, EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler, HeunDiscreteScheduler,
)

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

# ---------- Runtime ----------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32

# ---------- Model Registry (เบา/เร็ว เหมาะ ZeroGPU) ----------
# *SDXL ถูกตัดออกจากค่าเริ่มต้น (โหลดหนัก) แต่ยังรองรับถ้ากรอก Custom ID เอง
MODELS_TXT = [
    ("runwayml/stable-diffusion-v1-5", "SD 1.5 (base, fast)"),
    ("stabilityai/sd-turbo",          "SD-Turbo (ultra-fast)"),
    ("stabilityai/stable-diffusion-2-1", "SD 2.1 (landscape)"),
]
MODEL_IMG2IMG_DEFAULT = "runwayml/stable-diffusion-v1-5"
MODEL_INPAINT_DEFAULT = "runwayml/stable-diffusion-inpainting"  # legacy inpaint (เล็ก/เสถียร)

# ControlNet: เอาเฉพาะ Canny (เล็กและพอเพียง)
CONTROLNETS = [
    ("lllyasviel/sd-controlnet-canny", "Canny (edges)"),
]

PRESETS = {
    "Cinematic": ", cinematic lighting, bokeh, film grain",
    "Studio": ", studio photo, softbox lighting, sharp focus",
    "Anime": ", anime style, clean lines, vibrant colors",
}
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 with LRU ----------
MAX_PIPE_CACHE = 2
PIPE_CACHE: "OrderedDict[str, object]" = OrderedDict()
CONTROL_CACHE: Dict[str, ControlNetModel] = {}

def _lru_put(key, pipe):
    PIPE_CACHE[key] = pipe
    PIPE_CACHE.move_to_end(key)
    while len(PIPE_CACHE) > MAX_PIPE_CACHE:
        old_key, old_pipe = PIPE_CACHE.popitem(last=False)
        try:
            del old_pipe
        except Exception:
            pass
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

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

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

def _speed_tweaks(pipe):
    # ลดหน่วยความจำ/เพิ่มเสถียร
    try:
        if DEVICE == "cuda":
            pipe.enable_xformers_memory_efficient_attention()
            pipe.enable_vae_tiling()
            pipe.enable_vae_slicing()
        else:
            pipe.enable_sequential_cpu_offload()
            pipe.enable_attention_slicing()
    except Exception:
        pass

# ---------- Lazy loaders ----------
def get_controlnet(model_id: str):
    if model_id in CONTROL_CACHE:
        return CONTROL_CACHE[model_id]
    cn = ControlNetModel.from_pretrained(model_id, torch_dtype=DTYPE, use_safetensors=True)
    cn.to(DEVICE)
    CONTROL_CACHE[model_id] = cn
    return cn

def get_txt2img_pipe(model_id: str, use_control: bool, control_id: Optional[str]):
    key = f"t2i|{model_id}|{'cn' if use_control else 'none'}"
    if key in PIPE_CACHE:
        PIPE_CACHE.move_to_end(key)
        return PIPE_CACHE[key]

    if use_control and control_id:
        cn = get_controlnet(control_id)
        pipe = StableDiffusionControlNetPipeline.from_pretrained(
            model_id, controlnet=cn, torch_dtype=DTYPE,
            safety_checker=None, feature_extractor=None, use_safetensors=True
        )
    else:
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id, torch_dtype=DTYPE,
            safety_checker=None, feature_extractor=None, use_safetensors=True
        )
    pipe.to(DEVICE)
    _speed_tweaks(pipe)
    _lru_put(key, pipe)
    return pipe

def get_img2img_pipe(model_id: str):
    key = f"i2i|{model_id}"
    if key in PIPE_CACHE:
        PIPE_CACHE.move_to_end(key)
        return PIPE_CACHE[key]
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model_id, torch_dtype=DTYPE,
        safety_checker=None, feature_extractor=None, use_safetensors=True
    ).to(DEVICE)
    _speed_tweaks(pipe)
    _lru_put(key, pipe)
    return pipe

def get_inpaint_pipe(model_id: str):
    key = f"inpaint|{model_id}"
    if key in PIPE_CACHE:
        PIPE_CACHE.move_to_end(key)
        return PIPE_CACHE[key]
    pipe = StableDiffusionInpaintPipelineLegacy.from_pretrained(
        model_id, torch_dtype=DTYPE,
        safety_checker=None, feature_extractor=None, use_safetensors=True
    ).to(DEVICE)
    _speed_tweaks(pipe)
    _lru_put(key, pipe)
    return pipe

# ---------- Post process ----------
def remove_bg(img: Image.Image) -> Image.Image:
    if rembg_remove is None: return img
    try:
        return Image.open(io.BytesIO(rembg_remove(np.array(img))))
    except Exception:
        return img

# ---------- Auto-retry wrapper ----------
def run_with_retry(func, *, width: int, height: int, steps: int, max_time: float = 280.0):
    """ลองรันด้วยพารามิเตอร์เดิม → ถ้า OOM/Timeout จะลดขนาดภาพ/จำนวนสเต็ปแล้วรันซ้ำ"""
    t0 = time.time()
    w, h, s = width, height, steps
    for attempt in range(3):
        try:
            if time.time() - t0 > max_time:
                raise gr.Error("งานนานเกินกำหนด โปรดลองลดขนาดภาพหรือจำนวนสเต็ป")
            return func(w, h, s)
        except RuntimeError as e:
            msg = str(e).lower()
            if "out of memory" in msg or "cuda oom" in msg or "alloc" in msg:
                # ลดขนาดครึ่งหนึ่ง และลดสเต็ปเล็กน้อย
                w = max(384, int(w * 0.75) // 64 * 64)
                h = max(384, int(h * 0.75) // 64 * 64)
                s = max(10, s - 4)
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                continue
            raise
    raise gr.Error("หน่วยความจำไม่พอ แม้จะลดขนาดแล้ว — ลองลดพารามิเตอร์เพิ่มเติม")

# ---------- Generators ----------
def txt2img(
    model_id, custom_model, prompt, preset, negative,
    steps, cfg, width, height, scheduler, seed,
    use_control, control_choice, control_image,
    do_rembg
):
    if not prompt or not str(prompt).strip():
        raise gr.Error("กรุณากรอก prompt")

    model = (custom_model.strip() or model_id or MODELS_TXT[0][0]).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not str(negative).strip(): negative = NEG_DEFAULT
    width, height = int(width), int(height)
    use_control = bool(use_control and control_choice and control_image is not None)

    def _run(w, h, s):
        pipe = get_txt2img_pipe(model, use_control, CONTROLNETS[0][0] if use_control else None)
        set_scheduler(pipe, scheduler)
        gen = seed_gen(seed)
        if use_control:
            image = pipe(
                prompt=prompt, negative_prompt=negative,
                image=control_image, width=w, height=h,
                num_inference_steps=int(s), guidance_scale=float(cfg),
                generator=gen
            ).images[0]
        else:
            image = pipe(
                prompt=prompt, negative_prompt=negative,
                width=w, height=h,
                num_inference_steps=int(s), guidance_scale=float(cfg),
                generator=gen
            ).images[0]
        if do_rembg: image = remove_bg(image)
        meta = {
            "mode":"txt2img","model":model,"control":("canny" if use_control else None),
            "prompt":prompt,"neg":negative,"size":f"{w}x{h}",
            "steps":int(s),"cfg":float(cfg),"scheduler":scheduler,"seed":seed
        }
        return image, json.dumps(meta, ensure_ascii=False, indent=2)

    return run_with_retry(_run, width=width, height=height, steps=int(steps))

def img2img(
    model_id, custom_model, init_img, strength,
    prompt, preset, negative, steps, cfg, width, height, scheduler, seed,
    do_rembg
):
    if init_img is None: raise gr.Error("โปรดอัปโหลดภาพเริ่มต้น")
    model = (custom_model.strip() or model_id or MODEL_IMG2IMG_DEFAULT).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not str(negative).strip(): negative = NEG_DEFAULT
    width, height = int(width), int(height)

    def _run(w, h, s):
        pipe = get_img2img_pipe(model)
        set_scheduler(pipe, scheduler)
        gen = seed_gen(seed)
        image = pipe(
            prompt=prompt, negative_prompt=negative, image=init_img, strength=float(strength),
            num_inference_steps=int(s), guidance_scale=float(cfg),
            generator=gen
        ).images[0]
        if do_rembg: image = remove_bg(image)
        meta = {"mode":"img2img","model":model,"prompt":prompt,"neg":negative,
                "steps":int(s),"cfg":float(cfg),"seed":seed,"strength":float(strength)}
        return image, json.dumps(meta, ensure_ascii=False, indent=2)

    return run_with_retry(_run, width=width, height=height, steps=int(steps))

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); ImageDraw.Draw(mask).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); ImageDraw.Draw(mask).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); ImageDraw.Draw(mask).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); ImageDraw.Draw(mask).rectangle([0,h,w,h+expand_px], fill=255)
    return new.convert("RGB"), mask

def inpaint_outpaint(
    model_id, custom_model, base_img, mask_img, mode, expand_px, expand_dir,
    prompt, preset, negative, steps, cfg, width, height, scheduler, seed,
    strength, do_rembg
):
    if base_img is None: raise gr.Error("โปรดอัปโหลดภาพฐาน")
    model = (custom_model.strip() or model_id or MODEL_INPAINT_DEFAULT).strip()
    if preset and preset in PRESETS: prompt = prompt + PRESETS[preset]
    if not negative or not str(negative).strip(): negative = NEG_DEFAULT
    width, height = int(width), int(height)

    if mode == "Outpaint":
        base_img, mask_img = expand_canvas_for_outpaint(base_img, int(expand_px), expand_dir)

    def _run(w, h, s):
        pipe = get_inpaint_pipe(model)
        set_scheduler(pipe, scheduler)
        gen = seed_gen(seed)
        image = pipe(
            prompt=prompt, negative_prompt=negative,
            image=base_img, mask_image=mask_img, strength=float(strength),
            num_inference_steps=int(s), guidance_scale=float(cfg),
            generator=gen
        ).images[0]
        if do_rembg: image = remove_bg(image)
        meta = {"mode":mode,"model":model,"prompt":prompt,"steps":int(s),"cfg":float(cfg),"seed":seed}
        return image, json.dumps(meta, ensure_ascii=False, indent=2)

    return run_with_retry(_run, width=width, height=height, steps=int(steps))

# ---------- Clear cache ----------
def clear_runtime_caches():
    cache_root = os.environ.get("HF_HOME", "/tmp/hf")
    try:
        if os.path.isdir(cache_root):
            shutil.rmtree(cache_root, ignore_errors=True)
    except Exception as e:
        print("[ClearCache] remove cache failed:", e)
    PIPE_CACHE.clear()
    CONTROL_CACHE.clear()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return "✅ Cache cleared. Pipelines will be reloaded on demand."

# ---------- UI ----------
def build_ui():
    with gr.Blocks(theme=gr.themes.Soft(), title="ZeroGPU SD Studio") as demo:
        gr.Markdown("## 🖼️ ZeroGPU SD Studio — เบา เร็ว เสถียร (CPU/ZeroGPU)")
        with gr.Row():
            model_dd = gr.Dropdown([m[0] for m in MODELS_TXT], value=MODELS_TXT[0][0], label="Base model")
            model_custom = gr.Textbox(label="Custom model ID (optional)", placeholder="เช่น stabilityai/stable-diffusion-xl-base-1.0 (หนัก)")

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

        with gr.Row():
            steps = gr.Slider(10, 40, 18, 1, label="Steps (แนะนำ ≤20 บน ZeroGPU)")
            cfg   = gr.Slider(1.0, 12.0, 6.5, 0.1, label="CFG")
        with gr.Row():
            width  = gr.Slider(384, 768, 512, 64, label="Width")
            height = gr.Slider(384, 768, 512, 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)")

        # ControlNet (Canny เท่านั้น)
        with gr.Accordion("ControlNet (Canny)", open=False):
            use_control = gr.Checkbox(False, label="Enable Canny ControlNet")
            control_choice = gr.Dropdown([CONTROLNETS[0][1]], value=CONTROLNETS[0][1], label="Type")
            control_image  = gr.Image(type="pil", label="Edge image")

        with gr.Row():
            do_rembg = gr.Checkbox(False, label="Remove background (ถ้ามี rembg)")

        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 (white = edit)")
            mode_io = gr.Radio(["Inpaint","Outpaint"], value="Inpaint", label="Mode")
            expand_px = gr.Slider(64, 512, 192, 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)

        with gr.Row():
            btn_clear = gr.Button("🧹 Clear cache (runtime)")
            msg_clear = gr.Markdown()

        # Bindings
        btn_txt.click(
            fn=txt2img,
            inputs=[model_dd, model_custom, prompt_txt, preset, negative,
                    steps, cfg, width, height, scheduler, seed,
                    use_control, control_choice, control_image,
                    do_rembg],
            outputs=[out_img_txt, out_meta_txt],
            api_name="txt2img"
        )

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

        btn_io.click(
            fn=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_rembg],
            outputs=[out_img_io, out_meta_io],
            api_name="inpaint_outpaint"
        )

        btn_clear.click(fn=clear_runtime_caches, outputs=[msg_clear])

        gr.Markdown("ℹ️ โหมดนี้ออกแบบมาสำหรับ ZeroGPU/CPU: ถ้าต้องการ SDXL ให้กรอก Custom ID (จะช้าหนักขึ้น)")

    return demo

demo = build_ui()
# ลดโอกาส connection หลุดใน ZeroGPU
demo.queue(concurrency_count=1, max_size=8)
demo.launch(share=False, show_api=False, max_threads=1, prevent_thread_lock=True)