app.py

import gradio as gr
import numpy as np
from PIL import Image
import torch
import gc
import os
import warnings

# Suppress specific warnings
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', message='.*torch_dtype.*deprecated.*')
warnings.filterwarnings('ignore', message='.*CLIPFeatureExtractor.*deprecated.*')

# Performance optimizations
if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

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

print(f"🖥️  Device: {device} | dtype: {dtype}")

# Lazy import
from diffusers import (
    StableDiffusionControlNetPipeline, 
    ControlNetModel, 
    StableDiffusionPipeline,
    StableDiffusionXLPipeline,
    StableDiffusionXLControlNetPipeline,
    AutoPipelineForText2Image
)
from diffusers import UniPCMultistepScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler
from controlnet_aux import (
    LineartDetector, 
    LineartAnimeDetector,
    OpenposeDetector,
    MidasDetector,
    CannyDetector,
    MLSDdetector,
    HEDdetector,
    PidiNetDetector,
    NormalBaeDetector,
    ZoeDetector,
    MediapipeFaceDetector
)

# Memory optimization
if torch.cuda.is_available():
    torch.cuda.empty_cache()
    torch.cuda.set_per_process_memory_fraction(0.95)
    print(f"🔥 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
else:
    print("⚠️  Running on CPU - Image generation will be significantly slower")

# ===== Model & Config =====
CURRENT_CONTROLNET_PIPE = None
CURRENT_CONTROLNET_KEY = None
CURRENT_T2I_PIPE = None
CURRENT_T2I_MODEL = None
CURRENT_SDXL_REFINER = None
CURRENT_TURBO_PIPE = None
CURRENT_TURBO_MODEL = None

# Enhanced SDXL Models (including NSFW-capable)
SDXL_MODELS = [
    "stabilityai/stable-diffusion-xl-base-1.0",
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    "stabilityai/sdxl-turbo",  # เพิ่ม SDXL-Turbo
    "Laxhar/noobai-XL-1.1",
    "RunDiffusion/Juggernaut-XL-v9",
    "dataautogpt3/ProteusV0.4",
    "playgroundai/playground-v2.5-1024px-aesthetic",
    "misri/epicrealismXL_v10",
    "SG161222/RealVisXL_V4.0",
    "stablediffusionapi/juggernaut-xl-v8",
    "Lykon/dreamshaper-xl-1-0",
    "digiplay/Pony_Diffusion_V6_XL"
]

# Turbo models (fast generation)
TURBO_MODELS = [
    "stabilityai/sdxl-turbo",
    "stabilityai/sd-turbo"
]

# Enhanced SD1.5 Models (including NSFW-capable)
SD15_MODELS = [
    # Original models
    "digiplay/ChikMix_V3",
    "digiplay/chilloutmix_NiPrunedFp16Fix",
    "gsdf/Counterfeit-V2.5",
    "stablediffusionapi/anything-v5",
    "runwayml/stable-diffusion-v1-5",
    "stablediffusionapi/realistic-vision-v51",
    "stablediffusionapi/dreamshaper-v8",
    "stablediffusionapi/henmix-real-v11",
    "stablediffusionapi/rev-animated-v122",
    "stablediffusionapi/cyberrealistic-v33",
    "stablediffusionapi/meinamix-meina-v11",
    "prompthero/openjourney-v4",
    "wavymulder/Analog-Diffusion",
    "dreamlike-art/dreamlike-photoreal-2.0",
    "segmind/SSD-1B",
    "SG161222/Realistic_Vision_V5.1_noVAE",
    "Lykon/dreamshaper-8",
    "hakurei/waifu-diffusion",
    "andite/anything-v4.0",
    "Linaqruf/animagine-xl",
    # Additional NSFW-capable models
    "emilianJR/epiCRealism",
    "stablediffusionapi/deliberate-v2",
    "stablediffusionapi/edge-of-realism",
    "Yntec/epiCPhotoGasm",
    "digiplay/majicMIX_realistic_v7",
    "stablediffusionapi/perfect-world-v6",
    "stablediffusionapi/uber-realistic-merge",
    "XpucT/Deliberate",
    "prompthero/openjourney",
    "Lykon/absolute-reality-1.81",
    "digiplay/BeautyProMix_v2",
    "stablediffusionapi/3d-animation-diffusion",
    "nitrosocke/Ghibli-Diffusion",
    "nitrosocke/mo-di-diffusion",
    "Fictiverse/Stable_Diffusion_VoxelArt_Model"
]

# Specialized models
SPECIAL_MODELS = {
    "waifu_colorize": "ShinoharaHare/Waifu-Colorize-XL"  # โมเดลสำหรับ colorize โดยเฉพาะ
}

# Chinese Models
CHINESE_MODELS = [
    "AI-Chen/Chinese-Stable-Diffusion",
    "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1",
    "AI-ModelScope/stable-diffusion-v1-5-chinese"
]

# ControlNet models for SD1.5
CONTROLNET_MODELS_SD15 = {
    "lineart": "lllyasviel/control_v11p_sd15_lineart",
    "lineart_anime": "lllyasviel/control_v11p_sd15s2_lineart_anime",
    "canny": "lllyasviel/control_v11p_sd15_canny",
    "depth": "lllyasviel/control_v11p_sd15_depth",
    "normal": "lllyasviel/control_v11p_sd15_normalbae",
    "openpose": "lllyasviel/control_v11p_sd15_openpose",
    "softedge": "lllyasviel/control_v11p_sd15_softedge",
    "segmentation": "lllyasviel/control_v11p_sd15_seg",
    "mlsd": "lllyasviel/control_v11p_sd15_mlsd",
    "shuffle": "lllyasviel/control_v11p_sd15_shuffle",
    "scribble": "lllyasviel/control_v11p_sd15_scribble",
    "tile": "lllyasviel/control_v11f1e_sd15_tile"
}

# ControlNet models for SDXL
CONTROLNET_MODELS_SDXL = {
    "canny_sdxl": "diffusers/controlnet-canny-sdxl-1.0",
    "depth_sdxl": "diffusers/controlnet-depth-sdxl-1.0",
    "openpose_sdxl": "thibaud/controlnet-openpose-sdxl-1.0",
    "lineart_sdxl": "ShermanG/ControlNet-Standard-Lineart-for-SDXL"  # เพิ่ม ControlNet Lineart สำหรับ SDXL
}

# แก้ไข LoRA models ให้ใช้ชื่อที่ไม่มีช่องว่าง
LORA_MODELS = {
    "None": None,
    # Style LoRAs
    "lowpoly-game-character": "nerijs/lowpoly-game-character-lora",
    "pixel-art": "nerijs/pixel-art-xl",
    "watercolor-style": "OedoSoldier/watercolor-style-lora",
    "manga-style": "raemikk/Animerge_V3.0_LoRA",
    "cyberpunk": "artificialguybr/cyberpunk-anime-diffusion",
    "fantasy-art": "artificialguybr/fantasy-art-lora",
    "chinese-style": "yfszzx/Chinese_style_xl_LoRA",
    "traditional-painting": "artificialguybr/Traditional-Painting-Style-LoRA",
    "anime-art": "Linaqruf/anime-detailer-xl-lora",
    "cinematic": "artificialguybr/cinematic-diffusion",
    "oil-painting": "artificialguybr/oil-painting-style",
    # Character/Face LoRAs
    "japanese-doll": "Norod78/sd15-JapaneseDollLikeness_lora",
    "korean-doll": "Norod78/sd15-KoreanDollLikeness_lora",
    "detail-tweaker": "nitrosocke/detail-tweaker-lora",
    "beautiful-realistic-asians": "etok/Beautiful_Realistic_Asians",
    "asian-beauty": "digiplay/AsianBeauty_V1",
    "perfect-hands": "Sanster/perfect-hands",
    "face-detail": "ostris/face-detail-lora",
    # Body/Pose LoRAs
    "body-pose-control": "alvdansen/lora-body-pose",
    "dynamic-poses": "alvdansen/dynamic-poses-lora",
    "full-body": "artificialguybr/full-body-lora",
    # Realism LoRAs
    "photorealistic": "microsoft/lora-photorealistic",
    "hyper-realistic": "dallinmackay/hyper-realistic-lora",
    "ultra-realistic": "artificialguybr/ultra-realistic-lora",
    "realistic-vision": "SG161222/Realistic_Vision_V5.1_noVAE",
    # Lighting/Quality LoRAs
    "add-detail": "ostris/add-detail-lora",
    "sharp-details": "ostris/sharp-details-lora",
    "better-lighting": "artificialguybr/better-lighting-lora",
    "studio-lighting": "artificialguybr/studio-lighting",
    # NSFW-capable LoRAs
    "nsfw-master": "hearmeneigh/nsfw-master-lora",
    "realistic-nsfw": "digiplay/RealisticNSFW_v1",
    "anime-nsfw": "Linaqruf/anime-nsfw-lora",
    "hentai-diffusion": "Deltaadams/Hentai-Diffusion",
    "sexy-pose": "alvdansen/sexy-pose-lora",
    # Colorize LoRAs
    "colorize-xl": "ShinoharaHare/Waifu-Colorize-XL"
}

# VAE models for better quality
VAE_MODELS = {
    "None": None,
    "SD1.5 VAE": "stabilityai/sd-vae-ft-mse",
    "Anime VAE": "hakurei/waifu-diffusion-v1-4",
    "SDXL VAE": "madebyollin/sdxl-vae-fp16-fix",
    "Turbo VAE": "madebyollin/sdxl-vae-fp16-fix"
}

# Detector instances
DETECTORS = {}

def is_sdxl_model(model_name: str) -> bool:
    """Check if model is SDXL"""
    return model_name in SDXL_MODELS or "xl" in model_name.lower() or "XL" in model_name

def is_turbo_model(model_name: str) -> bool:
    """Check if model is Turbo"""
    return model_name in TURBO_MODELS or "turbo" in model_name.lower()

def load_detector(detector_type: str):
    """Lazy load detector"""
    global DETECTORS
    
    if detector_type in DETECTORS:
        return DETECTORS[detector_type]
    
    print(f"📥 Loading {detector_type} detector...")
    
    try:
        if detector_type == "lineart":
            DETECTORS[detector_type] = LineartDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "lineart_anime":
            DETECTORS[detector_type] = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "openpose":
            DETECTORS[detector_type] = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "depth":
            DETECTORS[detector_type] = MidasDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "canny":
            DETECTORS[detector_type] = CannyDetector()
        elif detector_type == "normal":
            DETECTORS[detector_type] = NormalBaeDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "hed":
            DETECTORS[detector_type] = HEDdetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "pidi":
            DETECTORS[detector_type] = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "mlsd":
            DETECTORS[detector_type] = MLSDdetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "zoe":
            DETECTORS[detector_type] = ZoeDetector.from_pretrained("lllyasviel/Annotators")
        elif detector_type == "face":
            DETECTORS[detector_type] = MediapipeFaceDetector()
        else:
            raise ValueError(f"Unknown detector type: {detector_type}")
        
        return DETECTORS[detector_type]
    except Exception as e:
        print(f"❌ Error loading {detector_type} detector: {e}")
        return None

def get_controlnet_model(controlnet_type: str, is_sdxl: bool = False):
    """Get ControlNet model based on type"""
    if is_sdxl:
        if controlnet_type not in CONTROLNET_MODELS_SDXL:
            raise ValueError(f"SDXL ControlNet type must be one of: {list(CONTROLNET_MODELS_SDXL.keys())}")
        return CONTROLNET_MODELS_SDXL[controlnet_type]
    else:
        if controlnet_type not in CONTROLNET_MODELS_SD15:
            raise ValueError(f"SD1.5 ControlNet type must be one of: {list(CONTROLNET_MODELS_SD15.keys())}")
        return CONTROLNET_MODELS_SD15[controlnet_type]

def prepare_condition_image(image, controlnet_type, is_sdxl=False):
    """Prepare condition image for ControlNet"""
    if "lineart" in controlnet_type:
        detector_type = "lineart_anime" if "anime" in controlnet_type else "lineart"
        detector = load_detector(detector_type)
        if detector:
            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
            return Image.fromarray(result) if isinstance(result, np.ndarray) else result
    
    elif "canny" in controlnet_type:
        detector = load_detector("canny")
        if detector:
            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
            return Image.fromarray(result) if isinstance(result, np.ndarray) else result
    
    elif "depth" in controlnet_type:
        detector = load_detector("depth")
        if detector:
            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
            return Image.fromarray(result) if isinstance(result, np.ndarray) else result
    
    elif controlnet_type == "normal":
        detector = load_detector("normal")
        if detector:
            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
            return Image.fromarray(result) if isinstance(result, np.ndarray) else result
    
    elif "openpose" in controlnet_type:
        detector = load_detector("openpose")
        if detector:
            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
            return Image.fromarray(result) if isinstance(result, np.ndarray) else result
    
    return image

def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model: str = None, 
                lora_weight: float = 0.8, vae_model: str = None):
    """Get or create a ControlNet pipeline with optional LoRA and VAE"""
    global CURRENT_CONTROLNET_PIPE, CURRENT_CONTROLNET_KEY
    
    key = (model_name, controlnet_type, lora_model, lora_weight, vae_model)
    
    if CURRENT_CONTROLNET_KEY == key and CURRENT_CONTROLNET_PIPE is not None:
        print(f"✅ Reusing existing ControlNet pipeline: {model_name}, type: {controlnet_type}")
        return CURRENT_CONTROLNET_PIPE
    
    if CURRENT_CONTROLNET_PIPE is not None:
        print(f"🗑️ Unloading old ControlNet pipeline: {CURRENT_CONTROLNET_KEY}")
        del CURRENT_CONTROLNET_PIPE
        CURRENT_CONTROLNET_PIPE = None
        CURRENT_CONTROLNET_KEY = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    print(f"📥 Loading ControlNet pipeline for model: {model_name}, type: {controlnet_type}")
    
    try:
        is_sdxl = is_sdxl_model(model_name)
        is_turbo = is_turbo_model(model_name)
        
        controlnet_model_name = get_controlnet_model(controlnet_type, is_sdxl)
        controlnet = ControlNetModel.from_pretrained(
            controlnet_model_name, 
            torch_dtype=dtype
        ).to(device)
        
        if is_sdxl:
            if is_turbo:
                # สำหรับ Turbo models ใช้ pipeline ที่เหมาะสม
                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                    model_name, 
                    controlnet=controlnet, 
                    torch_dtype=dtype,
                    safety_checker=None,
                    requires_safety_checker=False,
                    use_safetensors=True,
                    variant="fp16" if dtype == torch.float16 else None
                ).to(device)
            else:
                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                    model_name, 
                    controlnet=controlnet, 
                    torch_dtype=dtype,
                    safety_checker=None,
                    requires_safety_checker=False,
                    use_safetensors=True,
                    variant="fp16" if dtype == torch.float16 else None
                ).to(device)
        else:
            pipe = StableDiffusionControlNetPipeline.from_pretrained(
                model_name, 
                controlnet=controlnet, 
                torch_dtype=dtype,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=True,
                variant="fp16" if dtype == torch.float16 else None
            ).to(device)
        
        # Load custom VAE if specified
        if vae_model and vae_model != "None":
            try:
                from diffusers import AutoencoderKL
                print(f"🔄 Loading custom VAE: {vae_model}")
                vae = AutoencoderKL.from_pretrained(vae_model, torch_dtype=dtype).to(device)
                pipe.vae = vae
            except Exception as e:
                print(f"⚠️  Error loading VAE: {e}")
        
        # Apply LoRA if specified
        if lora_model and lora_model != "None":
            print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
            try:
                if lora_model in LORA_MODELS:
                    lora_path = LORA_MODELS[lora_model]
                    pipe.load_lora_weights(lora_path)
                    pipe.fuse_lora(lora_scale=lora_weight)
                else:
                    pipe.load_lora_weights(lora_model)
                    pipe.fuse_lora(lora_scale=lora_weight)
            except Exception as e:
                print(f"⚠️  Error loading LoRA: {e}")
        
        # Optimizations
        pipe.enable_attention_slicing(slice_size="max")
        
        if hasattr(pipe, 'vae') and hasattr(pipe.vae, 'enable_slicing'):
            pipe.vae.enable_slicing()
        else:
            try:
                pipe.enable_vae_slicing()
            except:
                pass
        
        if device.type == "cuda":
            try:
                pipe.enable_xformers_memory_efficient_attention()
                print("✅ xFormers enabled")
            except:
                pass
            pipe.enable_model_cpu_offload()
        
        # ใช้ scheduler ที่เหมาะสมสำหรับแต่ละโมเดล
        if is_turbo:
            # สำหรับ Turbo models ใช้ scheduler ที่เร็วขึ้น
            try:
                pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
                print("✅ Using UniPC scheduler for Turbo")
            except:
                pass
        else:
            try:
                pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
                print("✅ Using Euler Ancestral scheduler")
            except:
                try:
                    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
                except:
                    pass
        
        CURRENT_CONTROLNET_PIPE = pipe
        CURRENT_CONTROLNET_KEY = key
        return pipe
        
    except Exception as e:
        print(f"❌ Error loading ControlNet pipeline: {e}")
        CURRENT_CONTROLNET_PIPE = None
        CURRENT_CONTROLNET_KEY = None
        raise

def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float = 0.8, 
                   vae_model: str = None):
    """Load text-to-image model with optional LoRA and VAE"""
    global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL, CURRENT_SDXL_REFINER, CURRENT_TURBO_PIPE, CURRENT_TURBO_MODEL
    
    is_turbo = is_turbo_model(model_name)
    use_refiner = "refiner" in model_name.lower() and not is_turbo
    
    key = (model_name, lora_model, lora_weight, vae_model, use_refiner, is_turbo)
    
    if is_turbo:
        if CURRENT_TURBO_MODEL == key and CURRENT_TURBO_PIPE is not None:
            return
    else:
        if CURRENT_T2I_MODEL == key and CURRENT_T2I_PIPE is not None:
            return
    
    if is_turbo:
        if CURRENT_TURBO_PIPE is not None:
            print(f"🗑️ Unloading old Turbo model: {CURRENT_TURBO_MODEL}")
            del CURRENT_TURBO_PIPE
            CURRENT_TURBO_PIPE = None
    else:
        if CURRENT_T2I_PIPE is not None:
            print(f"🗑️ Unloading old T2I model: {CURRENT_T2I_MODEL}")
            del CURRENT_T2I_PIPE
            CURRENT_T2I_PIPE = None
            if CURRENT_SDXL_REFINER is not None:
                del CURRENT_SDXL_REFINER
                CURRENT_SDXL_REFINER = None
    
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    print(f"📥 Loading T2I model: {model_name}")
    
    try:
        if is_turbo:
            # โหลด Turbo model
            CURRENT_TURBO_PIPE = AutoPipelineForText2Image.from_pretrained(
                model_name,
                torch_dtype=dtype,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=True,
                variant="fp16" if dtype == torch.float16 else None
            ).to(device)
            
            CURRENT_TURBO_MODEL = key
            pipe = CURRENT_TURBO_PIPE
        elif is_sdxl_model(model_name):
            if use_refiner:
                CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
                    "stabilityai/stable-diffusion-xl-base-1.0",
                    torch_dtype=dtype,
                    safety_checker=None,
                    requires_safety_checker=False,
                    use_safetensors=True,
                    variant="fp16" if dtype == torch.float16 else None
                ).to(device)
                
                CURRENT_SDXL_REFINER = StableDiffusionXLPipeline.from_pretrained(
                    model_name,
                    torch_dtype=dtype,
                    safety_checker=None,
                    requires_safety_checker=False,
                    use_safetensors=True,
                    variant="fp16" if dtype == torch.float16 else None,
                    text_encoder_2=CURRENT_T2I_PIPE.text_encoder_2,
                    vae=CURRENT_T2I_PIPE.vae
                ).to(device)
            else:
                CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
                    model_name,
                    torch_dtype=dtype,
                    safety_checker=None,
                    requires_safety_checker=False,
                    use_safetensors=True,
                    variant="fp16" if dtype == torch.float16 else None
                ).to(device)
            
            CURRENT_T2I_MODEL = key
            pipe = CURRENT_T2I_PIPE
        else:
            CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
                model_name,
                torch_dtype=dtype,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=True,
                variant="fp16" if dtype == torch.float16 else None
            ).to(device)
            
            CURRENT_T2I_MODEL = key
            pipe = CURRENT_T2I_PIPE
        
        # Load custom VAE
        if vae_model and vae_model != "None":
            try:
                from diffusers import AutoencoderKL
                print(f"🔄 Loading custom VAE: {vae_model}")
                vae = AutoencoderKL.from_pretrained(vae_model, torch_dtype=dtype).to(device)
                pipe.vae = vae
            except Exception as e:
                print(f"⚠️  Error loading VAE: {e}")
        
        # Apply LoRA
        if lora_model and lora_model != "None":
            print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
            try:
                if lora_model in LORA_MODELS:
                    lora_path = LORA_MODELS[lora_model]
                    pipe.load_lora_weights(lora_path)
                    pipe.fuse_lora(lora_scale=lora_weight)
                else:
                    pipe.load_lora_weights(lora_model)
                    pipe.fuse_lora(lora_scale=lora_weight)
            except Exception as e:
                print(f"⚠️  Error loading LoRA: {e}")
        
        # Optimizations
        pipe.enable_attention_slicing(slice_size="max")
        
        if hasattr(pipe, 'vae') and hasattr(pipe.vae, 'enable_slicing'):
            pipe.vae.enable_slicing()
        else:
            try:
                pipe.enable_vae_slicing()
            except:
                pass
        
        if device.type == "cuda":
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except:
                pass
            pipe.enable_model_cpu_offload()
        
        # ตั้งค่า scheduler
        if is_turbo:
            try:
                pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
                print("✅ Using UniPC scheduler for Turbo")
            except:
                pass
        else:
            try:
                pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
            except:
                pass
        
    except Exception as e:
        print(f"❌ Error loading T2I model: {e}")
        if is_turbo:
            CURRENT_TURBO_PIPE = None
            CURRENT_TURBO_MODEL = None
        else:
            CURRENT_T2I_PIPE = None
            CURRENT_T2I_MODEL = None
        raise

def colorize_sd15(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
                  prompt, negative_prompt, seed, steps, scale, cn_weight):
    """Colorize function for SD1.5 models"""
    try:
        if base_model not in SD15_MODELS:
            error_img = Image.new('RGB', (512, 512), color='red')
            return error_img, Image.new('RGB', (512, 512), color='gray')
        
        if controlnet_type not in CONTROLNET_MODELS_SD15:
            error_img = Image.new('RGB', (512, 512), color='red')
            return error_img, Image.new('RGB', (512, 512), color='gray')
        
        pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight, vae_model)
        
        status_msg = f"🎨 Using: {base_model} + {controlnet_type}"
        if lora_model and lora_model != "None":
            status_msg += f" + {lora_model}"
        print(status_msg)
        
        condition_img = prepare_condition_image(sketch, controlnet_type, is_sdxl=False)
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        with torch.inference_mode():
            out = pipe(
                prompt,
                negative_prompt=negative_prompt,
                image=condition_img,
                num_inference_steps=int(steps),
                guidance_scale=float(scale),
                controlnet_conditioning_scale=float(cn_weight),
                generator=gen,
                height=512,
                width=512
            ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return out, condition_img
    except Exception as e:
        print(f"❌ Error in colorize_sd15: {e}")
        error_img = Image.new('RGB', (512, 512), color='red')
        return error_img, Image.new('RGB', (512, 512), color='gray')

def colorize_sdxl(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
                  prompt, negative_prompt, seed, steps, scale, cn_weight):
    """Colorize function for SDXL models"""
    try:
        if base_model not in SDXL_MODELS:
            error_img = Image.new('RGB', (1024, 1024), color='red')
            return error_img, Image.new('RGB', (1024, 1024), color='gray')
        
        if controlnet_type not in CONTROLNET_MODELS_SDXL:
            error_img = Image.new('RGB', (1024, 1024), color='red')
            return error_img, Image.new('RGB', (1024, 1024), color='gray')
        
        pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight, vae_model)
        
        status_msg = f"🎨 Using: {base_model} + {controlnet_type}"
        if lora_model and lora_model != "None":
            status_msg += f" + {lora_model}"
        print(status_msg)
        
        condition_img = prepare_condition_image(sketch, controlnet_type, is_sdxl=True)
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        with torch.inference_mode():
            out = pipe(
                prompt,
                negative_prompt=negative_prompt,
                image=condition_img,
                num_inference_steps=int(steps),
                guidance_scale=float(scale),
                controlnet_conditioning_scale=float(cn_weight),
                generator=gen,
                height=1024,
                width=1024
            ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return out, condition_img
    except Exception as e:
        print(f"❌ Error in colorize_sdxl: {e}")
        error_img = Image.new('RGB', (1024, 1024), color='red')
        return error_img, Image.new('RGB', (1024, 1024), color='gray')

def colorize_waifu_xl(sketch, lora_weight, vae_model, prompt, negative_prompt, seed, steps, scale, cn_weight):
    """Colorize function specifically for Waifu-Colorize-XL model"""
    try:
        model_name = "stabilityai/stable-diffusion-xl-base-1.0"
        controlnet_type = "lineart_sdxl"
        lora_model = "colorize-xl"
        
        pipe = get_pipeline(model_name, controlnet_type, lora_model, lora_weight, vae_model)
        
        status_msg = f"🎨 Using: Waifu-Colorize-XL (Lineart ControlNet)"
        print(status_msg)
        
        condition_img = prepare_condition_image(sketch, controlnet_type, is_sdxl=True)
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        # สำหรับ Waifu-Colorize-XL ใช้ prompt พิเศษ
        enhanced_prompt = f"colorized, vibrant colors, anime style, {prompt}" if prompt else "colorized, vibrant colors, anime style, masterpiece"
        
        with torch.inference_mode():
            out = pipe(
                enhanced_prompt,
                negative_prompt=negative_prompt,
                image=condition_img,
                num_inference_steps=int(steps),
                guidance_scale=float(scale),
                controlnet_conditioning_scale=float(cn_weight),
                generator=gen,
                height=1024,
                width=1024
            ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return out, condition_img
    except Exception as e:
        print(f"❌ Error in colorize_waifu_xl: {e}")
        error_img = Image.new('RGB', (1024, 1024), color='red')
        return error_img, Image.new('RGB', (1024, 1024), color='gray')

def t2i_sd15(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
             seed, steps, scale, w, h):
    """Text-to-image for SD1.5 models"""
    try:
        if model not in SD15_MODELS:
            error_img = Image.new('RGB', (int(w), int(h)), color='red')
            return error_img
        
        load_t2i_model(model, lora_model, lora_weight, vae_model)
        
        print(f"🖼️ Using SD1.5 model: {model}")
        if lora_model and lora_model != "None":
            print(f"   with LoRA: {lora_model} (weight: {lora_weight})")
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        with torch.inference_mode():
            result = CURRENT_T2I_PIPE(
                prompt,
                negative_prompt=negative_prompt,
                width=int(w),
                height=int(h),
                num_inference_steps=int(steps),
                guidance_scale=float(scale),
                generator=gen
            ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return result
    except Exception as e:
        print(f"❌ Error in t2i_sd15: {e}")
        error_img = Image.new('RGB', (int(w), int(h)), color='red')
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(error_img)
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except:
            font = ImageFont.load_default()
        draw.text((50, 50), f"Error: {str(e)[:50]}...", fill="white", font=font)
        return error_img

def t2i_sdxl(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
             seed, steps, scale, w, h, use_refiner=False):
    """Text-to-image for SDXL models"""
    try:
        if model not in SDXL_MODELS:
            error_img = Image.new('RGB', (int(w), int(h)), color='red')
            return error_img
        
        model_to_load = model
        if use_refiner and "refiner" not in model.lower() and not is_turbo_model(model):
            model_to_load = "stabilityai/stable-diffusion-xl-refiner-1.0"
        
        load_t2i_model(model_to_load, lora_model, lora_weight, vae_model)
        
        print(f"🖼️ Using SDXL model: {model}")
        if lora_model and lora_model != "None":
            print(f"   with LoRA: {lora_model} (weight: {lora_weight})")
        if use_refiner:
            print(f"   with refiner")
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        with torch.inference_mode():
            if use_refiner and CURRENT_SDXL_REFINER is not None and not is_turbo_model(model):
                image = CURRENT_T2I_PIPE(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=int(w),
                    height=int(h),
                    num_inference_steps=int(steps//2),
                    guidance_scale=float(scale),
                    generator=gen,
                    output_type="latent"
                ).images
                
                result = CURRENT_SDXL_REFINER(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    image=image,
                    num_inference_steps=int(steps//2),
                    guidance_scale=float(scale),
                    generator=gen
                ).images[0]
            else:
                if is_turbo_model(model):
                    # สำหรับ Turbo models ใช้ steps น้อยลง
                    turbo_steps = max(1, min(10, int(steps)))
                    result = CURRENT_TURBO_PIPE(
                        prompt,
                        negative_prompt=negative_prompt,
                        width=int(w),
                        height=int(h),
                        num_inference_steps=turbo_steps,
                        guidance_scale=float(scale),
                        generator=gen
                    ).images[0]
                else:
                    result = CURRENT_T2I_PIPE(
                        prompt,
                        negative_prompt=negative_prompt,
                        width=int(w),
                        height=int(h),
                        num_inference_steps=int(steps),
                        guidance_scale=float(scale),
                        generator=gen
                    ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return result
    except Exception as e:
        print(f"❌ Error in t2i_sdxl: {e}")
        error_img = Image.new('RGB', (int(w), int(h)), color='red')
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(error_img)
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except:
            font = ImageFont.load_default()
        draw.text((50, 50), f"Error: {str(e)[:50]}...", fill="white", font=font)
        return error_img

def t2i_turbo(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
              seed, steps, scale, w, h):
    """Text-to-image for Turbo models (fast generation)"""
    try:
        if model not in TURBO_MODELS:
            error_img = Image.new('RGB', (int(w), int(h)), color='red')
            return error_img
        
        load_t2i_model(model, lora_model, lora_weight, vae_model)
        
        print(f"⚡ Using Turbo model: {model}")
        if lora_model and lora_model != "None":
            print(f"   with LoRA: {lora_model} (weight: {lora_weight})")
        
        gen = torch.Generator(device=device).manual_seed(int(seed))
        
        with torch.inference_mode():
            # สำหรับ Turbo models ใช้ steps น้อยลง (1-10 steps)
            turbo_steps = max(1, min(10, int(steps)))
            
            if is_sdxl_model(model):
                # SDXL-Turbo
                result = CURRENT_TURBO_PIPE(
                    prompt,
                    negative_prompt=negative_prompt,
                    width=int(w),
                    height=int(h),
                    num_inference_steps=turbo_steps,
                    guidance_scale=float(scale),
                    generator=gen
                ).images[0]
            else:
                # SD-Turbo
                result = CURRENT_TURBO_PIPE(
                    prompt,
                    negative_prompt=negative_prompt,
                    width=int(w),
                    height=int(h),
                    num_inference_steps=turbo_steps,
                    guidance_scale=float(scale),
                    generator=gen
                ).images[0]
        
        if device.type == "cuda":
            torch.cuda.empty_cache()
        
        return result
    except Exception as e:
        print(f"❌ Error in t2i_turbo: {e}")
        error_img = Image.new('RGB', (int(w), int(h)), color='red')
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(error_img)
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except:
            font = ImageFont.load_default()
        draw.text((50, 50), f"Error: {str(e)[:50]}...", fill="white", font=font)
        return error_img

def unload_all_models():
    global CURRENT_CONTROLNET_PIPE, CURRENT_CONTROLNET_KEY
    global DETECTORS
    global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL, CURRENT_SDXL_REFINER
    global CURRENT_TURBO_PIPE, CURRENT_TURBO_MODEL
    
    print("🗑️ Unloading all models from memory...")
    
    try:
        if CURRENT_CONTROLNET_PIPE is not None:
            del CURRENT_CONTROLNET_PIPE
            CURRENT_CONTROLNET_PIPE = None
    except:
        pass
    CURRENT_CONTROLNET_KEY = None
    
    for detector_type in list(DETECTORS.keys()):
        try:
            del DETECTORS[detector_type]
        except:
            pass
    DETECTORS.clear()
    
    try:
        if CURRENT_T2I_PIPE is not None:
            del CURRENT_T2I_PIPE
            CURRENT_T2I_PIPE = None
    except:
        pass
    
    try:
        if CURRENT_SDXL_REFINER is not None:
            del CURRENT_SDXL_REFINER
            CURRENT_SDXL_REFINER = None
    except:
        pass
    
    try:
        if CURRENT_TURBO_PIPE is not None:
            del CURRENT_TURBO_PIPE
            CURRENT_TURBO_PIPE = None
    except:
        pass
    
    CURRENT_T2I_MODEL = None
    CURRENT_TURBO_MODEL = None
    
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"💾 GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
    
    return "✅ All models unloaded from memory!"

# ===== Gradio UI =====
with gr.Blocks(title="🎨 AI Image Generator Pro", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 AI Image Generator Pro - NSFW Capable")
    gr.Markdown("### Advanced Image Generation with ControlNet, LoRA & VAE Support")
    gr.Markdown("⚠️ **Content Warning:** This tool can generate NSFW content. Use responsibly and in compliance with applicable laws.")
    
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        gr.Markdown(f"**GPU:** {gpu_name} ({gpu_memory:.1f} GB)")
    else:
        gr.Markdown("**⚠️ Running on CPU** - Generation will be slower")
    
    with gr.Row():
        unload_btn = gr.Button("🗑️ Unload All Models", variant="stop", scale=1)
        status_text = gr.Textbox(label="Status", interactive=False, scale=3)
    unload_btn.click(unload_all_models, outputs=status_text)

    with gr.Tab("🎨 SD1.5 ControlNet"):
        gr.Markdown("""
        ### Transform sketches/images using SD1.5 with ControlNet
        - **Supports:** lineart, lineart_anime, canny, depth, normal, openpose, softedge, segmentation, mlsd, shuffle, scribble, tile
        - **Best Resolution:** 512x512
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                inp_sd15 = gr.Image(label="Input Sketch/Image", type="pil")
                
                gr.Markdown("### Model Settings")
                base_model_sd15 = gr.Dropdown(
                    choices=SD15_MODELS,
                    value="digiplay/ChikMix_V3",
                    label="SD1.5 Base Model"
                )
                controlnet_type_sd15 = gr.Dropdown(
                    choices=list(CONTROLNET_MODELS_SD15.keys()),
                    value="lineart_anime",
                    label="ControlNet Type"
                )
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    lora_model_sd15 = gr.Dropdown(
                        choices=list(LORA_MODELS.keys()),
                        value="None",
                        label="LoRA Model"
                    )
                    lora_weight_sd15 = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                
                vae_model_sd15 = gr.Dropdown(
                    choices=["None", "SD1.5 VAE", "Anime VAE"],
                    value="None",
                    label="VAE Model (Optional)"
                )
            
            with gr.Column(scale=1):
                out_sd15 = gr.Image(label="Generated Output")
                condition_out_sd15 = gr.Image(label="Processed Condition", type="pil")
        
        gr.Markdown("### Generation Parameters")
        with gr.Row():
            prompt_sd15 = gr.Textbox(
                label="Prompt",
                placeholder="masterpiece, best quality, 1girl, beautiful detailed eyes, long hair",
                lines=3
            )
            negative_prompt_sd15 = gr.Textbox(
                label="Negative Prompt",
                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers",
                lines=3
            )
        
        with gr.Row():
            seed_sd15 = gr.Number(value=-1, label="Seed (-1 for random)")
            steps_sd15 = gr.Slider(10, 100, 30, step=1, label="Steps")
            scale_sd15 = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
            cn_weight_sd15 = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
        
        run_sd15 = gr.Button("🎨 Generate (SD1.5)", variant="primary", size="lg")
        run_sd15.click(
            colorize_sd15,
            [inp_sd15, base_model_sd15, controlnet_type_sd15, lora_model_sd15, lora_weight_sd15, vae_model_sd15,
             prompt_sd15, negative_prompt_sd15, seed_sd15, steps_sd15, scale_sd15, cn_weight_sd15],
            [out_sd15, condition_out_sd15]
        )

    with gr.Tab("🎨 SDXL ControlNet"):
        gr.Markdown("""
        ### Transform sketches/images using SDXL with ControlNet
        - **Supports:** canny_sdxl, depth_sdxl, openpose_sdxl, lineart_sdxl (new!)
        - **Best Resolution:** 1024x1024
        - **Higher quality, more VRAM required**
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                inp_sdxl = gr.Image(label="Input Sketch/Image", type="pil")
                
                gr.Markdown("### Model Settings")
                base_model_sdxl = gr.Dropdown(
                    choices=SDXL_MODELS,
                    value="stabilityai/stable-diffusion-xl-base-1.0",
                    label="SDXL Base Model"
                )
                controlnet_type_sdxl = gr.Dropdown(
                    choices=list(CONTROLNET_MODELS_SDXL.keys()),
                    value="lineart_sdxl",
                    label="ControlNet Type"
                )
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    lora_model_sdxl = gr.Dropdown(
                        choices=list(LORA_MODELS.keys()),
                        value="None",
                        label="LoRA Model"
                    )
                    lora_weight_sdxl = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                
                vae_model_sdxl = gr.Dropdown(
                    choices=["None", "SDXL VAE", "Turbo VAE"],
                    value="None",
                    label="VAE Model (Optional)"
                )
            
            with gr.Column(scale=1):
                out_sdxl = gr.Image(label="Generated Output")
                condition_out_sdxl = gr.Image(label="Processed Condition", type="pil")
        
        gr.Markdown("### Generation Parameters")
        with gr.Row():
            prompt_sdxl = gr.Textbox(
                label="Prompt",
                placeholder="masterpiece, best quality, 8k, ultra-detailed, photorealistic, beautiful lighting",
                lines=3
            )
            negative_prompt_sdxl = gr.Textbox(
                label="Negative Prompt",
                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits",
                lines=3
            )
        
        with gr.Row():
            seed_sdxl = gr.Number(value=-1, label="Seed (-1 for random)")
            steps_sdxl = gr.Slider(10, 100, 30, step=1, label="Steps")
            scale_sdxl = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
            cn_weight_sdxl = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
        
        run_sdxl = gr.Button("🎨 Generate (SDXL)", variant="primary", size="lg")
        run_sdxl.click(
            colorize_sdxl,
            [inp_sdxl, base_model_sdxl, controlnet_type_sdxl, lora_model_sdxl, lora_weight_sdxl, vae_model_sdxl,
             prompt_sdxl, negative_prompt_sdxl, seed_sdxl, steps_sdxl, scale_sdxl, cn_weight_sdxl],
            [out_sdxl, condition_out_sdxl]
        )

    with gr.Tab("🌸 Waifu-Colorize-XL"):
        gr.Markdown("""
        ### Specialized Anime Lineart Colorization
        - **Model:** ShinoharaHare/Waifu-Colorize-XL
        - **Specialized for:** Anime/manga lineart coloring
        - **Features:** Automatic colorization with vibrant anime colors
        - **Best Resolution:** 1024x1024
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                inp_waifu = gr.Image(label="Input Lineart/Sketch", type="pil")
                
                gr.Markdown("### Model Settings")
                gr.Markdown("**Using:** ShinoharaHare/Waifu-Colorize-XL (Specialized Anime Colorization)")
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    lora_weight_waifu = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="LoRA Weight")
                
                vae_model_waifu = gr.Dropdown(
                    choices=["None", "SDXL VAE"],
                    value="None",
                    label="VAE Model (Optional)"
                )
            
            with gr.Column(scale=1):
                out_waifu = gr.Image(label="Colorized Output")
                condition_out_waifu = gr.Image(label="Processed Lineart", type="pil")
        
        gr.Markdown("### Generation Parameters")
        with gr.Row():
            prompt_waifu = gr.Textbox(
                label="Color Style Prompt (Optional)",
                placeholder="vibrant colors, anime style, beautiful coloring, masterpiece",
                lines=2
            )
            negative_prompt_waifu = gr.Textbox(
                label="Negative Prompt",
                placeholder="monochrome, grayscale, black and white, sketch, lineart only",
                lines=2,
                value="monochrome, grayscale, black and white, sketch, lineart only"
            )
        
        with gr.Row():
            seed_waifu = gr.Number(value=-1, label="Seed (-1 for random)")
            steps_waifu = gr.Slider(10, 50, 25, step=1, label="Steps")
            scale_waifu = gr.Slider(1, 20, 7.5, step=0.5, label="CFG Scale")
            cn_weight_waifu = gr.Slider(0.1, 2.0, 1.2, step=0.1, label="ControlNet Weight")
        
        gr.Markdown("### Tips for Best Results:")
        gr.Markdown("""
        1. Use clean lineart for best results
        2. Higher ControlNet weight (1.0-1.5) for better line following
        3. Lower CFG scale (5-8) for more natural coloring
        4. Add color hints in prompt (e.g., "blue hair, red eyes, pink dress")
        5. Keep prompts simple for this specialized model
        """)
        
        run_waifu = gr.Button("🌸 Colorize with Waifu-Colorize-XL", variant="primary", size="lg")
        run_waifu.click(
            colorize_waifu_xl,
            [inp_waifu, lora_weight_waifu, vae_model_waifu,
             prompt_waifu, negative_prompt_waifu, seed_waifu, steps_waifu, scale_waifu, cn_weight_waifu],
            [out_waifu, condition_out_waifu]
        )

    with gr.Tab("🖼️ SD1.5 Text-to-Image"):
        gr.Markdown("""
        ### Generate images from text descriptions using SD1.5
        - **Best Resolution:** 512x512, 512x768, 768x512
        - **Faster generation, lower VRAM usage**
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Model Configuration")
                t2i_model_sd15 = gr.Dropdown(
                    choices=SD15_MODELS,
                    value="digiplay/ChikMix_V3",
                    label="SD1.5 Base Model"
                )
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    t2i_lora_sd15 = gr.Dropdown(
                        choices=list(LORA_MODELS.keys()),
                        value="None",
                        label="LoRA Model"
                    )
                    t2i_lora_weight_sd15 = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                
                t2i_vae_sd15 = gr.Dropdown(
                    choices=["None", "SD1.5 VAE", "Anime VAE"],
                    value="None",
                    label="VAE Model"
                )
            
            with gr.Column(scale=1):
                t2i_out_sd15 = gr.Image(label="Generated Image", type="pil")
        
        gr.Markdown("### Prompts")
        with gr.Row():
            t2i_prompt_sd15 = gr.Textbox(
                label="Prompt",
                lines=4,
                placeholder="masterpiece, best quality, highly detailed, beautiful, 1girl"
            )
            t2i_negative_prompt_sd15 = gr.Textbox(
                label="Negative Prompt",
                lines=4,
                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
            )
        
        gr.Markdown("### Generation Parameters")
        with gr.Row():
            t2i_seed_sd15 = gr.Number(value=-1, label="Seed (-1 for random)")
            t2i_steps_sd15 = gr.Slider(10, 100, 30, step=1, label="Steps")
            t2i_scale_sd15 = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
        
        with gr.Row():
            w_sd15 = gr.Slider(256, 1024, 512, step=64, label="Width")
            h_sd15 = gr.Slider(256, 1024, 768, step=64, label="Height")
        
        gen_btn_sd15 = gr.Button("🖼️ Generate (SD1.5)", variant="primary", size="lg")
        gen_btn_sd15.click(
            t2i_sd15,
            [t2i_prompt_sd15, t2i_negative_prompt_sd15, t2i_model_sd15, t2i_lora_sd15, t2i_lora_weight_sd15,
             t2i_vae_sd15, t2i_seed_sd15, t2i_steps_sd15, t2i_scale_sd15, w_sd15, h_sd15],
            t2i_out_sd15
        )

    with gr.Tab("⚡ SDXL-Turbo Text-to-Image"):
        gr.Markdown("""
        ### Ultra-Fast Image Generation with SDXL-Turbo
        - **Model:** stabilityai/sdxl-turbo
        - **Features:** 1-4 steps generation, extremely fast
        - **Best Resolution:** 512x512 to 1024x1024
        - **Warning:** Lower quality than full SDXL but much faster
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Model Configuration")
                t2i_model_turbo = gr.Dropdown(
                    choices=TURBO_MODELS,
                    value="stabilityai/sdxl-turbo",
                    label="Turbo Model"
                )
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    t2i_lora_turbo = gr.Dropdown(
                        choices=list(LORA_MODELS.keys()),
                        value="None",
                        label="LoRA Model"
                    )
                    t2i_lora_weight_turbo = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                
                t2i_vae_turbo = gr.Dropdown(
                    choices=["None", "Turbo VAE", "SDXL VAE"],
                    value="None",
                    label="VAE Model"
                )
            
            with gr.Column(scale=1):
                t2i_out_turbo = gr.Image(label="Generated Image", type="pil")
        
        gr.Markdown("### Prompts")
        with gr.Row():
            t2i_prompt_turbo = gr.Textbox(
                label="Prompt",
                lines=4,
                placeholder="masterpiece, best quality, highly detailed, beautiful, 1girl"
            )
            t2i_negative_prompt_turbo = gr.Textbox(
                label="Negative Prompt",
                lines=4,
                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
            )
        
        gr.Markdown("### Generation Parameters (Turbo needs only 1-4 steps!)")
        with gr.Row():
            t2i_seed_turbo = gr.Number(value=-1, label="Seed (-1 for random)")
            t2i_steps_turbo = gr.Slider(1, 10, 4, step=1, label="Steps (1-4 recommended)")
            t2i_scale_turbo = gr.Slider(0, 10, 0.0, step=0.5, label="CFG Scale (0-2 recommended)")
        
        with gr.Row():
            w_turbo = gr.Slider(256, 1024, 512, step=64, label="Width")
            h_turbo = gr.Slider(256, 1024, 512, step=64, label="Height")
        
        gr.Markdown("### Turbo Model Tips:")
        gr.Markdown("""
        1. **Use very few steps:** 1-4 steps is enough!
        2. **Low CFG Scale:** 0.0-2.0 works best
        3. **Fast but lower quality:** For quick previews/testing
        4. **Works well with:** Simple prompts, concept testing
        """)
        
        gen_btn_turbo = gr.Button("⚡ Generate with Turbo (Fast!)", variant="primary", size="lg")
        gen_btn_turbo.click(
            t2i_turbo,
            [t2i_prompt_turbo, t2i_negative_prompt_turbo, t2i_model_turbo, t2i_lora_turbo, t2i_lora_weight_turbo,
             t2i_vae_turbo, t2i_seed_turbo, t2i_steps_turbo, t2i_scale_turbo, w_turbo, h_turbo],
            t2i_out_turbo
        )

    with gr.Tab("🖼️ SDXL Text-to-Image"):
        gr.Markdown("""
        ### Generate images from text descriptions using SDXL
        - **Best Resolution:** 1024x1024, 1024x1536, 1536x1024
        - **Higher quality, more detail, better composition**
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Model Configuration")
                t2i_model_sdxl = gr.Dropdown(
                    choices=[m for m in SDXL_MODELS if m not in TURBO_MODELS],  # ไม่รวม Turbo models
                    value="stabilityai/stable-diffusion-xl-base-1.0",
                    label="SDXL Base Model"
                )
                
                gr.Markdown("### Enhancement Options")
                with gr.Row():
                    t2i_lora_sdxl = gr.Dropdown(
                        choices=list(LORA_MODELS.keys()),
                        value="None",
                        label="LoRA Model"
                    )
                    t2i_lora_weight_sdxl = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                
                t2i_vae_sdxl = gr.Dropdown(
                    choices=["None", "SDXL VAE"],
                    value="None",
                    label="VAE Model"
                )
                
                use_refiner_sdxl = gr.Checkbox(
                    label="Use Refiner (for better quality)",
                    value=False
                )
            
            with gr.Column(scale=1):
                t2i_out_sdxl = gr.Image(label="Generated Image", type="pil")
        
        gr.Markdown("### Prompts")
        with gr.Row():
            t2i_prompt_sdxl = gr.Textbox(
                label="Prompt",
                lines=4,
                placeholder="masterpiece, best quality, 8k, ultra-detailed, photorealistic, cinematic lighting"
            )
            t2i_negative_prompt_sdxl = gr.Textbox(
                label="Negative Prompt",
                lines=4,
                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry"
            )
        
        gr.Markdown("### Generation Parameters")
        with gr.Row():
            t2i_seed_sdxl = gr.Number(value=-1, label="Seed (-1 for random)")
            t2i_steps_sdxl = gr.Slider(10, 100, 30, step=1, label="Steps")
            t2i_scale_sdxl = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
        
        with gr.Row():
            w_sdxl = gr.Slider(512, 2048, 1024, step=64, label="Width")
            h_sdxl = gr.Slider(512, 2048, 1024, step=64, label="Height")
        
        gen_btn_sdxl = gr.Button("🖼️ Generate (SDXL)", variant="primary", size="lg")
        gen_btn_sdxl.click(
            t2i_sdxl,
            [t2i_prompt_sdxl, t2i_negative_prompt_sdxl, t2i_model_sdxl, t2i_lora_sdxl, t2i_lora_weight_sdxl,
             t2i_vae_sdxl, t2i_seed_sdxl, t2i_steps_sdxl, t2i_scale_sdxl, w_sdxl, h_sdxl, use_refiner_sdxl],
            t2i_out_sdxl
        )
    
    with gr.Tab("📚 Quick Reference"):
        gr.Markdown("""
        # Model & Feature Guide
        
        ## 🎯 Model Comparison
        
        ### SD1.5 (Stable Diffusion 1.5)
        - **Pros:** Fast, low VRAM, many models
        - **Cons:** 512px max, lower quality
        - **Best for:** Quick tests, anime, lower-end hardware
        
        ### SDXL (Stable Diffusion XL)
        - **Pros:** 1024px+, high quality, better composition
        - **Cons:** High VRAM, slower
        - **Best for:** Final quality, professional work
        
        ### SDXL-Turbo
        - **Pros:** Extremely fast (1-4 steps!)
        - **Cons:** Lower quality than full SDXL
        - **Best for:** Quick previews, concept testing
        
        ### Waifu-Colorize-XL
        - **Pros:** Specialized for anime lineart coloring
        - **Cons:** Anime-only, requires clean lineart
        - **Best for:** Anime/manga colorization
        
        ## 🎨 New ControlNet for SDXL
        
        ### Lineart ControlNet for SDXL
        - **Model:** ShermanG/ControlNet-Standard-Lineart-for-SDXL
        - **Purpose:** Convert lineart to colored images
        - **Best with:** Clean lineart, anime/manga styles
        - **Used in:** Waifu-Colorize-XL Tab
        
        ## 💎 Recommended Workflows
        
        ### For Anime/Manga Artists
        1. Draw lineart
        2. Use **Waifu-Colorize-XL** Tab for automatic coloring
        3. Or use **SDXL ControlNet** with lineart_sdxl
        
        ### For Quick Concepts
        1. Use **SDXL-Turbo** Tab for 1-4 step generation
        2. Refine in **SDXL Text-to-Image** if needed
        
        ### For Professional Work
        1. Use **SDXL Text-to-Image** with Refiner
        2. High steps (40-50), CFG 7-9
        3. 1024x1024 resolution
        
        ## ⚡ Turbo Model Tips
        
        ### SDXL-Turbo Best Practices:
        - **Steps:** 1-4 only!
        - **CFG Scale:** 0.0-2.0
        - **Prompts:** Keep simple
        - **Resolution:** 512x512 to 1024x1024
        - **Use case:** Storyboarding, concept art, quick iterations
        
        ## 🌸 Waifu-Colorize-XL Tips
        
        ### For Best Results:
        1. **Clean lineart:** No stray marks
        2. **ControlNet weight:** 1.0-1.5
        3. **CFG Scale:** 5-8
        4. **Simple prompts:** "vibrant colors, anime style"
        5. **Resolution:** 1024x1024
        
        ### Example Workflow:
        1. Draw/sketch in your favorite app
        2. Export as clean lineart (black on white)
        3. Upload to Waifu-Colorize-XL Tab
        4. Adjust parameters as needed
        5. Generate and refine
        
        ## 🚀 Performance Optimization
        
        ### Low VRAM (<8GB)
        - Use SD1.5 models only
        - 512x512 resolution
        - Enable attention slicing
        
        ### Medium VRAM (8-12GB)
        - SD1.5 and SDXL (no refiner)
        - 1024x1024 for SDXL
        - Enable xFormers
        
        ### High VRAM (12GB+)
        - All models including SDXL with refiner
        - Higher resolutions
        - Multiple LoRAs
        
        ## 🔄 Memory Management
        
        ### When to Unload Models:
        1. Switching between SD1.5 and SDXL
        2. Getting "out of memory" errors
        3. Changing ControlNet types
        4. After long generation sessions
        
        ### Memory Saving Tips:
        1. Use "Unload All Models" button
        2. Generate in batches
        3. Lower resolution for testing
        4. Close other GPU applications
        """)

try:
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        quiet=False
    )
except Exception as e:
    print(f"❌ Error launching Gradio app: {e}")