Update app.py

app.py

@@ -36,106 +36,186 @@ class SteelBlueTheme(Soft):
 steel_blue_theme = SteelBlueTheme()
 
 print("=" * 50)
-print("🎨 Style2Paints - Multi-Model Edition")
+print("🎨 Style2Paints - Enhanced Multi-Model Edition")
 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, ControlNetModel, DDIMScheduler
+from diffusers import StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline, ControlNetModel, DDIMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline
 from controlnet_aux import LineartDetector, LineartAnimeDetector
 
-# Global variables for pipes
-current_pipe = None
-current_base_model = None
+# Model configurations
+AVAILABLE_MODELS = {
+    "Anything V3 (Anime)": {
+        "repo": "Linaqruf/anything-v3.0",
+        "type": "sd15",
+        "description": "Original anime model"
+    },
+    "ChikMix V3": {
+        "repo": "digiplay/ChikMix_V3",
+        "type": "sd15",
+        "description": "High quality anime/realistic mix"
+    },
+    "ChilloutMix": {
+        "repo": "digiplay/chilloutmix_NiPrunedFp16Fix",
+        "type": "sd15",
+        "description": "Popular realistic/anime hybrid"
+    },
+    "Pony Diffusion V6 XL": {
+        "repo": "LyliaEngine/Pony_Diffusion_V6_XL",
+        "type": "sdxl",
+        "description": "SDXL anime model"
+    },
+    "AbyssOrangeMix3": {
+        "repo": "wootwoot/abyssorangemix3-popupparade-fp16",
+        "type": "sd15",
+        "description": "Vibrant anime style"
+    },
+    "WAI-NSFW Illustrious XL": {
+        "repo": "John6666/wai-nsfw-illustrious-v80-sdxl",
+        "type": "sdxl",
+        "description": "SDXL NSFW-focused model"
+    }
+}
+
+# Global variables for models
+loaded_models = {}
 controlnet_standard = None
 controlnet_anime = None
 lineart_detector = None
 lineart_anime_detector = None
+txt2img_pipe = None
+current_txt2img_model = None
 
-BASE_MODELS = {
-    "Anything V3": "Linaqruf/anything-v3.0",
-    "AbyssOrangeMix": "WarriorMama777/OrangeMixs"
-}
-
-try:
-    print("🔄 Loading ControlNet Models...")
-    
-    # Load STANDARD lineart ControlNet
-    print("📦 Loading standard lineart ControlNet...")
-    controlnet_standard = ControlNetModel.from_pretrained(
-        "lllyasviel/control_v11p_sd15_lineart",
-        torch_dtype=dtype
-    ).to(device)
-    
-    # Load ANIME lineart ControlNet
-    print("📦 Loading anime lineart ControlNet...")
-    controlnet_anime = ControlNetModel.from_pretrained(
-        "lllyasviel/control_v11p_sd15s2_lineart_anime",
-        torch_dtype=dtype
-    ).to(device)
-    
-    # Load detectors
-    print("📦 Loading line art detectors...")
-    lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
-    lineart_anime_detector = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
-    
-    print("✅ ControlNet models loaded successfully!")
-    
-except Exception as e:
-    print(f"❌ Error loading ControlNet: {e}")
+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}")
 
-def load_base_model(base_model_name, lineart_type):
-    """Load base model on demand"""
-    global current_pipe, current_base_model
-    
-    # Check if we already have the right model loaded
-    if current_pipe is not None and current_base_model == (base_model_name, lineart_type):
-        print(f"✅ Model already loaded: {base_model_name} ({lineart_type})")
-        return current_pipe
-    
-    # Clear existing pipe
-    if current_pipe is not None:
-        print("🗑️ Clearing previous model...")
-        del current_pipe
-        if device.type == "cuda":
-            torch.cuda.empty_cache()
-        current_pipe = None
+def load_model_for_controlnet(model_name, lineart_type):
+    """Load a specific model with ControlNet"""
+    global loaded_models
     
-    # Select ControlNet
+    cache_key = f"{model_name}_{lineart_type}"
+    
+    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
-    model_path = BASE_MODELS[base_model_name]
     
-    print(f"📦 Loading {base_model_name} with {lineart_type} ControlNet...")
+    if controlnet is None:
+        raise gr.Error("ControlNet model not loaded")
+    
+    print(f"📦 Loading {model_name}...")
     
     try:
-        current_pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            model_path,
+        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.")
+        
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            model_info["repo"],
             controlnet=controlnet,
             torch_dtype=dtype,
             safety_checker=None,
             requires_safety_checker=False
         ).to(device)
         
-        # Configure pipe
-        current_pipe.scheduler = DDIMScheduler.from_config(current_pipe.scheduler.config)
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
         
         if device.type == "cuda":
-            current_pipe.enable_model_cpu_offload()
+            pipe.enable_model_cpu_offload()
             try:
-                current_pipe.enable_xformers_memory_efficient_attention()
+                pipe.enable_xformers_memory_efficient_attention()
             except:
                 pass
-        current_pipe.enable_attention_slicing()
+        pipe.enable_attention_slicing()
         
-        current_base_model = (base_model_name, lineart_type)
-        print(f"✅ {base_model_name} loaded successfully!")
+        loaded_models[cache_key] = pipe
+        print(f"✅ {model_name} loaded successfully")
+        return pipe
         
-        return current_pipe
+    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":
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                model_info["repo"],
+                torch_dtype=dtype,
+                safety_checker=None,
+                requires_safety_checker=False
+            ).to(device)
+        else:
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_info["repo"],
+                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()
+        
+        txt2img_pipe = pipe
+        current_txt2img_model = model_name
+        print(f"✅ {model_name} loaded for txt2img")
+        return pipe
         
     except Exception as e:
-        print(f"❌ Error loading {base_model_name}: {e}")
-        raise gr.Error(f"Failed to load {base_model_name}: {str(e)}")
+        print(f"❌ Error loading {model_name}: {e}")
+        raise gr.Error(f"Failed to load {model_name}: {str(e)}")
+
+# Load ControlNet models at startup
+load_controlnet_models()
 
 COLOR_STYLES = {
     "Anime Style": "anime, masterpiece, best quality, highly detailed, vibrant colors",
@@ -192,7 +272,7 @@ def extract_lineart(image, lineart_type="Standard", skip_if_lineart=True):
 
 def colorize_lineart(
     sketch_image,
-    base_model,
+    model_name,
     lineart_type,
     content_type,
     style,
@@ -206,16 +286,12 @@ def colorize_lineart(
     controlnet_strength,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """Colorize with model selection"""
+    """Colorize with explicit content support"""
     if sketch_image is None:
         raise gr.Error("Please upload a sketch/line art image")
     
     # Load the selected model
-    progress(0.1, desc=f"Loading {base_model}...")
-    pipe = load_base_model(base_model, lineart_type)
-    
-    if pipe is None:
-        raise gr.Error("Failed to load model")
+    pipe = load_model_for_controlnet(model_name, lineart_type)
     
     # Convert all numeric inputs to proper types
     seed = int(seed)
@@ -248,7 +324,7 @@ def colorize_lineart(
         new_height = (new_height // 8) * 8
         sketch_image = sketch_image.resize((new_width, new_height), Image.LANCZOS)
     
-    progress(0.2, desc="Extracting lineart...")
+    # Extract lineart with selected type
     lineart = extract_lineart(sketch_image, lineart_type=lineart_type, skip_if_lineart=True)
     
     # Build prompt intelligently
@@ -293,8 +369,7 @@ def colorize_lineart(
     if nsfw_level in ["Moderate", "Explicit"]:
         negative_prompt = negative_prompt.replace("nsfw, ", "")
     
-    print(f"🎨 Base Model: {base_model}")
-    print(f"🖼️ Lineart Type: {lineart_type}")
+    print(f"🎨 Model: {model_name}")
     print(f"🎨 Prompt: {full_prompt}")
     print(f"🎛️ ControlNet Strength: {controlnet_strength}")
     
@@ -314,8 +389,111 @@ def colorize_lineart(
         if device.type == "cuda":
             torch.cuda.empty_cache()
         
-        model_info = f"Model: {base_model} | Lineart: {lineart_type}"
-        return result, lineart, seed, f"{model_info}\n{full_prompt}"
+        return result, lineart, seed, full_prompt
+        
+    except Exception as e:
+        import traceback
+        print(f"❌ Full error: {traceback.format_exc()}")
+        raise gr.Error(f"Error: {str(e)}")
+
+def generate_txt2img(
+    model_name,
+    content_type,
+    style,
+    custom_prompt,
+    quality_tags,
+    nsfw_level,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    num_steps,
+    width,
+    height,
+    progress=gr.Progress(track_tqdm=True)
+):
+    """Generate image from text only"""
+    
+    # Load the selected model
+    pipe = load_txt2img_model(model_name)
+    
+    # Convert all numeric inputs to proper types
+    seed = int(seed)
+    guidance_scale = float(guidance_scale)
+    num_steps = int(num_steps)
+    width = int(width)
+    height = int(height)
+    
+    if randomize_seed:
+        import random
+        seed = random.randint(0, 2**32-1)
+    
+    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
+    if nsfw_level == "Safe":
+        nsfw_tags = "sfw, safe for work"
+    elif nsfw_level == "Suggestive":
+        nsfw_tags = "suggestive, slightly revealing"
+    elif nsfw_level == "Mild":
+        nsfw_tags = "nsfw, ecchi, revealing clothing"
+    elif nsfw_level == "Moderate":
+        nsfw_tags = "nsfw, nude, explicit"
+    else:  # Explicit
+        nsfw_tags = "nsfw, explicit, uncensored"
+    
+    prompt_parts.append(nsfw_tags)
+    
+    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"
+    
+    if nsfw_level in ["Moderate", "Explicit"]:
+        negative_prompt = negative_prompt.replace("nsfw, ", "")
+    
+    print(f"🎨 Model: {model_name}")
+    print(f"🎨 Prompt: {full_prompt}")
+    
+    try:
+        progress(0.3, desc="Generating image...")
+        
+        # Check if SDXL model
+        model_info = AVAILABLE_MODELS[model_name]
+        
+        result = pipe(
+            prompt=full_prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+        
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+        
+        return result, seed, full_prompt
         
     except Exception as e:
         import traceback
@@ -371,221 +549,207 @@ css="""
 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 Line Art Colorization with Model Selection")
+        gr.Markdown("### ✨ Professional AI Art Generation with Multiple Models")
         
         gr.HTML("""
         <div class="warning-box">
-            <strong>⚠️ Content Warning:</strong> This tool supports colorization of all content types including NSFW/adult material. 
+            <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.
         </div>
         """)
         
-        gr.HTML("""
-        <div class="feature-box">
-            <h3>✨ Key Features</h3>
-            <ul style="color:white; font-size:1.1em;">
-                <li>🎨 <strong>Multiple Base Models</strong> - Choose between Anything V3 and AbyssOrangeMix</li>
-                <li>🚀 <strong>Lazy Loading</strong> - Models load only when needed to save memory</li>
-                <li>🎭 <strong>Dual ControlNet</strong> - Standard and Anime-specific lineart detection</li>
-                <li>📝 <strong>Content Templates</strong> - Pre-built prompts for common scenarios</li>
-                <li>🎚️ <strong>NSFW Level Control</strong> - Explicit tags for better accuracy</li>
-            </ul>
-            <div style="margin-top:15px;">
-                <span class="model-badge">Anything V3</span>
-                <span class="model-badge">AbyssOrangeMix</span>
-                <span class="model-badge">Standard Lineart</span>
-                <span class="model-badge">Anime Lineart</span>
-            </div>
-        </div>
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_image = gr.Image(
-                    label="📤 Upload Line Art",
-                    type="pil",
-                    height=400
-                )
-                
-                gr.Markdown("### 🎨 Model & Content Settings")
-                
-                base_model = gr.Radio(
-                    choices=list(BASE_MODELS.keys()),
-                    label="🤖 Base Model",
-                    value="Anything V3",
-                    info="Model loads when you click Colorize"
-                )
-                
-                lineart_type = gr.Radio(
-                    choices=["Standard", "Anime"],
-                    label="🖊️ Lineart ControlNet",
-                    value="Anime",
-                    info="Anime model works better for anime/manga style"
-                )
-                
-                content_type = gr.Dropdown(
-                    choices=list(CONTENT_TEMPLATES.keys()),
-                    label="📋 Content Template",
-                    value="Character Portrait",
-                    info="Starting point for your prompt"
-                )
-                
-                custom_prompt = gr.Textbox(
-                    label="✍️ Detailed Description (IMPORTANT)",
-                    placeholder="Describe what you want: hair color, outfit, pose, background, body features, etc.",
-                    lines=3,
-                    info="Be specific! This is the most important field."
-                )
-                
+        with gr.Tabs():
+            # Tab 1: Line Art Colorization
+            with gr.Tab("🖊️ Line Art Colorization"):
                 gr.HTML("""
-                <div class="info-box">
-                <strong>💡 Base Model Characteristics:</strong><br>
-                • <strong>Anything V3</strong>: Clean anime style, good for general use<br>
-                • <strong>AbyssOrangeMix</strong>: More detailed, painterly style
+                <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>
                 """)
                 
                 with gr.Row():
-                    style = gr.Dropdown(
-                        choices=list(COLOR_STYLES.keys()),
-                        label="🎨 Color Style",
-                        value="Anime Style"
-                    )
+                    with gr.Column(scale=1):
+                        input_image = gr.Image(
+                            label="📤 Upload Line Art",
+                            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)"
+                        )
+                        
+                        lineart_type = gr.Radio(
+                            choices=["Standard", "Anime"],
+                            label="🖊️ Lineart Model",
+                            value="Anime",
+                            info="Anime works better for anime/manga"
+                        )
+                        
+                        content_type = gr.Dropdown(
+                            choices=list(CONTENT_TEMPLATES.keys()),
+                            label="📋 Content Template",
+                            value="Character Portrait"
+                        )
+                        
+                        custom_prompt = gr.Textbox(
+                            label="✍️ Detailed Description",
+                            placeholder="hair color, outfit, pose, features...",
+                            lines=3
+                        )
+                        
+                        with gr.Row():
+                            style = gr.Dropdown(
+                                choices=list(COLOR_STYLES.keys()),
+                                label="🎨 Style",
+                                value="Anime Style"
+                            )
+                            
+                            nsfw_level = gr.Dropdown(
+                                choices=["Safe", "Suggestive", "Mild", "Moderate", "Explicit"],
+                                label="🔞 Content Level",
+                                value="Moderate"
+                            )
+                        
+                        quality_tags = gr.Textbox(
+                            label="⭐ Quality Tags",
+                            value="masterpiece, best quality, highly detailed"
+                        )
+                        
+                        colorize_button = gr.Button("✨ Colorize!", variant="primary", size="lg")
                     
-                    nsfw_level = gr.Dropdown(
-                        choices=["Safe", "Suggestive", "Mild", "Moderate", "Explicit"],
-                        label="🔞 Content Level",
-                        value="Moderate",
-                        info="How explicit?"
-                    )
+                    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)
+                        
+                        colorized_prompt = gr.Textbox(
+                            label="📝 Generated Prompt",
+                            lines=3,
+                            interactive=False,
+                            show_copy_button=True
+                        )
+                        
+                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                            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)
+                            
+                            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)
+                            
+                            c_strength = gr.Slider(
+                                label="🎛️ Line Preservation",
+                                minimum=0.5,
+                                maximum=1.5,
+                                step=0.1,
+                                value=1.0
+                            )
                 
-                quality_tags = gr.Textbox(
-                    label="⭐ Quality Tags (Optional)",
-                    placeholder="masterpiece, best quality, highly detailed",
-                    value="masterpiece, best quality, highly detailed"
+                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
+                    ],
+                    outputs=[colorized_output, lineart_output, c_seed, colorized_prompt]
                 )
-                
-                colorize_button = gr.Button("✨ Colorize!", variant="primary", size="lg")
             
-            with gr.Column(scale=2):
-                with gr.Row():
-                    lineart_output = gr.Image(
-                        label="🖊️ Line Art",
-                        type="pil",
-                        height=380
-                    )
-                    output_image = gr.Image(
-                        label="🎨 Colorized Result",
-                        type="pil",
-                        height=380
-                    )
-                
-                generated_prompt = gr.Textbox(
-                    label="📝 Model & Prompt Info",
-                    lines=3,
-                    interactive=False,
-                    show_copy_button=True
-                )
+            # 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>
+                """)
                 
-                with gr.Accordion("⚙️ Advanced Settings", open=True):
-                    with gr.Row():
-                        seed = gr.Slider(
-                            label="🎲 Seed",
-                            minimum=0,
-                            maximum=2**32-1,
-                            step=1,
-                            value=42
+                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"
                         )
-                        randomize_seed = gr.Checkbox(
-                            label="🔀 Random Seed",
-                            value=True
+                        
+                        t2i_content = gr.Dropdown(
+                            choices=list(CONTENT_TEMPLATES.keys()),
+                            label="📋 Content Template",
+                            value="Character Portrait"
                         )
-                    
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="💬 Guidance Scale",
-                            minimum=5.0,
-                            maximum=15.0,
-                            step=0.5,
-                            value=8.0,
-                            info="7-9 recommended for NSFW"
+                        
+                        t2i_prompt = gr.Textbox(
+                            label="✍️ Your Description",
+                            placeholder="Describe what you want to see...",
+                            lines=4
                         )
                         
-                        num_steps = gr.Slider(
-                            label="🔢 Steps",
-                            minimum=10,
-                            maximum=30,
-                            step=5,
-                            value=20,
-                            info="20 is good balance"
+                        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"
+                        </div>
+                        """)
+                        
+                        with gr.Row():
+                            t2i_style = gr.Dropdown(
+                                choices=list(COLOR_STYLES.keys()),
+                                label="🎨 Style",
+                                value="Anime Style"
+                            )
+                            
+                            t2i_nsfw = gr.Dropdown(
+                                choices=["Safe", "Suggestive", "Mild", "Moderate", "Explicit"],
+                                label="🔞 Content Level",
+                                value="Safe"
+                            )
+                        
+                        t2i_quality = gr.Textbox(
+                            label="⭐ Quality Tags",
+                            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")
                     
-                    controlnet_strength = gr.Slider(
-                        label="🎛️ Line Preservation",
-                        minimum=0.5,
-                        maximum=1.5,
-                        step=0.1,
-                        value=1.0,
-                        info="How strictly to follow lines"
-                    )
-        
-        colorize_button.click(
-            fn=colorize_lineart,
-            inputs=[
-                input_image, base_model, lineart_type, content_type, style, custom_prompt, 
-                quality_tags, nsfw_level, seed, randomize_seed, guidance_scale, num_steps, 
-                controlnet_strength
-            ],
-            outputs=[output_image, lineart_output, seed, generated_prompt]
-        )
-    
-    gr.Markdown("""
-    ---
-    ## 📚 Quick Start Guide
-    
-    ### 🆕 **Model Selection Feature**
-    
-    Choose between two popular anime models:
-    - **Anything V3** - Clean, versatile anime style. Good all-rounder.
-    - **AbyssOrangeMix** - More detailed and painterly. Better for complex scenes.
-    
-    **💡 Memory Efficient**: Models load only when you click "Colorize", saving RAM!
-    
-    ### ✅ **How to Use**
-    
-    1. **Upload your line art** (black lines on white background works best)
-    2. **Select base model** - Try both to see which style you prefer!
-    3. **Choose lineart ControlNet** - Use "Anime" for anime/manga style
-    4. **Choose a content template** to get started
-    5. **Write a detailed description** - be specific about colors, features, clothing, pose
-    6. **Set the NSFW level** to match your content
-    7. **Click Colorize!** - Model will load automatically
-    
-    ### 🎨 **Model Comparison**
-    
-    | Feature | Anything V3 | AbyssOrangeMix |
-    |---------|-------------|----------------|
-    | Style | Clean anime | Painterly, detailed |
-    | Best for | General use | Complex scenes |
-    | Detail level | Standard | High |
-    | Colors | Vibrant | Rich, layered |
-    
-    ### 💡 **Tips for Best Results**
-    
-    - **First time**: Models download on first use (may take a few minutes)
-    - **Model switching**: Automatic - just select and click Colorize
-    - **Be detailed in prompts**: "blonde long hair, blue eyes, red dress" beats "girl"
-    - **Match NSFW level**: Use "Explicit" for nude, "Safe" for SFW
-    - **Clean line art is key**: Black lines on white background = best results
-    
-    ---
-    
-    <div style="text-align:center; color:#666; padding:20px;">
-        <strong>🔞 Responsible Use</strong><br>
-        This tool is for artistic purposes. Users must be 18+.<br>
-        Respect copyright, consent, and local laws.<br>
-        <em>Powered by Stable Diffusion + ControlNet (Multi-Model Support)</em>
-    </div>
+                    with gr.Column(scale=2):
+                        t2i_output = gr.Image(label="🎨 Generated Image", type="pil", height=600)
+                        
+                        t2i_prompt_output = gr.Textbox(
+                            label="📝 Full Prompt Used",
+                            lines=3,
+                            interactive=False,
+                            show_copy_button=True
+                        )
+                        
+                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                            with gr.Row():
+                                t2i_seed = gr.Slider(label="🎲 Seed", minimum=0, maximum=
     """)
 
 demo.launch(server_name="0.0.0.0", server_port=7860, share=False)