# ===== 必须首先导入spaces =====
try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ Spaces available - ZeroGPU mode")
except ImportError:
    SPACES_AVAILABLE = False
    print("⚠️ Spaces not available - running in regular mode")

# ===== 其他导入 =====
import os
import uuid
from datetime import datetime
import random
import torch
import gradio as gr
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
from PIL import Image
import traceback
import numpy as np

# ===== 新增：长提示词处理 =====
try:
    from compel import Compel, ReturnedEmbeddingsType
    COMPEL_AVAILABLE = True
    print("✅ Compel available for long prompt processing")
except ImportError:
    COMPEL_AVAILABLE = False
    print("⚠️ Compel not available - using standard prompt processing")

# ===== 🔥 优化后的LoRA配置 =====
LORA_MODELS = {
    "sdxl_lightning": {
        "repo_id": "ByteDance/SDXL-Lightning", 
        "filename": "sdxl_lightning_4step_lora.safetensors",
        "scale": 1.0,  # Lightning LoRA建议使用1.0
        "description": "SDXL Lightning 4-step acceleration",
        "recommended_steps": 4,
        "recommended_cfg": 0.0,  # Lightning LoRA需要CFG=0
        "enabled": True
    },
    "quality_enhancer": {
        "repo_id": "artificialguybr/LogoRedmond-LogoLoraForSDXL-V2",
        "filename": "LogoRedAF.safetensors", 
        "scale": 0.8,
        "description": "Quality and realism enhancer",
        "recommended_steps": 20,
        "recommended_cfg": 7.0,
        "enabled": True
    },
    "detail_enhancer": {
        "repo_id": "stabilityai/stable-diffusion-xl-base-1.0",  # 使用基础模型作为fallback
        "filename": None,  # 不实际加载，只是占位
        "scale": 0.6,
        "description": "Detail enhancement (placeholder)",
        "recommended_steps": 28,
        "recommended_cfg": 7.0,
        "enabled": False  # 默认禁用
    }
}

# ===== 配置 =====
STYLE_PRESETS = {
    "None": "",
    "Realistic": "photorealistic, 8k, ultra-detailed, cinematic lighting, masterpiece, realistic skin texture, detailed anatomy",
    "Anime": "anime style, detailed, high quality, masterpiece, best quality, detailed eyes, perfect anatomy",
    "Oil Painting": "oil painting, rich brush strokes, canvas texture, artistic, baroque lighting, detailed composition", 
    "Comic": "comic book style, bold outlines, vibrant colors, cel shading, dynamic pose",
    "Watercolor": "watercolor illustration, soft gradients, pastel palette, artistic brush strokes",
    "Portrait": "portrait photography, professional lighting, detailed face, studio quality, sharp focus",
    "Fantasy": "fantasy art, magical, ethereal, detailed background, mystical atmosphere, high fantasy"
}

# 🔥 增强的NSFW专用提示词（模拟LoRA效果）
NSFW_ENHANCERS = [
    "detailed anatomy", "(perfect anatomy:1.2)", "soft skin", "natural lighting", 
    "high resolution", "(masterpiece:1.3)", "(best quality:1.2)",
    "professional photography", "artistic composition",
    "(perfect proportions:1.1)", "smooth textures", "intimate lighting",
    "realistic skin texture", "(detailed face:1.1)", "natural pose"
]

# 🔥 风格专用增强词
STYLE_ENHANCERS = {
    "Realistic": ["photorealistic", "(ultra realistic:1.2)", "natural lighting", "detailed skin", "professional photography"],
    "Anime": ["anime style", "(high quality anime:1.2)", "detailed eyes", "perfect face", "clean art style"],
    "Oil Painting": ["oil painting style", "artistic", "brush strokes", "classical art"],
    "Portrait": ["portrait photography", "professional lighting", "studio quality", "(detailed face:1.3)"],
    "Fantasy": ["fantasy art", "magical", "ethereal", "mystical atmosphere"]
}

SAVE_DIR = "generated_images"
os.makedirs(SAVE_DIR, exist_ok=True)

# ===== 模型相关变量 =====
pipeline = None
compel_processor = None
device = None
model_loaded = False
current_loras = {}  # 🔥 新增：跟踪当前加载的LoRA
current_model = None

# ===== NSFW优化模型列表 =====
NSFW_MODELS = {
    "PornMasterPro": "votepurchase/pornmasterPro_noobV3VAE",
    "CounterfeitV3": "votepurchase/counterfeitV30_v30", 
    "PornMaster SDXL": "John6666/pornmaster-pro-sdxlv1-sdxl",
    "MeinaMix": "Meina/MeinaMix_V11",
    "MeinaUnreal": "Meina/MeinaUnreal_V5",
    "SDXL Base": "stabilityai/stable-diffusion-xl-base-1.0"
}

def load_lora_weights(pipeline, lora_name, lora_config):
    """🔥 优化的LoRA加载函数"""
    if not lora_config["enabled"]:
        return False
        
    try:
        repo_id = lora_config["repo_id"]
        filename = lora_config["filename"]
        
        if filename is None:  # 跳过占位符
            return False
            
        scale = lora_config["scale"]
        
        print(f"🔧 Loading LoRA: {lora_name} - {repo_id}/{filename} (scale: {scale})")
        
        # 检查是否已经加载
        if lora_name in current_loras:
            print(f"✅ LoRA {lora_name} already loaded")
            return True
        
        # 尝试加载LoRA权重
        pipeline.load_lora_weights(
            repo_id, 
            weight_name=filename,
            adapter_name=lora_name
        )
        
        current_loras[lora_name] = lora_config
        print(f"✅ LoRA loaded successfully: {lora_config['description']}")
        return True
        
    except Exception as e:
        print(f"⚠️ Failed to load LoRA {lora_name}: {e}")
        return False

def get_best_nsfw_model():
    """尝试获取最佳可用的NSFW模型"""
    for model_name, model_id in NSFW_MODELS.items():
        try:
            print(f"🔍 Checking model availability: {model_name} ({model_id})")
            return model_id, model_name
        except Exception as e:
            print(f"❌ Model {model_name} check failed: {e}")
            continue
    
    return list(NSFW_MODELS.values())[0], list(NSFW_MODELS.keys())[0]

def initialize_model(selected_model=None):
    """🔥 优化的模型初始化函数"""
    global pipeline, compel_processor, device, model_loaded, current_loras, current_model
    
    # 如果模型已加载且是相同模型，直接返回
    if model_loaded and pipeline is not None and current_model == selected_model:
        return True
    
    # 如果切换模型，清理之前的状态
    if current_model != selected_model and pipeline is not None:
        del pipeline
        pipeline = None
        current_loras.clear()
        model_loaded = False
        torch.cuda.empty_cache()
        
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"🖥️ Using device: {device}")
        
        # 选择模型
        if selected_model and selected_model in NSFW_MODELS.values():
            model_id = selected_model
            model_name = [k for k, v in NSFW_MODELS.items() if v == selected_model][0]
        else:
            model_id, model_name = get_best_nsfw_model()
        
        print(f"📦 Loading model: {model_name} ({model_id})")
        
        # 基础模型加载
        pipeline = StableDiffusionXLPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            variant="fp16" if torch.cuda.is_available() else None,
            use_safetensors=True,
            safety_checker=None,
            requires_safety_checker=False
        )
        
        # 🔥 优化：使用EulerDiscreteScheduler（SDXL-Lightning推荐）
        pipeline.scheduler = EulerDiscreteScheduler.from_config(
            pipeline.scheduler.config,
            timestep_spacing="trailing"  # Lightning LoRA需要
        )
        pipeline = pipeline.to(device)
        
        # 统一数据类型
        if torch.cuda.is_available():
            pipeline.text_encoder.to(torch.float16)
            pipeline.text_encoder_2.to(torch.float16)
            pipeline.vae.to(torch.float16)
            pipeline.unet.to(torch.float16)
        
        # 🔥 优化：减少内存优化以提高速度
        if torch.cuda.is_available():
            try:
                pipeline.enable_vae_slicing()
                pipeline.enable_attention_slicing()
                # 移除CPU offload以提高速度
                # pipeline.enable_model_cpu_offload()
                try:
                    pipeline.enable_xformers_memory_efficient_attention()
                except:
                    pass
            except Exception as mem_error:
                print(f"⚠️ Memory optimization warning: {mem_error}")
        
        # 初始化Compel
        if COMPEL_AVAILABLE:
            try:
                compel_processor = Compel(
                    tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
                    text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
                    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
                    requires_pooled=[False, True],
                    truncate_long_prompts=False
                )
                print("✅ Long prompt processor (Compel) initialized successfully")
            except Exception as compel_error:
                print(f"⚠️ Compel initialization failed: {compel_error}")
                compel_processor = None
        
        current_model = model_id
        model_loaded = True
        print(f"✅ Successfully loaded model: {model_name}")
        return True
        
    except Exception as e:
        print(f"❌ Critical model loading error: {e}")
        print(traceback.format_exc())
        model_loaded = False
        return False

def update_lora_settings(lora_settings):
    """🔥 新增：更新LoRA设置"""
    for lora_name, settings in lora_settings.items():
        if lora_name in LORA_MODELS:
            LORA_MODELS[lora_name].update(settings)

def apply_loras(pipeline, lora_settings):
    """🔥 新增：应用选中的LoRA"""
    active_loras = []
    
    for lora_name, config in LORA_MODELS.items():
        # 更新配置
        if lora_name in lora_settings:
            config.update(lora_settings[lora_name])
        
        if config.get("enabled", False):
            if load_lora_weights(pipeline, lora_name, config):
                active_loras.append((lora_name, config))
    
    # 设置adapter权重
    if active_loras:
        try:
            adapter_names = [name for name, _ in active_loras]
            adapter_weights = [config["scale"] for _, config in active_loras]
            
            pipeline.set_adapters(adapter_names, adapter_weights=adapter_weights)
            print(f"✅ Activated {len(active_loras)} LoRA adapters: {adapter_names}")
            return active_loras
        except Exception as adapter_error:
            print(f"⚠️ LoRA adapter setup warning: {adapter_error}")
            return []
    
    return []

def enhance_nsfw_prompt(prompt: str, style: str) -> str:
    """🔥 增强NSFW提示词（LoRA优化版）"""
    # 基础增强词
    quality_terms = ", ".join(NSFW_ENHANCERS)
    
    # 风格专用增强词
    style_terms = ""
    if style in STYLE_ENHANCERS:
        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style])
    
    # 风格预设
    style_suffix = STYLE_PRESETS.get(style, "")
    
    # 组合所有增强
    enhanced_parts = [prompt.strip()]
    
    if style_suffix:
        enhanced_parts.append(style_suffix)
    
    if style_terms:
        enhanced_parts.append(style_terms.lstrip(", "))
    
    enhanced_parts.append(quality_terms)
    
    enhanced_prompt = ", ".join(filter(None, enhanced_parts))
    return enhanced_prompt

def process_long_prompt(prompt, negative_prompt=""):
    """处理长提示词"""
    if not compel_processor:
        return None, None
    
    try:
        conditioning, pooled = compel_processor([prompt, negative_prompt])
        return conditioning, pooled
    except Exception as e:
        print(f"Long prompt processing failed: {e}")
        return None, None

def apply_spaces_decorator(func):
    """应用spaces装饰器"""
    if SPACES_AVAILABLE:
        return spaces.GPU(duration=60)(func)
    return func

@apply_spaces_decorator
def generate_image(prompt: str, selected_model: str, style: str, lora_settings: dict, negative_prompt: str = "", steps: int = 28, cfg_scale: float = 7.0, progress=gr.Progress()):
    """🔥 优化的NSFW图像生成函数"""
    if not prompt or prompt.strip() == "":
        return None, "❌ Please enter a prompt!"
    
    # 初始化模型
    progress(0.1, desc="🔥 Loading optimized NSFW model...")
    if not initialize_model(selected_model):
        return None, "❌ Failed to load NSFW model"
    
    progress(0.2, desc="🎨 Applying LoRA configurations...")
    
    try:
        # 🔥 应用LoRA设置
        active_loras = apply_loras(pipeline, lora_settings)
        
        # 🔥 根据LoRA调整参数
        adjusted_steps = steps
        adjusted_cfg = cfg_scale
        
        # 如果使用Lightning LoRA，使用推荐参数
        if any(name == "sdxl_lightning" for name, _ in active_loras):
            adjusted_steps = 4
            adjusted_cfg = 0.0
            print("🔥 Using SDXL-Lightning optimized parameters: steps=4, cfg=0.0")
        
        # 🔥 LoRA增强的提示词处理
        enhanced_prompt = enhance_nsfw_prompt(prompt.strip(), style)
        
        # 🔥 增强的负面提示词
        if not negative_prompt.strip():
            negative_prompt = "(low quality, worst quality:1.4), (bad anatomy, bad hands:1.2), blurry, watermark, signature, text, error, missing limbs, extra limbs, cropped, normal quality, jpeg artifacts, deformed, mutated"
        
        # 生成参数
        seed = random.randint(0, np.iinfo(np.int32).max)
        generator = torch.Generator(device).manual_seed(seed)
        
        progress(0.4, desc="🚀 Generating with LoRA-enhanced pipeline...")
        
        # 长提示词处理
        use_long_prompt = len(enhanced_prompt.split()) > 60 or len(enhanced_prompt) > 300
        
        if use_long_prompt and compel_processor:
            print("Using long prompt processing with LoRA...")
            progress(0.5, desc="📝 Processing long prompt with LoRA enhancements...")
            
            conditioning, pooled = process_long_prompt(enhanced_prompt, negative_prompt)
            
            if conditioning is not None:
                result = pipeline(
                    prompt_embeds=conditioning[0:1],
                    pooled_prompt_embeds=pooled[0:1], 
                    negative_prompt_embeds=conditioning[1:2],
                    negative_pooled_prompt_embeds=pooled[1:2],
                    num_inference_steps=adjusted_steps,
                    guidance_scale=adjusted_cfg,
                    width=1024,
                    height=1024,
                    generator=generator
                )
                image = result.images[0]
            else:
                # 回退到标准处理
                result = pipeline(
                    prompt=enhanced_prompt,
                    negative_prompt=negative_prompt,
                    num_inference_steps=adjusted_steps,
                    guidance_scale=adjusted_cfg,
                    width=1024,
                    height=1024,
                    generator=generator
                )
                image = result.images[0]
        else:
            # 标准处理
            progress(0.5, desc="🎯 Generating with LoRA-enhanced standard processing...")
            result = pipeline(
                prompt=enhanced_prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=adjusted_steps,
                guidance_scale=adjusted_cfg,
                width=1024,
                height=1024,
                generator=generator
            )
            image = result.images[0]
        
        progress(0.9, desc="💾 Saving high-quality image...")
        
        # 保存图像
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"nsfw_lora_{timestamp}_{uuid.uuid4().hex[:8]}.png"
        filepath = os.path.join(SAVE_DIR, filename)
        image.save(filepath, quality=95, optimize=True)
        
        # 保存元数据
        try:
            with open(os.path.join(SAVE_DIR, "metadata.txt"), "a", encoding="utf-8") as f:
                f.write(f"{filename}|{enhanced_prompt}|{seed}|{timestamp}|{adjusted_steps}|{adjusted_cfg}|LoRAs: {len(active_loras)}\n")
        except Exception as save_error:
            print(f"Metadata save warning: {save_error}")
        
        progress(1.0, desc="✅ Complete!")
        
        # 🔥 增强的状态信息（包含LoRA信息）
        lora_info = "\n".join([f"• {name}: {config['description']} (scale: {config['scale']})" 
                              for name, config in active_loras]) if active_loras else "None active"
        
        status_info = f"""✅ LoRA-Enhanced NSFW image generated successfully!
🎯 Seed: {seed}
🖥️ Device: {device}
🏭 Model: {current_model.split('/')[-1] if current_model else 'Unknown'}
📐 Resolution: 1024x1024
⚙️ Steps: {adjusted_steps} | CFG: {adjusted_cfg}
🚀 Long Prompt: {'Yes' if use_long_prompt and compel_processor else 'No'}
📝 Enhanced Prompt Length: {len(enhanced_prompt)} chars
🎨 Active LoRA Adapters ({len(active_loras)}):
{lora_info}
💡 Style Enhancement: {style}"""
        
        return image, status_info
        
    except Exception as e:
        error_msg = str(e)
        print(f"Generation error: {error_msg}")
        print(traceback.format_exc())
        return None, f"❌ Generation failed: {error_msg}"

# ===== CSS样式（保持原有设计）=====
css = """
/* 全局容器 */
.gradio-container {
    max-width: 100% !important;
    margin: 0 auto !important;
    background: linear-gradient(135deg, #e6a4f2 0%, #1197e4 100%) !important;
    min-height: 100vh !important;
    font-family: 'Segoe UI', Arial, sans-serif !important;
}

/* 主要内容区域 */
.main-content {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 25px !important;
    padding: 35px !important;
    margin: 25px !important;
    box-shadow: 0 15px 35px rgba(0,0,0,0.3) !important;
    backdrop-filter: blur(10px) !important;
}

/* 标题 */
.title {
    text-align: center !important;
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    font-size: 3rem !important;
    margin-bottom: 30px !important;
    font-weight: bold !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1) !important;
}

/* 警告信息 */
.warning-box {
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    color: white !important;
    padding: 15px !important;
    border-radius: 10px !important;
    margin-bottom: 20px !important;
    text-align: center !important;
    font-weight: bold !important;
}

/* 输入框 */
.prompt-box textarea {
    min-height: 150px !important;
    border-radius: 12px !important;
    border: 2px solid #bb6ded !important;
    padding: 18px !important;
    font-size: 16px !important;
    resize: vertical !important;
    background: rgba(255,255,255,0.2) !important;
}

.prompt-box textarea:focus {
    border-color: #08676b !important;
    box-shadow: 0 0 15px rgba(77, 8, 161, 0.3) !important;
}

/* 参数控制区域 */
.controls-section {
    background: #f8f9fa !important;
    border-radius: 15px !important;
    padding: 20px !important;
    margin: 15px 0 !important;
}

/* 新增：模型和LoRA选择器 */
.model-selector, .lora-selector {
    background: rgba(89, 74, 232, 0.1) !important;
    border-radius: 12px !important;
    padding: 18px !important;
    border: 2px solid #bb6ded !important;
    margin-bottom: 15px !important;
}

.model-selector label, .lora-selector label {
    font-weight: 600 !important;
    color: #08676b !important;
}

/* 样式选择器 */
.style-selector {
    background: rgba(89, 74, 232, 0.8) !important;
    border-radius: 12px !important;
    padding: 18px !important;
    border: 2px solid #bb6ded !important;
}

.style-selector label {
    font-weight: 600 !important;
    color: #08676b !important;
}

/* 生成按钮 */
.generate-btn {
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    color: white !important;
    border: none !important;
    padding: 20px 40px !important;
    border-radius: 30px !important;
    font-size: 20px !important;
    font-weight: bold !important;
    width: 100% !important;
    margin: 25px 0 !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
}

.generate-btn:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 8px 25px rgba(231, 76, 60, 0.5) !important;
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
}

/* 图片输出区域 */
.image-output {
    border-radius: 20px !important;
    overflow: hidden !important;
    max-width: 1024px !important;
    margin: 0 auto !important;
    border: 3px solid #085a66 !important;
    box-shadow: 0 10px 30px rgba(0,0,0,0.2) !important;
}

.image-output img {
    max-width: 100% !important;
    height: auto !important;
    border-radius: 17px !important;
}

/* 状态信息 */
.status-box {
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    color: white !important;
    border-radius: 10px !important;
    padding: 15px !important;
    margin-top: 20px !important;
    font-weight: 500 !important;
    text-align: center !important;
}

/* 滑块样式 */
.slider-container input[type="range"] {
    accent-color: #555555 !important;
}

/* 响应式设计 */
@media (max-width: 768px) {
    .main-content {
        margin: 15px !important;
        padding: 25px !important;
    }
    
    .title {
        font-size: 2.2rem !important;
    }
    
    .prompt-box textarea {
        min-height: 120px !important;
        font-size: 14px !important;
    }
}

/* 隐藏占位符 */
.gr-image .image-container:empty::before {
    content: "🔥 LoRA-Enhanced NSFW image will appear here 🔥" !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    height: 400px !important;
    background: linear-gradient(45deg, #f8f9fa, #e9ecef) !important;
    border-radius: 15px !important;
    color: #666 !important;
    font-size: 18px !important;
    font-weight: bold !important;
}
"""

# ===== 创建UI（新增模型和LoRA选项卡）=====
def create_interface():
    with gr.Blocks(css=css, title="LoRA-Enhanced NSFW Image Generator") as interface:
        with gr.Column(elem_classes=["main-content"]):
            # 标题
            gr.HTML('<div class="title">🔥 LoRA-Enhanced NSFW Image Generator 🔥</div>')
            
            # 警告信息
            gr.HTML('''
                <div class="warning-box">
                    ⚠️ 18+ CONTENT WARNING ⚠️<br>
                    🎨 LORA-ENHANCED: Multiple Models + Configurable LoRA Adapters + Long Prompt Support<br>
                    ✨ Features: Model Selection + LoRA Parameter Tuning + Style-specific enhancements<br>
                    By using this tool, you confirm you are 18+ and understand the content nature.
                </div>
            ''')
            
            # 主要输入区域
            with gr.Row():
                with gr.Column(scale=2):
                    prompt_input = gr.Textbox(
                        label="🎨 Detailed NSFW Prompt (LoRA-Enhanced Processing)",
                        placeholder="Enter your detailed NSFW prompt here...\n\nExample: beautiful woman, detailed anatomy, soft lighting, artistic pose, bedroom setting, intimate atmosphere\n\n🎨 NEW: LoRA-enhanced processing for superior quality and anatomy!\n✨ Supports very long prompts (300+ characters) for detailed descriptions!\nTip: Be specific about lighting, pose, setting, mood, and artistic style for best results with LoRA enhancements.",
                        lines=8,
                        elem_classes=["prompt-box"]
                    )
                    
                    negative_prompt_input = gr.Textbox(
                        label="❌ Negative Prompt (LoRA-Optimized)",
                        placeholder="Things you don't want in the image...\nExample: low quality, blurry, distorted, bad anatomy\n\n✨ Auto-enhanced with LoRA-optimized negative terms if left empty",
                        lines=3,
                        elem_classes=["prompt-box"]
                    )
                
                with gr.Column(scale=1):
                    style_input = gr.Radio(
                        label="🎭 Style Preset (LoRA-Enhanced)",
                        choices=list(STYLE_PRESETS.keys()),
                        value="Realistic",
                        elem_classes=["style-selector"]
                    )
                    
                    # 🔥 新增：模型选择器
                    with gr.Group(elem_classes=["model-selector"]):
                        gr.Markdown("### 🏭 Model Selection")
                        model_input = gr.Radio(
                            label="Choose Base Model",
                            choices=list(NSFW_MODELS.keys()),
                            value="PornMasterPro",
                            info="Select the base NSFW model for generation"
                        )
                    
                    # 🔥 新增：LoRA配置选项卡
                    with gr.Group(elem_classes=["lora-selector"]):
                        gr.Markdown("### 🎨 LoRA Configuration")
                        
                        with gr.Tabs():
                            # SDXL Lightning LoRA配置
                            with gr.TabItem("⚡ Lightning"):
                                lightning_enabled = gr.Checkbox(
                                    label="Enable SDXL Lightning",
                                    value=LORA_MODELS["sdxl_lightning"]["enabled"],
                                    info="4-step acceleration LoRA"
                                )
                                lightning_scale = gr.Slider(
                                    label="Lightning Scale",
                                    minimum=0.0,
                                    maximum=2.0,
                                    value=LORA_MODELS["sdxl_lightning"]["scale"],
                                    step=0.1,
                                    info="Recommended: 1.0"
                                )
                            
                            # Quality Enhancer LoRA配置
                            with gr.TabItem("✨ Quality"):
                                quality_enabled = gr.Checkbox(
                                    label="Enable Quality Enhancer",
                                    value=LORA_MODELS["quality_enhancer"]["enabled"],
                                    info="Realism and quality enhancement"
                                )
                                quality_scale = gr.Slider(
                                    label="Quality Scale",
                                    minimum=0.0,
                                    maximum=2.0,
                                    value=LORA_MODELS["quality_enhancer"]["scale"],
                                    step=0.1,
                                    info="Recommended: 0.8"
                                )
                            
                            # Detail Enhancer LoRA配置
                            with gr.TabItem("🔍 Detail"):
                                detail_enabled = gr.Checkbox(
                                    label="Enable Detail Enhancer",
                                    value=LORA_MODELS["detail_enhancer"]["enabled"],
                                    info="Detail enhancement (experimental)"
                                )
                                detail_scale = gr.Slider(
                                    label="Detail Scale",
                                    minimum=0.0,
                                    maximum=2.0,
                                    value=LORA_MODELS["detail_enhancer"]["scale"],
                                    step=0.1,
                                    info="Recommended: 0.6"
                                )
                    
                    # 参数控制
                    with gr.Group(elem_classes=["controls-section"]):
                        gr.Markdown("### ⚙️ Generation Parameters")
                        
                        steps_input = gr.Slider(
                            label="🔥 Steps (Auto-adjusted for Lightning)",
                            minimum=4,
                            maximum=50,
                            value=28,
                            step=1,
                            elem_classes=["slider-container"],
                            info="Lightning LoRA will override to 4 steps"
                        )
                        
                        cfg_input = gr.Slider(
                            label="🎯 CFG Scale (Auto-adjusted for Lightning)",
                            minimum=0.0,
                            maximum=15.0,
                            value=7.0,
                            step=0.1,
                            elem_classes=["slider-container"],
                            info="Lightning LoRA will override to 0.0"
                        )
            
            # 生成按钮
            generate_button = gr.Button(
                "🚀 Generate LoRA-Enhanced NSFW Image",
                elem_classes=["generate-btn"],
                variant="primary"
            )
            
            # 输出区域
            with gr.Row():
                image_output = gr.Image(
                    label="Generated LoRA-Enhanced NSFW Image",
                    elem_classes=["image-output"],
                    show_label=False,
                    container=True,
                    width=768,
                    height=1024
                )
            
            # 状态信息
            status_output = gr.Textbox(
                label="LoRA-Enhanced Status",
                interactive=False,
                elem_classes=["status-box"],
                show_label=False
            )
        
        # 🔥 LoRA设置处理函数
        def collect_lora_settings(lightning_enabled, lightning_scale, quality_enabled, quality_scale, detail_enabled, detail_scale):
            return {
                "sdxl_lightning": {
                    "enabled": lightning_enabled,
                    "scale": lightning_scale
                },
                "quality_enhancer": {
                    "enabled": quality_enabled,
                    "scale": quality_scale
                },
                "detail_enhancer": {
                    "enabled": detail_enabled,
                    "scale": detail_scale
                }
            }
        
        # 🔥 模型ID获取函数
        def get_model_id(model_name):
            return NSFW_MODELS.get(model_name, list(NSFW_MODELS.values())[0])
        
        # 绑定事件
        generate_button.click(
            fn=lambda prompt, model_name, style, neg_prompt, steps, cfg, l_en, l_sc, q_en, q_sc, d_en, d_sc: generate_image(
                prompt=prompt,
                selected_model=get_model_id(model_name),
                style=style,
                lora_settings=collect_lora_settings(l_en, l_sc, q_en, q_sc, d_en, d_sc),
                negative_prompt=neg_prompt,
                steps=steps,
                cfg_scale=cfg
            ),
            inputs=[
                prompt_input, model_input, style_input, negative_prompt_input, 
                steps_input, cfg_input,
                lightning_enabled, lightning_scale,
                quality_enabled, quality_scale,
                detail_enabled, detail_scale
            ],
            outputs=[image_output, status_output],
            show_progress=True
        )
        
        # 支持Enter键触发
        prompt_input.submit(
            fn=lambda prompt, model_name, style, neg_prompt, steps, cfg, l_en, l_sc, q_en, q_sc, d_en, d_sc: generate_image(
                prompt=prompt,
                selected_model=get_model_id(model_name),
                style=style,
                lora_settings=collect_lora_settings(l_en, l_sc, q_en, q_sc, d_en, d_sc),
                negative_prompt=neg_prompt,
                steps=steps,
                cfg_scale=cfg
            ),
            inputs=[
                prompt_input, model_input, style_input, negative_prompt_input, 
                steps_input, cfg_input,
                lightning_enabled, lightning_scale,
                quality_enabled, quality_scale,
                detail_enabled, detail_scale
            ],
            outputs=[image_output, status_output],
            show_progress=True
        )
        
        # 启动时显示LoRA增强欢迎信息
        interface.load(
            fn=lambda: (None, """🎨 Optimized LoRA-Enhanced NSFW Generator ready! 
✨ Available Features:
• 6 High-quality NSFW base models
• Configurable LoRA adapters with parameter tuning
• SDXL-Lightning acceleration (4 steps, CFG=0)
• Quality enhancement LoRAs
• Long prompt support (300+ characters)
• Auto-parameter adjustment for optimal results

🎯 Select your model, configure LoRA settings, and enter your detailed prompt!
💡 Lightning LoRA will automatically optimize speed with 4-step generation.
🔧 Each LoRA can be individually enabled and fine-tuned for your needs."""),
            outputs=[image_output, status_output]
        )
    
    return interface

# ===== 启动应用 =====
if __name__ == "__main__":
    print("🎨 Starting Optimized LoRA-Enhanced NSFW Image Generator...")
    
    # 显示LoRA增强信息
    print("🔧 LoRA Enhancement Features:")
    print("   • Configurable LoRA adapters with individual settings")
    print("   • SDXL-Lightning acceleration support")
    print("   • Auto-parameter adjustment for optimal performance")
    print("   • Model selection with 6 NSFW-optimized options")
    print("   • Enhanced negative prompting")
    print("   • Long prompt processing with LoRA support")
    print("   • Reduced GPU usage through optimized loading")
    
    # 显示配置的LoRA模型
    print("🎨 Available LoRA Models:")
    for name, config in LORA_MODELS.items():
        print(f"   • {name}: {config['description']} (default scale: {config['scale']})")
    
    # 显示可用模型
    print("📋 Available base models:")
    for i, (name, model_id) in enumerate(NSFW_MODELS.items(), 1):
        print(f"   {i}. {name}: {model_id}")
    
    # 检查关键组件
    print(f"🔧 Spaces GPU: {'✅ Available' if SPACES_AVAILABLE else '❌ Not Available'}")
    print(f"🔧 Compel Library: {'✅ Available' if COMPEL_AVAILABLE else '❌ Not Available'}")
    print(f"🔧 CUDA: {'✅ Available' if torch.cuda.is_available() else '❌ Not Available'}")
    
    # 创建并启动界面
    app = create_interface()
    app.queue(max_size=15, default_concurrency_limit=2)
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=False
    )