Update app.py

app.py

@@ -124,10 +124,8 @@ CHINESE_MODELS = [
     "AI-ModelScope/stable-diffusion-v1-5-chinese"
 ]
 
-ALL_MODELS = SD15_MODELS + SDXL_MODELS + CHINESE_MODELS
-
-# ControlNet models
-CONTROLNET_MODELS = {
+# ControlNet models for SD1.5
+CONTROLNET_MODELS_SD15 = {
     "lineart": "lllyasviel/control_v11p_sd15_lineart",
     "lineart_anime": "lllyasviel/control_v11p_sd15s2_lineart_anime",
     "canny": "lllyasviel/control_v11p_sd15_canny",
@@ -142,8 +140,8 @@ CONTROLNET_MODELS = {
     "tile": "lllyasviel/control_v11f1e_sd15_tile"
 }
 
-# SDXL ControlNet models
-SDXL_CONTROLNET_MODELS = {
+# ControlNet models for SDXL
+CONTROLNET_MODELS_SDXL = {
     "canny_sdxl": "diffusers/controlnet-canny-sdxl-1.0",
     "depth_sdxl": "diffusers/controlnet-depth-sdxl-1.0",
     "openpose_sdxl": "thibaud/controlnet-openpose-sdxl-1.0"
@@ -249,45 +247,43 @@ def load_detector(detector_type: str):
         print(f"❌ Error loading {detector_type} detector: {e}")
         return None
 
-def get_controlnet_model(controlnet_type: str):
+def get_controlnet_model(controlnet_type: str, is_sdxl: bool = False):
     """Get ControlNet model based on type"""
-    if controlnet_type in CONTROLNET_MODELS:
-        return CONTROLNET_MODELS[controlnet_type]
-    elif controlnet_type in SDXL_CONTROLNET_MODELS:
-        return SDXL_CONTROLNET_MODELS[controlnet_type]
+    if is_sdxl:
+        return CONTROLNET_MODELS_SDXL[controlnet_type]
     else:
-        raise ValueError(f"Unknown ControlNet type: {controlnet_type}")
+        return CONTROLNET_MODELS_SD15[controlnet_type]
 
-def prepare_condition_image(image, controlnet_type):
+def prepare_condition_image(image, controlnet_type, is_sdxl=False):
     """Prepare condition image for ControlNet"""
     if controlnet_type in ["lineart", "lineart_anime"]:
         detector = load_detector("lineart_anime" if controlnet_type == "lineart_anime" else "lineart")
         if detector:
-            result = detector(image, detect_resolution=512, image_resolution=512)
+            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
-    elif controlnet_type == "canny":
+    elif "canny" in controlnet_type:
         detector = load_detector("canny")
         if detector:
-            result = detector(image, detect_resolution=512, image_resolution=512)
+            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
-    elif controlnet_type == "depth":
+    elif "depth" in controlnet_type:
         detector = load_detector("depth")
         if detector:
-            result = detector(image, detect_resolution=512, image_resolution=512)
+            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
     elif controlnet_type == "normal":
         detector = load_detector("normal")
         if detector:
-            result = detector(image, detect_resolution=512, image_resolution=512)
+            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
-    elif controlnet_type == "openpose":
+    elif "openpose" in controlnet_type:
         detector = load_detector("openpose")
         if detector:
-            result = detector(image, detect_resolution=512, image_resolution=512)
+            result = detector(image, detect_resolution=512 if not is_sdxl else 1024, image_resolution=512 if not is_sdxl else 1024)
             return Image.fromarray(result) if isinstance(result, np.ndarray) else result
     
     return image
@@ -315,32 +311,30 @@ def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model:
     print(f"📥 Loading ControlNet pipeline for model: {model_name}, type: {controlnet_type}")
     
     try:
-        if is_sdxl_model(model_name):
-            if controlnet_type in SDXL_CONTROLNET_MODELS:
-                controlnet_model_name = get_controlnet_model(controlnet_type)
-                controlnet = ControlNetModel.from_pretrained(
-                    controlnet_model_name, 
-                    torch_dtype=dtype
-                ).to(device)
-                
-                pipe = StableDiffusionXLPipeline.from_pretrained(
-                    model_name, 
-                    controlnet=controlnet, 
-                    torch_dtype=dtype,
-                    safety_checker=None,
-                    requires_safety_checker=False,
-                    use_safetensors=True,
-                    variant="fp16" if dtype == torch.float16 else None
-                ).to(device)
-            else:
-                raise ValueError(f"SDXL model only supports: {list(SDXL_CONTROLNET_MODELS.keys())}")
-        else:
-            controlnet_model_name = get_controlnet_model(controlnet_type)
-            controlnet = ControlNetModel.from_pretrained(
-                controlnet_model_name, 
-                torch_dtype=dtype
+        is_sdxl = is_sdxl_model(model_name)
+        
+        if is_sdxl and controlnet_type not in CONTROLNET_MODELS_SDXL:
+            raise ValueError(f"SDXL model only supports: {list(CONTROLNET_MODELS_SDXL.keys())}")
+        elif not is_sdxl and controlnet_type not in CONTROLNET_MODELS_SD15:
+            raise ValueError(f"SD1.5 model only supports: {list(CONTROLNET_MODELS_SD15.keys())}")
+        
+        controlnet_model_name = get_controlnet_model(controlnet_type, is_sdxl)
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_model_name, 
+            torch_dtype=dtype
+        ).to(device)
+        
+        if is_sdxl:
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                model_name, 
+                controlnet=controlnet, 
+                torch_dtype=dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
+                use_safetensors=True,
+                variant="fp16" if dtype == torch.float16 else None
             ).to(device)
-            
+        else:
             pipe = StableDiffusionControlNetPipeline.from_pretrained(
                 model_name, 
                 controlnet=controlnet, 
@@ -365,13 +359,11 @@ def get_pipeline(model_name: str, controlnet_type: str = "lineart", lora_model:
         if lora_model and lora_model != "None":
             print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
             try:
-                # ตรวจสอบว่า lora_model อยู่ใน dictionary หรือไม่
                 if lora_model in LORA_MODELS:
                     lora_path = LORA_MODELS[lora_model]
                     pipe.load_lora_weights(lora_path)
                     pipe.fuse_lora(lora_scale=lora_weight)
                 else:
-                    # ถ้าเป็น full path โดยตรง
                     pipe.load_lora_weights(lora_model)
                     pipe.fuse_lora(lora_scale=lora_weight)
             except Exception as e:
@@ -494,13 +486,11 @@ def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float =
         if lora_model and lora_model != "None":
             print(f"🔄 Applying LoRA: {lora_model} with weight: {lora_weight}")
             try:
-                # ตรวจสอบว่า lora_model อยู่ใน dictionary หรือไม่
                 if lora_model in LORA_MODELS:
                     lora_path = LORA_MODELS[lora_model]
                     CURRENT_T2I_PIPE.load_lora_weights(lora_path)
                     CURRENT_T2I_PIPE.fuse_lora(lora_scale=lora_weight)
                 else:
-                    # ถ้าเป็น full path โดยตรง
                     CURRENT_T2I_PIPE.load_lora_weights(lora_model)
                     CURRENT_T2I_PIPE.fuse_lora(lora_scale=lora_weight)
             except Exception as e:
@@ -537,21 +527,17 @@ def load_t2i_model(model_name: str, lora_model: str = None, lora_weight: float =
         CURRENT_T2I_MODEL = None
         raise
 
-def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
-             prompt, negative_prompt, seed, steps, scale, cn_weight):
+def colorize_sd15(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
+                  prompt, negative_prompt, seed, steps, scale, cn_weight):
+    """Colorize function for SD1.5 models"""
     try:
-        if is_sdxl_model(base_model) and controlnet_type not in SDXL_CONTROLNET_MODELS:
+        if base_model not in SD15_MODELS:
             error_img = Image.new('RGB', (512, 512), color='red')
-            error_msg_img = Image.new('RGB', (512, 512), color='yellow')
-            from PIL import ImageDraw, ImageFont
-            draw = ImageDraw.Draw(error_msg_img)
-            try:
-                font = ImageFont.truetype("arial.ttf", 20)
-            except:
-                font = ImageFont.load_default()
-            draw.text((50, 200), f"SDXL model only supports:", fill="black", font=font)
-            draw.text((50, 230), f"{', '.join(SDXL_CONTROLNET_MODELS.keys())}", fill="black", font=font)
-            return error_img, error_msg_img
+            return error_img, Image.new('RGB', (512, 512), color='gray')
+        
+        if controlnet_type not in CONTROLNET_MODELS_SD15:
+            error_img = Image.new('RGB', (512, 512), color='red')
+            return error_img, Image.new('RGB', (512, 512), color='gray')
         
         pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight, vae_model)
         
@@ -560,7 +546,7 @@ def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_m
             status_msg += f" + {lora_model}"
         print(status_msg)
         
-        condition_img = prepare_condition_image(sketch, controlnet_type)
+        condition_img = prepare_condition_image(sketch, controlnet_type, is_sdxl=False)
         
         gen = torch.Generator(device=device).manual_seed(int(seed))
         
@@ -582,22 +568,117 @@ def colorize(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_m
         
         return out, condition_img
     except Exception as e:
-        print(f"❌ Error in colorize: {e}")
+        print(f"❌ Error in colorize_sd15: {e}")
         error_img = Image.new('RGB', (512, 512), color='red')
         return error_img, Image.new('RGB', (512, 512), color='gray')
 
-def t2i(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
-        seed, steps, scale, w, h, use_refiner=False):
+def colorize_sdxl(sketch, base_model, controlnet_type, lora_model, lora_weight, vae_model,
+                  prompt, negative_prompt, seed, steps, scale, cn_weight):
+    """Colorize function for SDXL models"""
+    try:
+        if base_model not in SDXL_MODELS:
+            error_img = Image.new('RGB', (1024, 1024), color='red')
+            return error_img, Image.new('RGB', (1024, 1024), color='gray')
+        
+        if controlnet_type not in CONTROLNET_MODELS_SDXL:
+            error_img = Image.new('RGB', (1024, 1024), color='red')
+            return error_img, Image.new('RGB', (1024, 1024), color='gray')
+        
+        pipe = get_pipeline(base_model, controlnet_type, lora_model, lora_weight, vae_model)
+        
+        status_msg = f"🎨 Using: {base_model} + {controlnet_type}"
+        if lora_model and lora_model != "None":
+            status_msg += f" + {lora_model}"
+        print(status_msg)
+        
+        condition_img = prepare_condition_image(sketch, controlnet_type, is_sdxl=True)
+        
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        
+        with torch.inference_mode():
+            out = pipe(
+                prompt,
+                negative_prompt=negative_prompt,
+                image=condition_img,
+                num_inference_steps=int(steps),
+                guidance_scale=float(scale),
+                controlnet_conditioning_scale=float(cn_weight),
+                generator=gen,
+                height=1024,
+                width=1024
+            ).images[0]
+        
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+        
+        return out, condition_img
+    except Exception as e:
+        print(f"❌ Error in colorize_sdxl: {e}")
+        error_img = Image.new('RGB', (1024, 1024), color='red')
+        return error_img, Image.new('RGB', (1024, 1024), color='gray')
+
+def t2i_sd15(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
+             seed, steps, scale, w, h):
+    """Text-to-image for SD1.5 models"""
+    try:
+        if model not in SD15_MODELS:
+            error_img = Image.new('RGB', (int(w), int(h)), color='red')
+            return error_img
+        
+        load_t2i_model(model, lora_model, lora_weight, vae_model)
+        
+        print(f"🖼️ Using SD1.5 model: {model}")
+        if lora_model and lora_model != "None":
+            print(f"   with LoRA: {lora_model} (weight: {lora_weight})")
+        
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        
+        with torch.inference_mode():
+            result = CURRENT_T2I_PIPE(
+                prompt,
+                negative_prompt=negative_prompt,
+                width=int(w),
+                height=int(h),
+                num_inference_steps=int(steps),
+                guidance_scale=float(scale),
+                generator=gen
+            ).images[0]
+        
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+        
+        return result
+    except Exception as e:
+        print(f"❌ Error in t2i_sd15: {e}")
+        error_img = Image.new('RGB', (int(w), int(h)), color='red')
+        from PIL import ImageDraw, ImageFont
+        draw = ImageDraw.Draw(error_img)
+        try:
+            font = ImageFont.truetype("arial.ttf", 20)
+        except:
+            font = ImageFont.load_default()
+        draw.text((50, 50), f"Error: {str(e)[:50]}...", fill="white", font=font)
+        return error_img
+
+def t2i_sdxl(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
+             seed, steps, scale, w, h, use_refiner=False):
+    """Text-to-image for SDXL models"""
     try:
+        if model not in SDXL_MODELS:
+            error_img = Image.new('RGB', (int(w), int(h)), color='red')
+            return error_img
+        
         model_to_load = model
         if use_refiner and "refiner" not in model.lower():
             model_to_load = "stabilityai/stable-diffusion-xl-refiner-1.0"
         
         load_t2i_model(model_to_load, lora_model, lora_weight, vae_model)
         
-        print(f"🖼️ Using T2I model: {model}")
+        print(f"🖼️ Using SDXL model: {model}")
         if lora_model and lora_model != "None":
             print(f"   with LoRA: {lora_model} (weight: {lora_weight})")
+        if use_refiner:
+            print(f"   with refiner")
         
         gen = torch.Generator(device=device).manual_seed(int(seed))
         
@@ -623,35 +704,22 @@ def t2i(prompt, negative_prompt, model, lora_model, lora_weight, vae_model,
                     generator=gen
                 ).images[0]
             else:
-                if is_sdxl_model(model):
-                    width = max(int(w), 512)
-                    height = max(int(h), 512)
-                    result = CURRENT_T2I_PIPE(
-                        prompt,
-                        negative_prompt=negative_prompt,
-                        width=width,
-                        height=height,
-                        num_inference_steps=int(steps),
-                        guidance_scale=float(scale),
-                        generator=gen
-                    ).images[0]
-                else:
-                    result = CURRENT_T2I_PIPE(
-                        prompt,
-                        negative_prompt=negative_prompt,
-                        width=int(w),
-                        height=int(h),
-                        num_inference_steps=int(steps),
-                        guidance_scale=float(scale),
-                        generator=gen
-                    ).images[0]
+                result = CURRENT_T2I_PIPE(
+                    prompt,
+                    negative_prompt=negative_prompt,
+                    width=int(w),
+                    height=int(h),
+                    num_inference_steps=int(steps),
+                    guidance_scale=float(scale),
+                    generator=gen
+                ).images[0]
         
         if device.type == "cuda":
             torch.cuda.empty_cache()
         
         return result
     except Exception as e:
-        print(f"❌ Error in t2i: {e}")
+        print(f"❌ Error in t2i_sdxl: {e}")
         error_img = Image.new('RGB', (int(w), int(h)), color='red')
         from PIL import ImageDraw, ImageFont
         draw = ImageDraw.Draw(error_img)
@@ -727,129 +795,187 @@ with gr.Blocks(title="🎨 AI Image Generator Pro", theme=gr.themes.Soft()) as d
         status_text = gr.Textbox(label="Status", interactive=False, scale=3)
     unload_btn.click(unload_all_models, outputs=status_text)
 
-    with gr.Tab("🎨 ControlNet Image-to-Image"):
+    with gr.Tab("🎨 SD1.5 ControlNet"):
         gr.Markdown("""
-        ### Transform sketches/images using ControlNet
-        - **SD1.5 Models:** Support all ControlNet types
-        - **SDXL Models:** Support canny_sdxl, depth_sdxl, openpose_sdxl only
+        ### Transform sketches/images using SD1.5 with ControlNet
+        - **Supports:** lineart, lineart_anime, canny, depth, normal, openpose, softedge, segmentation, mlsd, shuffle, scribble, tile
+        - **Best Resolution:** 512x512
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                inp = gr.Image(label="Input Sketch/Image", type="pil")
+                inp_sd15 = gr.Image(label="Input Sketch/Image", type="pil")
                 
                 gr.Markdown("### Model Settings")
-                base_model = gr.Dropdown(
-                    choices=ALL_MODELS,
+                base_model_sd15 = gr.Dropdown(
+                    choices=SD15_MODELS,
                     value="digiplay/ChikMix_V3",
-                    label="Base Model"
+                    label="SD1.5 Base Model"
                 )
-                controlnet_type = gr.Dropdown(
-                    choices=list(CONTROLNET_MODELS.keys()) + list(SDXL_CONTROLNET_MODELS.keys()),
+                controlnet_type_sd15 = gr.Dropdown(
+                    choices=list(CONTROLNET_MODELS_SD15.keys()),
                     value="lineart_anime",
                     label="ControlNet Type"
                 )
                 
                 gr.Markdown("### Enhancement Options")
                 with gr.Row():
-                    lora_model = gr.Dropdown(
+                    lora_model_sd15 = gr.Dropdown(
                         choices=list(LORA_MODELS.keys()),
                         value="None",
                         label="LoRA Model"
                     )
-                    lora_weight = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                    lora_weight_sd15 = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                 
-                vae_model = gr.Dropdown(
-                    choices=list(VAE_MODELS.keys()),
+                vae_model_sd15 = gr.Dropdown(
+                    choices=["None", "SD1.5 VAE", "Anime VAE"],
                     value="None",
                     label="VAE Model (Optional)"
                 )
             
             with gr.Column(scale=1):
-                out = gr.Image(label="Generated Output")
-                condition_out = gr.Image(label="Processed Condition", type="pil")
+                out_sd15 = gr.Image(label="Generated Output")
+                condition_out_sd15 = gr.Image(label="Processed Condition", type="pil")
         
         gr.Markdown("### Generation Parameters")
         with gr.Row():
-            prompt = gr.Textbox(
+            prompt_sd15 = gr.Textbox(
                 label="Prompt",
                 placeholder="masterpiece, best quality, 1girl, beautiful detailed eyes, long hair",
                 lines=3
             )
-            negative_prompt = gr.Textbox(
+            negative_prompt_sd15 = gr.Textbox(
                 label="Negative Prompt",
                 placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers",
                 lines=3
             )
         
         with gr.Row():
-            seed = gr.Number(value=-1, label="Seed (-1 for random)")
-            steps = gr.Slider(10, 150, 30, step=1, label="Steps")
-            scale = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
-            cn_weight = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
-        
-        run = gr.Button("🎨 Generate", variant="primary", size="lg")
-        run.click(
-            colorize,
-            [inp, base_model, controlnet_type, lora_model, lora_weight, vae_model,
-             prompt, negative_prompt, seed, steps, scale, cn_weight],
-            [out, condition_out]
-        )
+            seed_sd15 = gr.Number(value=-1, label="Seed (-1 for random)")
+            steps_sd15 = gr.Slider(10, 100, 30, step=1, label="Steps")
+            scale_sd15 = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
+            cn_weight_sd15 = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
         
+        run_sd15 = gr.Button("🎨 Generate (SD1.5)", variant="primary", size="lg")
+        run_sd15.click(
+            colorize_sd15,
+            [inp_sd15, base_model_sd15, controlnet_type_sd15, lora_model_sd15, lora_weight_sd15, vae_model_sd15,
+             prompt_sd15, negative_prompt_sd15, seed_sd15, steps_sd15, scale_sd15, cn_weight_sd15],
+            [out_sd15, condition_out_sd15]
+        )
+
+    with gr.Tab("🎨 SDXL ControlNet"):
         gr.Markdown("""
-        ### Tips for Better Results:
-        - Use detailed prompts for better control
-        - Adjust ControlNet weight to balance between condition and creativity
-        - Try different LoRA models for various styles
-        - Higher steps = better quality but slower generation
+        ### Transform sketches/images using SDXL with ControlNet
+        - **Supports:** canny_sdxl, depth_sdxl, openpose_sdxl
+        - **Best Resolution:** 1024x1024
+        - **Higher quality, more VRAM required**
         """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                inp_sdxl = gr.Image(label="Input Sketch/Image", type="pil")
+                
+                gr.Markdown("### Model Settings")
+                base_model_sdxl = gr.Dropdown(
+                    choices=SDXL_MODELS,
+                    value="stabilityai/stable-diffusion-xl-base-1.0",
+                    label="SDXL Base Model"
+                )
+                controlnet_type_sdxl = gr.Dropdown(
+                    choices=list(CONTROLNET_MODELS_SDXL.keys()),
+                    value="canny_sdxl",
+                    label="ControlNet Type"
+                )
+                
+                gr.Markdown("### Enhancement Options")
+                with gr.Row():
+                    lora_model_sdxl = gr.Dropdown(
+                        choices=list(LORA_MODELS.keys()),
+                        value="None",
+                        label="LoRA Model"
+                    )
+                    lora_weight_sdxl = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                
+                vae_model_sdxl = gr.Dropdown(
+                    choices=["None", "SDXL VAE"],
+                    value="None",
+                    label="VAE Model (Optional)"
+                )
+            
+            with gr.Column(scale=1):
+                out_sdxl = gr.Image(label="Generated Output")
+                condition_out_sdxl = gr.Image(label="Processed Condition", type="pil")
+        
+        gr.Markdown("### Generation Parameters")
+        with gr.Row():
+            prompt_sdxl = gr.Textbox(
+                label="Prompt",
+                placeholder="masterpiece, best quality, 8k, ultra-detailed, photorealistic, beautiful lighting",
+                lines=3
+            )
+            negative_prompt_sdxl = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits",
+                lines=3
+            )
+        
+        with gr.Row():
+            seed_sdxl = gr.Number(value=-1, label="Seed (-1 for random)")
+            steps_sdxl = gr.Slider(10, 100, 30, step=1, label="Steps")
+            scale_sdxl = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
+            cn_weight_sdxl = gr.Slider(0.1, 2.0, 1.0, step=0.1, label="ControlNet Weight")
+        
+        run_sdxl = gr.Button("🎨 Generate (SDXL)", variant="primary", size="lg")
+        run_sdxl.click(
+            colorize_sdxl,
+            [inp_sdxl, base_model_sdxl, controlnet_type_sdxl, lora_model_sdxl, lora_weight_sdxl, vae_model_sdxl,
+             prompt_sdxl, negative_prompt_sdxl, seed_sdxl, steps_sdxl, scale_sdxl, cn_weight_sdxl],
+            [out_sdxl, condition_out_sdxl]
+        )
 
-    with gr.Tab("🖼️ Text-to-Image Generation"):
+    with gr.Tab("🖼️ SD1.5 Text-to-Image"):
         gr.Markdown("""
-        ### Generate images from text descriptions
-        Supports both SD1.5 and SDXL models with advanced features
+        ### Generate images from text descriptions using SD1.5
+        - **Best Resolution:** 512x512, 512x768, 768x512
+        - **Faster generation, lower VRAM usage**
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
                 gr.Markdown("### Model Configuration")
-                t2i_model = gr.Dropdown(
-                    choices=ALL_MODELS,
+                t2i_model_sd15 = gr.Dropdown(
+                    choices=SD15_MODELS,
                     value="digiplay/ChikMix_V3",
-                    label="Base Model"
+                    label="SD1.5 Base Model"
                 )
                 
                 gr.Markdown("### Enhancement Options")
                 with gr.Row():
-                    t2i_lora = gr.Dropdown(
+                    t2i_lora_sd15 = gr.Dropdown(
                         choices=list(LORA_MODELS.keys()),
                         value="None",
                         label="LoRA Model"
                     )
-                    t2i_lora_weight = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                    t2i_lora_weight_sd15 = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
                 
-                t2i_vae = gr.Dropdown(
-                    choices=list(VAE_MODELS.keys()),
+                t2i_vae_sd15 = gr.Dropdown(
+                    choices=["None", "SD1.5 VAE", "Anime VAE"],
                     value="None",
                     label="VAE Model"
                 )
-                
-                use_refiner = gr.Checkbox(
-                    label="Use SDXL Refiner (SDXL only)",
-                    value=False
-                )
             
             with gr.Column(scale=1):
-                t2i_out = gr.Image(label="Generated Image", type="pil")
+                t2i_out_sd15 = gr.Image(label="Generated Image", type="pil")
         
         gr.Markdown("### Prompts")
         with gr.Row():
-            t2i_prompt = gr.Textbox(
+            t2i_prompt_sd15 = gr.Textbox(
                 label="Prompt",
                 lines=4,
-                placeholder="masterpiece, best quality, highly detailed, 8k, photorealistic, beautiful lighting"
+                placeholder="masterpiece, best quality, highly detailed, beautiful, 1girl"
             )
-            t2i_negative_prompt = gr.Textbox(
+            t2i_negative_prompt_sd15 = gr.Textbox(
                 label="Negative Prompt",
                 lines=4,
                 placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
@@ -857,116 +983,194 @@ with gr.Blocks(title="🎨 AI Image Generator Pro", theme=gr.themes.Soft()) as d
         
         gr.Markdown("### Generation Parameters")
         with gr.Row():
-            t2i_seed = gr.Number(value=-1, label="Seed (-1 for random)")
-            t2i_steps = gr.Slider(10, 150, 30, step=1, label="Steps")
-            t2i_scale = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
+            t2i_seed_sd15 = gr.Number(value=-1, label="Seed (-1 for random)")
+            t2i_steps_sd15 = gr.Slider(10, 100, 30, step=1, label="Steps")
+            t2i_scale_sd15 = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
         
         with gr.Row():
-            w = gr.Slider(256, 2048, 512, step=64, label="Width")
-            h = gr.Slider(256, 2048, 768, step=64, label="Height")
-        
-        gen_btn = gr.Button("🖼️ Generate Image", variant="primary", size="lg")
-        gen_btn.click(
-            t2i,
-            [t2i_prompt, t2i_negative_prompt, t2i_model, t2i_lora, t2i_lora_weight,
-             t2i_vae, t2i_seed, t2i_steps, t2i_scale, w, h, use_refiner],
-            t2i_out
-        )
+            w_sd15 = gr.Slider(256, 1024, 512, step=64, label="Width")
+            h_sd15 = gr.Slider(256, 1024, 768, step=64, label="Height")
         
+        gen_btn_sd15 = gr.Button("🖼️ Generate (SD1.5)", variant="primary", size="lg")
+        gen_btn_sd15.click(
+            t2i_sd15,
+            [t2i_prompt_sd15, t2i_negative_prompt_sd15, t2i_model_sd15, t2i_lora_sd15, t2i_lora_weight_sd15,
+             t2i_vae_sd15, t2i_seed_sd15, t2i_steps_sd15, t2i_scale_sd15, w_sd15, h_sd15],
+            t2i_out_sd15
+        )
+
+    with gr.Tab("🖼️ SDXL Text-to-Image"):
         gr.Markdown("""
-        ### Pro Tips:
-        - **SDXL models** produce higher quality at 1024x1024
-        - **SD1.5 models** work best at 512x512 or 512x768
-        - Use **LoRA** for specific styles (anime, realistic, etc.)
-        - Use **VAE** for better colors and details
-        - **Refiner** adds extra polish to SDXL generations
-        - Higher **CFG Scale** = more prompt adherence
+        ### Generate images from text descriptions using SDXL
+        - **Best Resolution:** 1024x1024, 1024x1536, 1536x1024
+        - **Higher quality, more detail, better composition**
         """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Model Configuration")
+                t2i_model_sdxl = gr.Dropdown(
+                    choices=SDXL_MODELS,
+                    value="stabilityai/stable-diffusion-xl-base-1.0",
+                    label="SDXL Base Model"
+                )
+                
+                gr.Markdown("### Enhancement Options")
+                with gr.Row():
+                    t2i_lora_sdxl = gr.Dropdown(
+                        choices=list(LORA_MODELS.keys()),
+                        value="None",
+                        label="LoRA Model"
+                    )
+                    t2i_lora_weight_sdxl = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="LoRA Weight")
+                
+                t2i_vae_sdxl = gr.Dropdown(
+                    choices=["None", "SDXL VAE"],
+                    value="None",
+                    label="VAE Model"
+                )
+                
+                use_refiner_sdxl = gr.Checkbox(
+                    label="Use Refiner (for better quality)",
+                    value=False
+                )
+            
+            with gr.Column(scale=1):
+                t2i_out_sdxl = gr.Image(label="Generated Image", type="pil")
+        
+        gr.Markdown("### Prompts")
+        with gr.Row():
+            t2i_prompt_sdxl = gr.Textbox(
+                label="Prompt",
+                lines=4,
+                placeholder="masterpiece, best quality, 8k, ultra-detailed, photorealistic, cinematic lighting"
+            )
+            t2i_negative_prompt_sdxl = gr.Textbox(
+                label="Negative Prompt",
+                lines=4,
+                placeholder="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry"
+            )
+        
+        gr.Markdown("### Generation Parameters")
+        with gr.Row():
+            t2i_seed_sdxl = gr.Number(value=-1, label="Seed (-1 for random)")
+            t2i_steps_sdxl = gr.Slider(10, 100, 30, step=1, label="Steps")
+            t2i_scale_sdxl = gr.Slider(1, 30, 7.5, step=0.5, label="CFG Scale")
+        
+        with gr.Row():
+            w_sdxl = gr.Slider(512, 2048, 1024, step=64, label="Width")
+            h_sdxl = gr.Slider(512, 2048, 1024, step=64, label="Height")
+        
+        gen_btn_sdxl = gr.Button("🖼️ Generate (SDXL)", variant="primary", size="lg")
+        gen_btn_sdxl.click(
+            t2i_sdxl,
+            [t2i_prompt_sdxl, t2i_negative_prompt_sdxl, t2i_model_sdxl, t2i_lora_sdxl, t2i_lora_weight_sdxl,
+             t2i_vae_sdxl, t2i_seed_sdxl, t2i_steps_sdxl, t2i_scale_sdxl, w_sdxl, h_sdxl, use_refiner_sdxl],
+            t2i_out_sdxl
+        )
     
     with gr.Tab("📚 Quick Reference"):
         gr.Markdown("""
         # Model & Feature Guide
         
-        ## 🎯 Recommended Models for Different Purposes
-        
-        ### Realistic/Photorealistic
-        - `emilianJR/epiCRealism` - Excellent for realistic portraits
-        - `stablediffusionapi/realistic-vision-v51` - High quality realistic images
-        - `digiplay/majicMIX_realistic_v7` - Great for realistic characters
-        - `SG161222/RealVisXL_V4.0` - SDXL realistic model
-        
-        ### Anime/Cartoon
-        - `digiplay/ChikMix_V3` - Versatile anime style
-        - `gsdf/Counterfeit-V2.5` - High quality anime
-        - `stablediffusionapi/anything-v5` - Popular anime model
-        - `digiplay/Pony_Diffusion_V6_XL` - SDXL anime model
-        
-        ### Artistic/Stylized
-        - `stablediffusionapi/dreamshaper-v8` - Dream-like artistic style
-        - `wavymulder/Analog-Diffusion` - Analog photo aesthetic
-        - `Lykon/dreamshaper-xl-1-0` - SDXL artistic model
-        
-        ## 🎨 ControlNet Types Explained
-        
-        - **lineart/lineart_anime**: Convert line drawings to colored images
-        - **canny**: Edge detection based generation
-        - **depth**: Depth map based generation
-        - **openpose**: Human pose based generation
-        - **normal**: Normal map based generation
-        - **softedge**: Soft edge detection
+        ## 🎯 SD1.5 vs SDXL Comparison
+        
+        ### SD1.5 (Stable Diffusion 1.5)
+        - **Pros:** Faster, lower VRAM usage, many specialized models
+        - **Cons:** Lower quality, less detail
+        - **Best for:** Quick generations, testing, lower-end hardware
+        - **Resolution:** 512x512 optimal
+        
+        ### SDXL (Stable Diffusion XL)
+        - **Pros:** Higher quality, better composition, more detail
+        - **Cons:** Slower, higher VRAM usage
+        - **Best for:** Final quality images, professional work
+        - **Resolution:** 1024x1024 optimal
+        
+        ## 🎨 ControlNet Types
+        
+        ### SD1.5 ControlNet (12 types)
+        - **lineart/lineart_anime**: Line art to image
+        - **canny**: Edge detection
+        - **depth**: Depth map based
+        - **openpose**: Human pose
+        - **normal**: Normal map
+        - **softedge**: Soft edges
+        - **segmentation**: Semantic segmentation
+        - **mlsd**: Straight line detection
+        - **shuffle**: Color shuffle
         - **scribble**: Scribble to image
-        - **tile**: Upscaling and detail enhancement
+        - **tile**: Upscaling
+        
+        ### SDXL ControlNet (3 types)
+        - **canny_sdxl**: Edge detection
+        - **depth_sdxl**: Depth map
+        - **openpose_sdxl**: Human pose
         
-        ## 💎 Popular LoRA Combinations
+        ## 💎 Recommended Model Combinations
         
-        ### For Portraits
+        ### SD1.5 - Realistic Portraits
         - Base: `digiplay/majicMIX_realistic_v7`
-        - LoRA: `detail-tweaker` or `face-detail`
+        - LoRA: `detail-tweaker`
         - VAE: `SD1.5 VAE`
+        - Steps: 30-40
+        - CFG: 7-9
         
-        ### For Anime Characters
+        ### SD1.5 - Anime Style
         - Base: `digiplay/ChikMix_V3`
-        - LoRA: `anime-art` or `manga-style`
+        - LoRA: `anime-art`
         - VAE: `Anime VAE`
+        - Steps: 25-35
+        - CFG: 6-8
+        
+        ### SDXL - Photorealistic
+        - Base: `SG161222/RealVisXL_V4.0`
+        - LoRA: `photorealistic`
+        - VAE: `SDXL VAE`
+        - Use Refiner: Yes
+        - Steps: 35-45
+        - CFG: 7-10
         
-        ### For NSFW Content
-        - Base: Any NSFW-capable model
-        - LoRA: `nsfw-master`, `realistic-nsfw`, or `anime-nsfw`
-        - Note: Always use responsibly and legally
+        ### SDXL - Artistic
+        - Base: `Lykon/dreamshaper-xl-1-0`
+        - LoRA: `watercolor-style`
+        - VAE: `SDXL VAE`
+        - Steps: 30-40
+        - CFG: 6-8
         
         ## ⚙️ Parameter Guidelines
         
         ### Steps
-        - **20-30**: Fast, good quality
-        - **30-50**: Balanced
-        - **50-100**: High quality, slow
+        - **SD1.5:** 20-35 (fast), 35-50 (quality)
+        - **SDXL:** 30-45 (fast), 45-60 (quality)
         
         ### CFG Scale
-        - **5-7**: Creative, loose interpretation
-        - **7-10**: Balanced
-        - **10-15**: Strict prompt adherence
-        - **15+**: Very strict, may oversaturate
+        - **4-6:** Creative, loose interpretation
+        - **6-9:** Balanced (recommended)
+        - **9-12:** Strict prompt adherence
+        - **12+:** Very strict
         
         ### Resolution
-        - **SD1.5**: 512x512, 512x768, 768x512
-        - **SDXL**: 1024x1024, 1024x1536, 1536x1024
+        - **SD1.5:** 512x512, 512x768, 768x512 (max 1024x1024)
+        - **SDXL:** 1024x1024, 1024x1536, 1536x1024 (max 2048x2048)
         
-        ## 🔞 NSFW Generation Guidelines
+        ## 🚀 Performance Tips
         
-        1. Use NSFW-capable base models
-        2. Apply relevant LoRA for style enhancement
-        3. Use detailed prompts
-        4. Adjust CFG scale (7-12 recommended)
-        5. Consider using higher steps (40-60)
-        6. **Always comply with local laws and regulations**
+        ### For Low VRAM (<8GB)
+        - Use SD1.5 models only
+        - Enable attention slicing
+        - Use lower resolutions (512x512)
+        - Reduce steps (20-30)
         
-        ## 🚀 Performance Tips
+        ### For Medium VRAM (8-12GB)
+        - Can use SD1.5 and some SDXL
+        - For SDXL, use 1024x1024
+        - Enable xFormers
         
-        - Unload models when switching between different types
-        - Use lower resolutions for testing
-        - Enable xFormers if available (automatic)
-        - Use appropriate batch sizes for your GPU
-        - Monitor GPU memory usage
+        ### For High VRAM (12GB+)
+        - Can use all models
+        - SDXL with refiner
+        - Higher resolutions
         """)
 
 try: