Update app.py

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 import torch
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
+import gc
 
 colors.steel_blue = colors.Color(
     name="steel_blue",
@@ -36,186 +37,208 @@ class SteelBlueTheme(Soft):
 steel_blue_theme = SteelBlueTheme()
 
 print("=" * 50)
-print("🎨 Style2Paints - Enhanced Multi-Model Edition")
+print("🎨 Style2Paints - Uncensored Line Art Colorization & Text-to-Image")
 print("=" * 50)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-from diffusers import StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline, ControlNetModel, DDIMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, DDIMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline, EulerDiscreteScheduler
 from controlnet_aux import LineartDetector, LineartAnimeDetector
 
-# Model configurations
+# ===== 模型配置 =====
 AVAILABLE_MODELS = {
-    "Anything V3 (Anime)": {
-        "repo": "Linaqruf/anything-v3.0",
+    "LineArt Colorization": ["Anything V3 (ControlNet)"],
+    "Text-to-Image": [
+        "Linaqruf/anything-v3.0",
+        "digiplay/ChikMix_V3",
+        "digiplay/chilloutmix_NiPrunedFp16Fix",
+        "LyliaEngine/Pony_Diffusion_V6_XL",
+        "wootwoot/abyssorangemix3-popupparade-fp16",
+        "John6666/wai-nsfw-illustrious-v80-sdxl"
+    ]
+}
+
+MODEL_CONFIGS = {
+    "Linaqruf/anything-v3.0": {
         "type": "sd15",
-        "description": "Original anime model"
+        "description": "Anything V3 - 全能模型",
+        "default_resolution": (512, 768)
     },
-    "ChikMix V3": {
-        "repo": "digiplay/ChikMix_V3",
+    "digiplay/ChikMix_V3": {
         "type": "sd15",
-        "description": "High quality anime/realistic mix"
+        "description": "ChikMix V3 - 高质量动漫模型",
+        "default_resolution": (512, 768)
     },
-    "ChilloutMix": {
-        "repo": "digiplay/chilloutmix_NiPrunedFp16Fix",
+    "digiplay/chilloutmix_NiPrunedFp16Fix": {
         "type": "sd15",
-        "description": "Popular realistic/anime hybrid"
+        "description": "ChilloutMix - 真人风格",
+        "default_resolution": (512, 768)
     },
-    "Pony Diffusion V6 XL": {
-        "repo": "LyliaEngine/Pony_Diffusion_V6_XL",
+    "LyliaEngine/Pony_Diffusion_V6_XL": {
         "type": "sdxl",
-        "description": "SDXL anime model"
+        "description": "Pony Diffusion V6 XL - SDXL动漫模型",
+        "default_resolution": (1024, 1024)
     },
-    "AbyssOrangeMix3": {
-        "repo": "wootwoot/abyssorangemix3-popupparade-fp16",
+    "wootwoot/abyssorangemix3-popupparade-fp16": {
         "type": "sd15",
-        "description": "Vibrant anime style"
+        "description": "AbyssOrangeMix3 - 色彩鲜艳",
+        "default_resolution": (512, 768)
     },
-    "WAI-NSFW Illustrious XL": {
-        "repo": "John6666/wai-nsfw-illustrious-v80-sdxl",
+    "John6666/wai-nsfw-illustrious-v80-sdxl": {
         "type": "sdxl",
-        "description": "SDXL NSFW-focused model"
+        "description": "WAI NSFW Illustrious - SDXL成人内容优化",
+        "default_resolution": (1024, 1024)
     }
 }
 
-# Global variables for models
-loaded_models = {}
-controlnet_standard = None
-controlnet_anime = None
+# ===== 全局模型变量 =====
+pipe_standard = None
+pipe_anime = None
 lineart_detector = None
 lineart_anime_detector = None
-txt2img_pipe = None
-current_txt2img_model = None
-
-def load_controlnet_models():
-    """Load ControlNet models"""
-    global controlnet_standard, controlnet_anime, lineart_detector, lineart_anime_detector
-    
-    print("📦 Loading ControlNet models...")
-    
-    try:
-        controlnet_standard = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_lineart",
-            torch_dtype=dtype
-        ).to(device)
-        print("✅ Standard ControlNet loaded")
-    except Exception as e:
-        print(f"❌ Failed to load standard ControlNet: {e}")
-    
-    try:
-        controlnet_anime = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15s2_lineart_anime",
-            torch_dtype=dtype
-        ).to(device)
-        print("✅ Anime ControlNet loaded")
-    except Exception as e:
-        print(f"❌ Failed to load anime ControlNet: {e}")
-    
-    try:
-        lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
-        lineart_anime_detector = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
-        print("✅ Line art detectors loaded")
-    except Exception as e:
-        print(f"❌ Failed to load detectors: {e}")
+current_t2i_model = None
+current_t2i_pipe = None
 
-def load_model_for_controlnet(model_name, lineart_type):
-    """Load a specific model with ControlNet"""
-    global loaded_models
-    
-    cache_key = f"{model_name}_{lineart_type}"
+def load_text_to_image_model(model_name, progress=gr.Progress()):
+    """动态加载文本到图像模型"""
+    global current_t2i_model, current_t2i_pipe
     
-    if cache_key in loaded_models:
-        print(f"♻️ Using cached model: {cache_key}")
-        return loaded_models[cache_key]
-    
-    model_info = AVAILABLE_MODELS[model_name]
-    controlnet = controlnet_anime if lineart_type == "Anime" else controlnet_standard
-    
-    if controlnet is None:
-        raise gr.Error("ControlNet model not loaded")
-    
-    print(f"📦 Loading {model_name}...")
+    if model_name == current_t2i_model and current_t2i_pipe is not None:
+        print(f"✅ 模型 {model_name} 已加载，跳过重新加载")
+        return True
     
     try:
-        if model_info["type"] == "sdxl":
-            # SDXL models don't support SD1.5 ControlNet currently
-            raise gr.Error("SDXL models are not yet supported for line art colorization. Use SD1.5 models.")
+        # 清理之前的模型
+        if current_t2i_pipe is not None:
+            del current_t2i_pipe
+            current_t2i_pipe = None
+            current_t2i_model = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
         
-        pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            model_info["repo"],
-            controlnet=controlnet,
-            torch_dtype=dtype,
-            safety_checker=None,
-            requires_safety_checker=False
-        ).to(device)
-        
-        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-        
-        if device.type == "cuda":
-            pipe.enable_model_cpu_offload()
-            try:
-                pipe.enable_xformers_memory_efficient_attention()
-            except:
-                pass
-        pipe.enable_attention_slicing()
+        print(f"🔄 正在加载模型: {model_name}")
+        progress(0.3, desc=f"正在加载 {model_name}")
         
-        loaded_models[cache_key] = pipe
-        print(f"✅ {model_name} loaded successfully")
-        return pipe
+        model_config = MODEL_CONFIGS.get(model_name, {})
+        model_type = model_config.get("type", "sd15")
         
-    except Exception as e:
-        print(f"❌ Error loading {model_name}: {e}")
-        raise gr.Error(f"Failed to load {model_name}: {str(e)}")
-
-def load_txt2img_model(model_name):
-    """Load model for text-to-image"""
-    global txt2img_pipe, current_txt2img_model
-    
-    if current_txt2img_model == model_name and txt2img_pipe is not None:
-        print(f"♻️ Using cached txt2img model: {model_name}")
-        return txt2img_pipe
-    
-    model_info = AVAILABLE_MODELS[model_name]
-    print(f"📦 Loading {model_name} for text-to-image...")
-    
-    try:
-        if model_info["type"] == "sdxl":
+        if model_type == "sdxl":
+            # SDXL 模型
             pipe = StableDiffusionXLPipeline.from_pretrained(
-                model_info["repo"],
+                model_name,
                 torch_dtype=dtype,
                 safety_checker=None,
-                requires_safety_checker=False
+                requires_safety_checker=False,
+                use_safetensors=True,
+                variant="fp16" if dtype == torch.float16 else None
             ).to(device)
+            
+            # SDXL 推荐使用 Euler scheduler
+            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+            
         else:
+            # SD1.5 模型
             pipe = StableDiffusionPipeline.from_pretrained(
-                model_info["repo"],
+                model_name,
                 torch_dtype=dtype,
                 safety_checker=None,
-                requires_safety_checker=False
+                requires_safety_checker=False,
+                use_safetensors=True,
+                variant="fp16" if dtype == torch.float16 else None
             ).to(device)
+            
+            # SD1.5 使用 DDIM
+            pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
         
-        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-        
+        # 优化设置
         if device.type == "cuda":
             pipe.enable_model_cpu_offload()
             try:
                 pipe.enable_xformers_memory_efficient_attention()
             except:
-                pass
-        pipe.enable_attention_slicing()
+                print("⚠️ XFormers 不可用，跳过")
+            pipe.enable_attention_slicing()
+        
+        current_t2i_model = model_name
+        current_t2i_pipe = pipe
+        
+        print(f"✅ 模型 {model_name} 加载成功！")
+        progress(1.0, desc="模型加载完成")
+        return True
+        
+    except Exception as e:
+        import traceback
+        print(f"❌ 加载模型失败: {str(e)}")
+        print(f"详细错误: {traceback.format_exc()}")
+        return False
+
+def load_lineart_models():
+    """加载线稿着色模型"""
+    global pipe_standard, pipe_anime, lineart_detector, lineart_anime_detector
+    
+    try:
+        print("🔄 加载线稿着色模型...")
+        
+        # Load STANDARD lineart ControlNet
+        print("📦 加载标准线稿模型...")
+        controlnet_standard = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_lineart",
+            torch_dtype=dtype
+        ).to(device)
+        
+        # Load ANIME lineart ControlNet
+        print("📦 加载动漫线稿模型...")
+        controlnet_anime = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15s2_lineart_anime",
+            torch_dtype=dtype
+        ).to(device)
+        
+        # 使用 Anything V3 作为基础模型
+        print("📦 加载基础模型 (Anything V3)...")
+        pipe_standard = StableDiffusionControlNetPipeline.from_pretrained(
+            "Linaqruf/anything-v3.0",
+            controlnet=controlnet_standard,
+            torch_dtype=dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(device)
+        
+        pipe_anime = StableDiffusionControlNetPipeline.from_pretrained(
+            "Linaqruf/anything-v3.0",
+            controlnet=controlnet_anime,
+            torch_dtype=dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(device)
+        
+        # 配置两个管道
+        for pipe in [pipe_standard, pipe_anime]:
+            pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+            
+            if device.type == "cuda":
+                pipe.enable_model_cpu_offload()
+                try:
+                    pipe.enable_xformers_memory_efficient_attention()
+                except:
+                    pass
+            pipe.enable_attention_slicing()
+        
+        # 加载线稿检测器
+        print("📦 加载线稿检测器...")
+        lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
+        lineart_anime_detector = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
         
-        txt2img_pipe = pipe
-        current_txt2img_model = model_name
-        print(f"✅ {model_name} loaded for txt2img")
-        return pipe
+        print("✅ 线稿着色模型加载成功！")
+        return True
         
     except Exception as e:
-        print(f"❌ Error loading {model_name}: {e}")
-        raise gr.Error(f"Failed to load {model_name}: {str(e)}")
+        print(f"❌ 加载线稿模型失败: {e}")
+        return False
 
-# Load ControlNet models at startup
-load_controlnet_models()
+# 加载线稿着色模型
+load_lineart_models()
 
 COLOR_STYLES = {
     "Anime Style": "anime, masterpiece, best quality, highly detailed, vibrant colors",
@@ -239,7 +262,7 @@ CONTENT_TEMPLATES = {
 }
 
 def is_already_lineart(image):
-    """Check if image is already line art"""
+    """检查图像是否已经是线稿"""
     if isinstance(image, Image.Image):
         image = np.array(image)
     
@@ -250,15 +273,15 @@ def is_already_lineart(image):
     return black_white_ratio > 0.7 or unique_vals < 30
 
 def extract_lineart(image, lineart_type="Standard", skip_if_lineart=True):
-    """Extract line art from image"""
+    """从图像中提取线稿"""
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     
     if skip_if_lineart and is_already_lineart(image):
-        print("✅ Already line art, skipping extraction")
+        print("✅ 已经是线稿，跳过提取")
         return image.convert('RGB')
     
-    print(f"🔄 Extracting line art ({lineart_type})...")
+    print(f"🔄 提取线稿 ({lineart_type})...")
     
     if lineart_type == "Anime" and lineart_anime_detector is not None:
         lineart = lineart_anime_detector(image, detect_resolution=512, image_resolution=512)
@@ -272,7 +295,6 @@ def extract_lineart(image, lineart_type="Standard", skip_if_lineart=True):
 
 def colorize_lineart(
     sketch_image,
-    model_name,
     lineart_type,
     content_type,
     style,
@@ -286,14 +308,22 @@ def colorize_lineart(
     controlnet_strength,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """Colorize with explicit content support"""
+    """线稿着色"""
     if sketch_image is None:
-        raise gr.Error("Please upload a sketch/line art image")
+        raise gr.Error("请上传线稿图像")
     
-    # Load the selected model
-    pipe = load_model_for_controlnet(model_name, lineart_type)
+    # 根据线稿类型选择管道
+    if lineart_type == "Anime" and pipe_anime is not None:
+        pipe = pipe_anime
+        print("🎨 使用动漫线稿模型")
+    else:
+        pipe = pipe_standard
+        print("🎨 使用标准线稿模型")
     
-    # Convert all numeric inputs to proper types
+    if pipe is None:
+        raise gr.Error("模型未加载")
+    
+    # 转换数值输入
     seed = int(seed)
     guidance_scale = float(guidance_scale)
     num_steps = int(num_steps)
@@ -305,7 +335,7 @@ def colorize_lineart(
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Convert and resize
+    # 转换和调整大小
     if isinstance(sketch_image, np.ndarray):
         sketch_image = Image.fromarray(sketch_image)
     
@@ -324,30 +354,30 @@ def colorize_lineart(
         new_height = (new_height // 8) * 8
         sketch_image = sketch_image.resize((new_width, new_height), Image.LANCZOS)
     
-    # Extract lineart with selected type
+    # 提取线稿
     lineart = extract_lineart(sketch_image, lineart_type=lineart_type, skip_if_lineart=True)
     
-    # Build prompt intelligently
+    # 构建提示词
     prompt_parts = []
     
-    # Content template
+    # 内容模板
     content_template = CONTENT_TEMPLATES.get(content_type, "")
     if content_template:
         prompt_parts.append(content_template)
     
-    # Custom prompt (most important)
+    # 自定义提示词
     if custom_prompt.strip():
         prompt_parts.append(custom_prompt.strip())
     
-    # Quality tags
+    # 质量标签
     if quality_tags:
         prompt_parts.append(quality_tags)
     
-    # Style
+    # 风格
     style_prompt = COLOR_STYLES.get(style, COLOR_STYLES["Anime Style"])
     prompt_parts.append(style_prompt)
     
-    # NSFW level tags
+    # NSFW 级别标签
     if nsfw_level == "Safe":
         nsfw_tags = "sfw, safe for work"
     elif nsfw_level == "Suggestive":
@@ -363,18 +393,18 @@ def colorize_lineart(
     
     full_prompt = ", ".join(prompt_parts)
     
-    # Negative prompt
+    # 负面提示词
     negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry, artist name, black and white, monochrome"
     
     if nsfw_level in ["Moderate", "Explicit"]:
         negative_prompt = negative_prompt.replace("nsfw, ", "")
     
-    print(f"🎨 Model: {model_name}")
-    print(f"🎨 Prompt: {full_prompt}")
-    print(f"🎛️ ControlNet Strength: {controlnet_strength}")
+    print(f"🎨 提示词: {full_prompt}")
+    print(f"🎛️ ControlNet 强度: {controlnet_strength}")
+    print(f"🖼️ 线稿类型: {lineart_type}")
     
     try:
-        progress(0.3, desc="Generating colors...")
+        progress(0.3, desc="正在生成颜色...")
         
         result = pipe(
             prompt=full_prompt,
@@ -393,10 +423,10 @@ def colorize_lineart(
         
     except Exception as e:
         import traceback
-        print(f"❌ Full error: {traceback.format_exc()}")
-        raise gr.Error(f"Error: {str(e)}")
+        print(f"❌ 完整错误: {traceback.format_exc()}")
+        raise gr.Error(f"错误: {str(e)}")
 
-def generate_txt2img(
+def generate_text_to_image(
     model_name,
     content_type,
     style,
@@ -411,12 +441,12 @@ def generate_txt2img(
     height,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """Generate image from text only"""
-    
-    # Load the selected model
-    pipe = load_txt2img_model(model_name)
+    """文本到图像生成"""
+    # 加载模型
+    if not load_text_to_image_model(model_name, progress):
+        raise gr.Error(f"无法加载模型: {model_name}")
     
-    # Convert all numeric inputs to proper types
+    # 转换数值输入
     seed = int(seed)
     guidance_scale = float(guidance_scale)
     num_steps = int(num_steps)
@@ -429,27 +459,27 @@ def generate_txt2img(
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Build prompt
+    # 构建提示词
     prompt_parts = []
     
-    # Content template
+    # 内容模板
     content_template = CONTENT_TEMPLATES.get(content_type, "")
     if content_template:
         prompt_parts.append(content_template)
     
-    # Custom prompt
+    # 自定义提示词
     if custom_prompt.strip():
         prompt_parts.append(custom_prompt.strip())
     
-    # Quality tags
+    # 质量标签
     if quality_tags:
         prompt_parts.append(quality_tags)
     
-    # Style
+    # 风格
     style_prompt = COLOR_STYLES.get(style, COLOR_STYLES["Anime Style"])
     prompt_parts.append(style_prompt)
     
-    # NSFW level tags
+    # NSFW 级别标签
     if nsfw_level == "Safe":
         nsfw_tags = "sfw, safe for work"
     elif nsfw_level == "Suggestive":
@@ -465,28 +495,27 @@ def generate_txt2img(
     
     full_prompt = ", ".join(prompt_parts)
     
-    # Negative prompt
-    negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry, artist name"
+    # 负面提示词
+    negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry, artist name, black and white, monochrome"
     
     if nsfw_level in ["Moderate", "Explicit"]:
         negative_prompt = negative_prompt.replace("nsfw, ", "")
     
-    print(f"🎨 Model: {model_name}")
-    print(f"🎨 Prompt: {full_prompt}")
+    print(f"���� 提示词: {full_prompt}")
+    print(f"📐 分辨率: {width}x{height}")
+    print(f"🎛️ 引导尺度: {guidance_scale}")
+    print(f"🔄 步数: {num_steps}")
     
     try:
-        progress(0.3, desc="Generating image...")
-        
-        # Check if SDXL model
-        model_info = AVAILABLE_MODELS[model_name]
+        progress(0.5, desc="正在生成图像...")
         
-        result = pipe(
+        result = current_t2i_pipe(
             prompt=full_prompt,
             negative_prompt=negative_prompt,
-            num_inference_steps=num_steps,
-            guidance_scale=guidance_scale,
             width=width,
             height=height,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
             generator=generator,
         ).images[0]
         
@@ -497,8 +526,21 @@ def generate_txt2img(
         
     except Exception as e:
         import traceback
-        print(f"❌ Full error: {traceback.format_exc()}")
-        raise gr.Error(f"Error: {str(e)}")
+        print(f"❌ 完整错误: {traceback.format_exc()}")
+        raise gr.Error(f"错误: {str(e)}")
+
+def update_resolution_from_model(model_name):
+    """根据选择的模型更新推荐分辨率"""
+    config = MODEL_CONFIGS.get(model_name, {})
+    default_res = config.get("default_resolution", (512, 768))
+    description = config.get("description", "通用模型")
+    
+    width, height = default_res
+    return (
+        gr.update(value=width, minimum=256, maximum=2048, step=8),
+        gr.update(value=height, minimum=256, maximum=2048, step=8),
+        gr.update(value=f"📊 推荐分辨率: {width}x{height} ({description})")
+    )
 
 css="""
 #col-container {
@@ -544,212 +586,411 @@ css="""
     font-size: 0.9em;
     margin: 5px;
 }
+.tab-buttons {
+    margin-bottom: 20px;
+}
+.tab-nav {
+    border-bottom: 2px solid #e0e0e0;
+}
 """
 
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# 🎨 Style2Paints - Multi-Model Edition", elem_id="main-title")
-        gr.Markdown("### ✨ Professional AI Art Generation with Multiple Models")
+        gr.Markdown("# 🎨 Style2Paints - 全能图像生成", elem_id="main-title")
+        gr.Markdown("### ✨ 专业线稿着色与文本生成图像")
         
         gr.HTML("""
         <div class="warning-box">
-            <strong>⚠️ Content Warning:</strong> This tool supports generation of all content types including NSFW/adult material. 
-            Use responsibly and ensure compliance with local laws. Users must be 18+ for explicit content.
+            <strong>⚠️ 内容警告:</strong> 此工具支持所有类型内容的生成，包括 NSFW/成人内容。
+            请负责任地使用，并确保符合当地法律。生成成人内容必须年满18岁。
         </div>
         """)
         
-        with gr.Tabs():
-            # Tab 1: Line Art Colorization
-            with gr.Tab("🖊️ Line Art Colorization"):
-                gr.HTML("""
-                <div class="feature-box">
-                    <h3>✨ Colorize Your Line Art</h3>
-                    <ul style="color:white; font-size:1.1em;">
-                        <li>🎨 <strong>6 Anime Models</strong> - Choose your favorite style</li>
-                        <li>🖊️ <strong>Dual Lineart Detection</strong> - Standard & Anime optimized</li>
-                        <li>🎭 <strong>Smart Prompt Building</strong> - Easy content creation</li>
-                        <li>🎚️ <strong>NSFW Support</strong> - All content levels supported</li>
-                    </ul>
-                </div>
-                """)
-                
+        gr.HTML("""
+        <div class="feature-box">
+            <h3>✨ 核心功能</h3>
+            <ul style="color:white; font-size:1.1em;">
+                <li>🎨 <strong>双线稿模型</strong> - 标准和动漫专用线稿检测</li>
+                <li>🖼️ <strong>文本生成图像</strong> - 从文本描述直接生成图像</li>
+                <li>🎭 <strong>多模型支持</strong> - 6种不同风格的模型可选</li>
+                <li>📝 <strong>内容模板</strong> - 常见场景的预设提示词</li>
+                <li>🎚️ <strong>NSFW级别控制</strong> - 精确的内容级别控制</li>
+                <li>⚡ <strong>智能模型加载</strong> - 按需加载，节省显存</li>
+            </ul>
+            <div style="margin-top:15px;">
+                <span class="model-badge">6种文本生成模型</span>
+                <span class="model-badge">2种线稿模型</span>
+            </div>
+        </div>
+        """)
+        
+        with gr.Tabs() as tabs:
+            # ===== 标签页 1: 线稿着色 =====
+            with gr.TabItem("🎨 线稿着色"):
                 with gr.Row():
                     with gr.Column(scale=1):
                         input_image = gr.Image(
-                            label="📤 Upload Line Art",
+                            label="📤 上传线稿",
                             type="pil",
                             height=400
                         )
                         
-                        # Get SD1.5 models only for ControlNet
-                        sd15_models = [name for name, info in AVAILABLE_MODELS.items() if info["type"] == "sd15"]
-                        
-                        model_selector = gr.Dropdown(
-                            choices=sd15_models,
-                            label="🤖 Select Model",
-                            value="Anything V3 (Anime)",
-                            info="Choose your preferred model (SD1.5 only for line art)"
-                        )
+                        gr.Markdown("### 🎨 内容设置")
                         
                         lineart_type = gr.Radio(
                             choices=["Standard", "Anime"],
-                            label="🖊️ Lineart Model",
+                            label="🖊️ 线稿模型",
                             value="Anime",
-                            info="Anime works better for anime/manga"
+                            info="动漫模型更适合动漫/漫画风格"
                         )
                         
                         content_type = gr.Dropdown(
                             choices=list(CONTENT_TEMPLATES.keys()),
-                            label="📋 Content Template",
-                            value="Character Portrait"
+                            label="📋 内容模板",
+                            value="Character Portrait",
+                            info="提示词起点"
                         )
                         
                         custom_prompt = gr.Textbox(
-                            label="✍️ Detailed Description",
-                            placeholder="hair color, outfit, pose, features...",
-                            lines=3
+                            label="✍️ 详细描述 (重要)",
+                            placeholder="描述您想要的内容：发色、服装、姿势、背景、身体特征等",
+                            lines=3,
+                            info="请具体描述！这是最重要的字段。"
                         )
                         
+                        gr.HTML("""
+                        <div class="info-box">
+                        <strong>💡 提示词示例:</strong><br>
+                        • "金发，蓝眼睛，女仆装，丰满"<br>
+                        • "红色马尾辫，校服，短裙，过膝袜"<br>
+                        • "白发，猫耳，裸体，躺在床上"<br>
+                        • "两个女孩，接吻，亲密，卧室"
+                        </div>
+                        """)
+                        
                         with gr.Row():
                             style = gr.Dropdown(
                                 choices=list(COLOR_STYLES.keys()),
-                                label="🎨 Style",
+                                label="🎨 颜色风格",
                                 value="Anime Style"
                             )
                             
                             nsfw_level = gr.Dropdown(
                                 choices=["Safe", "Suggestive", "Mild", "Moderate", "Explicit"],
-                                label="🔞 Content Level",
-                                value="Moderate"
+                                label="🔞 内容级别",
+                                value="Moderate",
+                                info="内容明确程度"
                             )
                         
                         quality_tags = gr.Textbox(
-                            label="⭐ Quality Tags",
+                            label="⭐ 质量标签 (可选)",
+                            placeholder="masterpiece, best quality, highly detailed",
                             value="masterpiece, best quality, highly detailed"
                         )
                         
-                        colorize_button = gr.Button("✨ Colorize!", variant="primary", size="lg")
+                        colorize_button = gr.Button("✨ 开始着色!", variant="primary", size="lg")
                     
                     with gr.Column(scale=2):
                         with gr.Row():
-                            lineart_output = gr.Image(label="🖊️ Line Art", type="pil", height=380)
-                            colorized_output = gr.Image(label="🎨 Result", type="pil", height=380)
+                            lineart_output = gr.Image(
+                                label="🖊️ 提取的线稿",
+                                type="pil",
+                                height=380
+                            )
+                            output_image = gr.Image(
+                                label="🎨 着色结果",
+                                type="pil",
+                                height=380
+                            )
                         
-                        colorized_prompt = gr.Textbox(
-                            label="📝 Generated Prompt",
+                        generated_prompt = gr.Textbox(
+                            label="📝 生成的提示词",
                             lines=3,
                             interactive=False,
                             show_copy_button=True
                         )
                         
-                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                        with gr.Accordion("⚙️ 高级设置", open=True):
                             with gr.Row():
-                                c_seed = gr.Slider(label="🎲 Seed", minimum=0, maximum=2**32-1, step=1, value=42)
-                                c_random = gr.Checkbox(label="🔀 Random", value=True)
+                                seed = gr.Slider(
+                                    label="🎲 种子",
+                                    minimum=0,
+                                    maximum=2**32-1,
+                                    step=1,
+                                    value=42
+                                )
+                                randomize_seed = gr.Checkbox(
+                                    label="🔀 随机种子",
+                                    value=True
+                                )
                             
                             with gr.Row():
-                                c_guidance = gr.Slider(label="💬 Guidance", minimum=5.0, maximum=15.0, step=0.5, value=8.0)
-                                c_steps = gr.Slider(label="🔢 Steps", minimum=10, maximum=30, step=5, value=20)
+                                guidance_scale = gr.Slider(
+                                    label="💬 引导尺度",
+                                    minimum=5.0,
+                                    maximum=15.0,
+                                    step=0.5,
+                                    value=8.0,
+                                    info="7-9 推荐用于 NSFW"
+                                )
+                                
+                                num_steps = gr.Slider(
+                                    label="🔢 步数",
+                                    minimum=10,
+                                    maximum=30,
+                                    step=5,
+                                    value=20,
+                                    info="20 是良好平衡"
+                                )
                             
-                            c_strength = gr.Slider(
-                                label="🎛️ Line Preservation",
+                            controlnet_strength = gr.Slider(
+                                label="🎛️ 线稿保留强度",
                                 minimum=0.5,
                                 maximum=1.5,
                                 step=0.1,
-                                value=1.0
+                                value=1.0,
+                                info="严格遵循线稿的程度"
                             )
                 
                 colorize_button.click(
                     fn=colorize_lineart,
                     inputs=[
-                        input_image, model_selector, lineart_type, content_type, style,
-                        custom_prompt, quality_tags, nsfw_level, c_seed, c_random,
-                        c_guidance, c_steps, c_strength
+                        input_image, lineart_type, content_type, style, custom_prompt, quality_tags, nsfw_level,
+                        seed, randomize_seed, guidance_scale, num_steps, controlnet_strength
                     ],
-                    outputs=[colorized_output, lineart_output, c_seed, colorized_prompt]
+                    outputs=[output_image, lineart_output, seed, generated_prompt]
                 )
             
-            # Tab 2: Text to Image
-            with gr.Tab("✍️ Text to Image"):
-                gr.HTML("""
-                <div class="feature-box">
-                    <h3>✨ Generate Images from Text</h3>
-                    <ul style="color:white; font-size:1.1em;">
-                        <li>🎨 <strong>6 Models Including SDXL</strong> - Maximum variety</li>
-                        <li>✍️ <strong>Pure Text Generation</strong> - No line art needed</li>
-                        <li>📐 <strong>Flexible Sizes</strong> - From 512x512 to 1024x1024</li>
-                        <li>🎭 <strong>All Content Types</strong> - SFW to NSFW</li>
-                    </ul>
-                </div>
-                """)
-                
+            # ===== 标签页 2: 文本生成图像 =====
+            with gr.TabItem("🖼️ 文本生成图像"):
                 with gr.Row():
                     with gr.Column(scale=1):
-                        t2i_model = gr.Dropdown(
-                            choices=list(AVAILABLE_MODELS.keys()),
-                            label="🤖 Select Model",
-                            value="Anything V3 (Anime)",
-                            info="All models available for txt2img"
+                        gr.Markdown("### 🤖 模型选择")
+                        
+                        model_selector = gr.Dropdown(
+                            choices=AVAILABLE_MODELS["Text-to-Image"],
+                            label="🎯 选择模型",
+                            value="Linaqruf/anything-v3.0",
+                            info="选择要使用的生成模型"
+                        )
+                        
+                        model_info = gr.Textbox(
+                            label="📊 模型信息",
+                            value="📊 推荐分辨率: 512x768 (Anything V3 - 全能模型)",
+                            interactive=False
+                        )
+                        
+                        load_model_btn = gr.Button("🔄 加载模型", variant="secondary")
+                        model_status = gr.Textbox(
+                            label="✅ 状态",
+                            value="✅ 模型已就绪",
+                            interactive=False
                         )
                         
-                        t2i_content = gr.Dropdown(
+                        gr.Markdown("### 🎨 内容设置")
+                        
+                        t2i_content_type = gr.Dropdown(
                             choices=list(CONTENT_TEMPLATES.keys()),
-                            label="📋 Content Template",
-                            value="Character Portrait"
+                            label="📋 内容模板",
+                            value="Character Portrait",
+                            info="提示词起点"
                         )
                         
-                        t2i_prompt = gr.Textbox(
-                            label="✍️ Your Description",
-                            placeholder="Describe what you want to see...",
-                            lines=4
+                        t2i_custom_prompt = gr.Textbox(
+                            label="✍️ 详细描述 (重要)",
+                            placeholder="详细描述您想要生成的图像：角色特征、服装、姿势、场景等",
+                            lines=3,
+                            info="描述越详细，生成效果越好"
                         )
                         
                         gr.HTML("""
                         <div class="info-box">
-                        <strong>💡 Example Prompts:</strong><br>
-                        • "girl with long silver hair, blue eyes, wearing kimono, cherry blossoms"<br>
-                        • "cyberpunk girl, neon lights, futuristic city, night scene"<br>
-                        • "maid outfit, cat ears, blushing, bedroom"
+                        <strong>💡 提示词示例:</strong><br>
+                        • "美丽的女孩，金色长发，蓝色眼睛，穿着白色连衣裙，站在花园里"<br>
+                        • "性感的女战士，红色铠甲，手持长剑，动态姿势，战场背景"<br>
+                        • "两个女孩在咖啡馆约会，温馨的氛围，详细的面部表情"<br>
+                        • "幻想风格的女精灵，尖耳朵，魔法光效，森林背景"
                         </div>
                         """)
                         
                         with gr.Row():
                             t2i_style = gr.Dropdown(
                                 choices=list(COLOR_STYLES.keys()),
-                                label="🎨 Style",
+                                label="🎨 艺术风格",
                                 value="Anime Style"
                             )
                             
-                            t2i_nsfw = gr.Dropdown(
+                            t2i_nsfw_level = gr.Dropdown(
                                 choices=["Safe", "Suggestive", "Mild", "Moderate", "Explicit"],
-                                label="🔞 Content Level",
-                                value="Safe"
+                                label="🔞 内容级别",
+                                value="Moderate",
+                                info="内容明确程度"
                             )
                         
-                        t2i_quality = gr.Textbox(
-                            label="⭐ Quality Tags",
+                        t2i_quality_tags = gr.Textbox(
+                            label="⭐ 质量标签 (可选)",
+                            placeholder="masterpiece, best quality, highly detailed",
                             value="masterpiece, best quality, highly detailed"
                         )
                         
-                        with gr.Row():
-                            t2i_width = gr.Slider(label="📏 Width", minimum=512, maximum=1024, step=64, value=512)
-                            t2i_height = gr.Slider(label="📏 Height", minimum=512, maximum=1024, step=64, value=512)
-                        
-                        generate_button = gr.Button("🎨 Generate!", variant="primary", size="lg")
+                        generate_button = gr.Button("✨ 生成图像!", variant="primary", size="lg")
                     
                     with gr.Column(scale=2):
-                        t2i_output = gr.Image(label="🎨 Generated Image", type="pil", height=600)
+                        t2i_output_image = gr.Image(
+                            label="🖼️ 生成的图像",
+                            type="pil",
+                            height=500
+                        )
                         
-                        t2i_prompt_output = gr.Textbox(
-                            label="📝 Full Prompt Used",
+                        t2i_generated_prompt = gr.Textbox(
+                            label="📝 生成的提示词",
                             lines=3,
                             interactive=False,
                             show_copy_button=True
                         )
                         
-                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                        with gr.Accordion("⚙️ 高级设置", open=True):
+                            with gr.Row():
+                                t2i_seed = gr.Slider(
+                                    label="🎲 种子",
+                                    minimum=0,
+                                    maximum=2**32-1,
+                                    step=1,
+                                    value=42
+                                )
+                                t2i_randomize_seed = gr.Checkbox(
+                                    label="🔀 随机种子",
+                                    value=True
+                                )
+                            
+                            with gr.Row():
+                                t2i_guidance_scale = gr.Slider(
+                                    label="💬 引导尺度",
+                                    minimum=5.0,
+                                    maximum=15.0,
+                                    step=0.5,
+                                    value=7.5,
+                                    info="控制提示词影响力"
+                                )
+                                
+                                t2i_num_steps = gr.Slider(
+                                    label="🔢 生成步数",
+                                    minimum=10,
+                                    maximum=50,
+                                    step=5,
+                                    value=30,
+                                    info="步数越多质量越高但越慢"
+                                )
+                            
                             with gr.Row():
-                                t2i_seed = gr.Slider(label="🎲 Seed", minimum=0, maximum=
+                                t2i_width = gr.Slider(
+                                    label="📏 宽度",
+                                    minimum=256,
+                                    maximum=2048,
+                                    step=8,
+                                    value=512,
+                                    info="图像宽度"
+                                )
+                                
+                                t2i_height = gr.Slider(
+                                    label="📐 高度",
+                                    minimum=256,
+                                    maximum=2048,
+                                    step=8,
+                                    value=768,
+                                    info="图像高度"
+                                )
+                
+                # 事件处理
+                model_selector.change(
+                    fn=update_resolution_from_model,
+                    inputs=[model_selector],
+                    outputs=[t2i_width, t2i_height, model_info]
+                )
+                
+                load_model_btn.click(
+                    fn=lambda model_name: (
+                        load_text_to_image_model(model_name, gr.Progress()) and 
+                        gr.update(value=f"✅ {model_name} 加载成功")
+                    ),
+                    inputs=[model_selector],
+                    outputs=[model_status]
+                )
+                
+                generate_button.click(
+                    fn=generate_text_to_image,
+                    inputs=[
+                        model_selector,
+                        t2i_content_type,
+                        t2i_style,
+                        t2i_custom_prompt,
+                        t2i_quality_tags,
+                        t2i_nsfw_level,
+                        t2i_seed,
+                        t2i_randomize_seed,
+                        t2i_guidance_scale,
+                        t2i_num_steps,
+                        t2i_width,
+                        t2i_height
+                    ],
+                    outputs=[t2i_output_image, t2i_seed, t2i_generated_prompt]
+                )
+    
+    gr.Markdown("""
+    ---
+    ## 📚 快速开始指南
+    
+    ### 🆕 **新功能: 文本生成图像**
+    
+    此版本新增文本生成图像功能，支持6种不同的模型：
+    
+    #### 🤖 **可用模型:**
+    
+    1. **Anything V3** (`Linaqruf/anything-v3.0`) - 全能动漫模型
+    2. **ChikMix V3** (`digiplay/ChikMix_V3`) - 高质量动漫模型
+    3. **ChilloutMix** (`digiplay/chilloutmix_NiPrunedFp16Fix`) - 真人风格模型
+    4. **Pony Diffusion V6 XL** (`LyliaEngine/Pony_Diffusion_V6_XL`) - SDXL动漫模型 (高分辨率)
+    5. **AbyssOrangeMix3** (`wootwoot/abyssorangemix3-popupparade-fp16`) - 色彩鲜艳的动漫模型
+    6. **WAI NSFW Illustrious** (`John6666/wai-nsfw-illustrious-v80-sdxl`) - SDXL成人内容优化模型
+    
+    ### 🎨 **线稿着色模型**
+    
+    线稿着色功能提供两种线稿模型：
+    - **标准线稿** (`control_v11p_sd15_lineart`) - 适合一般艺术作品
+    - **动漫线稿** (`control_v11p_sd15s2_lineart_anime`) - 专为动漫/漫画风格优化 ✨
+    
+    ### ✅ **如何使用**
+    
+    #### **线稿着色:**
+    1. 上传您的线稿（黑白线条在白底上效果最好）
+    2. 选择线稿模型 - 动漫风格使用"Anime"模型
+    3. 选择内容模板作为起点
+    4. 编写详细描述 - 具体说明颜色、特征、服装等
+    5. 设置NSFW级别以匹配您的内容
+    6. 点击"开始着色!"
+    
+    #### **文本生成图像:**
+    1. 选择您想要使用的模型
+    2. 点击"加载模型"按钮（首次使用或切换模型时需要）
+    3. 编写详细描述您想要生成的图像
+    4. 调整分辨率和生成参数
+    5. 点击"生成图像!"
+    
+    ### 💡 **最佳实践提示**
+    
+    - **模型选择**: SDXL模型需要更多显存但生成质量更高
+    - **详细描述**: 描述越详细，生成效果越好
+    - **分辨率设置**: SDXL模型推荐使用1024x1024，SD1.5模型推荐512x768
+    - **显存管理**: 模型按需加载，切换模型时会自动清理之前的模型
+    
+    ---
+    
+    <div style="text-align:center; color:#666; padding:20px;">
+        <strong>🔞 负责任使用</strong><br>
+        此工具用于艺术创作目的。用户必须年满18岁。<br>
+        请尊重版权、同意和当地法律。<br>
+        <em>由 Stable Diffusion + ControlNet + 多种生成模型驱动</em>
+    </div>
     """)
 
 demo.launch(server_name="0.0.0.0", server_port=7860, share=False)