Update app.py

app.py

@@ -26,6 +26,7 @@ print(f"🖥️  Device: {device} | dtype: {dtype}")
 
 # Lazy import (to avoid long startup if unused)
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionPipeline
+from diffusers import StableDiffusionXLPipeline  # สำหรับ SDXL models
 from controlnet_aux import LineartDetector, LineartAnimeDetector
 
 # Memory optimization
@@ -46,6 +47,22 @@ LINEART_ANIME_DETECTOR = None
 CURRENT_T2I_PIPE = None
 CURRENT_T2I_MODEL = None
 
+# Define model types
+SDXL_MODELS = ["Laxhar/noobai-XL-1.1"]  # เพิ่ม SDXL models ตรงนี้
+SD15_MODELS = [
+    "digiplay/ChikMix_V3",
+    "digiplay/chilloutmix_NiPrunedFp16Fix",
+    "gsdf/Counterfeit-V2.5",
+    "stablediffusionapi/anything-v5",
+    "digiplay/CleanLinearMix_nsfw"
+]
+
+ALL_MODELS = SD15_MODELS + SDXL_MODELS
+
+def is_sdxl_model(model_name: str) -> bool:
+    """ตรวจสอบว่าโมเดลเป็น SDXL หรือไม่"""
+    return model_name in SDXL_MODELS
+
 def get_pipeline(model_name: str, anime_model: bool = False):
     """Get or create a ControlNet pipeline for the given model and anime flag"""
     global CURRENT_CONTROLNET_PIPE, CURRENT_CONTROLNET_KEY
@@ -57,6 +74,11 @@ def get_pipeline(model_name: str, anime_model: bool = False):
         print(f"✅ Reusing existing ControlNet pipeline: {model_name}, anime: {anime_model}")
         return CURRENT_CONTROLNET_PIPE
     
+    # ถ้าเป็น SDXL model ให้แจ้งเตือนว่าไม่รองรับ ControlNet
+    if is_sdxl_model(model_name):
+        print(f"⚠️  SDXL model {model_name} is not compatible with ControlNet")
+        raise ValueError(f"SDXL model {model_name} is not compatible with ControlNet. Please use SD1.5 models for ControlNet.")
+    
     # ถ้าเป็นโมเดลใหม่ ลบอันเก่าก่อน
     if CURRENT_CONTROLNET_PIPE is not None:
         print(f"🗑️ Unloading old ControlNet pipeline: {CURRENT_CONTROLNET_KEY}")
@@ -220,14 +242,30 @@ def load_t2i_model(model_name: str):
                 torch.cuda.empty_cache()
         
         print(f"📥 Loading T2I model: {model_name}")
-        CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
-            model_name, 
-            torch_dtype=dtype, 
-            safety_checker=None, 
-            requires_safety_checker=False,
-            use_safetensors=True,
-            variant="fp16" if dtype == torch.float16 else None
-        ).to(device)
+        
+        # ตรวจสอบว่าเป็น SDXL หรือ SD1.5 model
+        if is_sdxl_model(model_name):
+            # โหลด SDXL model
+            CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
+                model_name, 
+                torch_dtype=dtype, 
+                safety_checker=None, 
+                requires_safety_checker=False,
+                use_safetensors=True,
+                variant="fp16" if dtype == torch.float16 else None
+            ).to(device)
+            print(f"✅ Loaded SDXL model: {model_name}")
+        else:
+            # โหลด SD1.5 model
+            CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
+                model_name, 
+                torch_dtype=dtype, 
+                safety_checker=None, 
+                requires_safety_checker=False,
+                use_safetensors=True,
+                variant="fp16" if dtype == torch.float16 else None
+            ).to(device)
+            print(f"✅ Loaded SD1.5 model: {model_name}")
         
         # Optimizations
         CURRENT_T2I_PIPE.enable_attention_slicing(slice_size="max")
@@ -265,12 +303,20 @@ def load_t2i_model(model_name: str):
         
         # Retry without use_safetensors
         try:
-            CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
-                model_name, 
-                torch_dtype=dtype, 
-                safety_checker=None, 
-                requires_safety_checker=False
-            ).to(device)
+            if is_sdxl_model(model_name):
+                CURRENT_T2I_PIPE = StableDiffusionXLPipeline.from_pretrained(
+                    model_name, 
+                    torch_dtype=dtype, 
+                    safety_checker=None, 
+                    requires_safety_checker=False
+                ).to(device)
+            else:
+                CURRENT_T2I_PIPE = StableDiffusionPipeline.from_pretrained(
+                    model_name, 
+                    torch_dtype=dtype, 
+                    safety_checker=None, 
+                    requires_safety_checker=False
+                ).to(device)
             
             CURRENT_T2I_PIPE.enable_attention_slicing(slice_size="max")
             if hasattr(CURRENT_T2I_PIPE, 'vae') and hasattr(CURRENT_T2I_PIPE.vae, 'enable_slicing'):
@@ -331,6 +377,22 @@ def resize_image(image, max_size=512):
 # ===== Functions =====
 def colorize(sketch, base_model, anime_model, prompt, seed, steps, scale, cn_weight):
     try:
+        # ตรวจสอบว่าเป็น SDXL model หรือไม่
+        if is_sdxl_model(base_model):
+            error_img = Image.new('RGB', (512, 512), color='red')
+            error_msg_img = Image.new('RGB', (512, 512), color='yellow')
+            # สร้างภาพแสดงข้อความ error
+            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 not compatible", fill="black", font=font)
+            draw.text((50, 230), f"with ControlNet", fill="black", font=font)
+            draw.text((50, 260), f"Use SD1.5 models instead", fill="black", font=font)
+            return error_img, error_msg_img
+        
         # โหลด pipeline ที่เหมาะสม (จะลบอันเก่าออกอัตโนมัติถ้าเปลี่ยนโมเดล)
         pipe = get_pipeline(base_model, anime_model)
         
@@ -373,14 +435,28 @@ def t2i(prompt, model, seed, steps, scale, w, h):
         gen = torch.Generator(device=device).manual_seed(int(seed))
         
         with torch.inference_mode():
-            result = CURRENT_T2I_PIPE(
-                prompt, 
-                width=int(w), 
-                height=int(h),
-                num_inference_steps=int(steps), 
-                guidance_scale=float(scale),
-                generator=gen
-            ).images[0]
+            # สำหรับ SDXL model ใช้ขนาดเริ่มต้นที่ใหญ่กว่า
+            if is_sdxl_model(model):
+                # SDXL ต้องการขนาดขั้นต่ำ 1024x1024 สำหรับผลลัพธ์ที่ดี
+                width = max(int(w), 512)
+                height = max(int(h), 512)
+                result = CURRENT_T2I_PIPE(
+                    prompt, 
+                    width=width, 
+                    height=height,
+                    num_inference_steps=int(steps), 
+                    guidance_scale=float(scale),
+                    generator=gen
+                ).images[0]
+            else:
+                result = CURRENT_T2I_PIPE(
+                    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()
@@ -389,6 +465,14 @@ def t2i(prompt, model, seed, steps, scale, w, h):
     except Exception as e:
         print(f"❌ Error in t2i: {e}")
         error_img = Image.new('RGB', (int(w), int(h)), color='red')
+        # สร้างภาพแสดงข้อความ error
+        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
 
 # ===== Function to unload all models =====
@@ -446,6 +530,7 @@ def unload_all_models():
 with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🎨 Advanced Image Generation & Editing Suite")
     gr.Markdown("### Powered by Stable Diffusion & ControlNet")
+    gr.Markdown("**Note:** SDXL models (noobai-XL) work only in Text-to-Image tab, not in ControlNet")
     
     # Add system info
     if torch.cuda.is_available():
@@ -464,7 +549,7 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
     with gr.Tab("🎨 Colorize Sketch"):
         gr.Markdown("""
         ### Convert your sketches to colored images using ControlNet
-        Upload a sketch or line art, and the AI will automatically colorize it based on your prompt.
+        **Note:** SDXL models are not compatible with ControlNet. Please use SD1.5 models only.
         """)
         
         with gr.Row():
@@ -476,16 +561,9 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
         
         with gr.Row():
             base_model = gr.Dropdown(
-                choices=[
-                    "digiplay/ChikMix_V3",
-                    "digiplay/chilloutmix_NiPrunedFp16Fix",
-                    "gsdf/Counterfeit-V2.5",
-                    "stablediffusionapi/anything-v5",
-                    "digiplay/CleanLinearMix_nsfw",
-                    "Laxhar/noobai-XL-1.1"  # เพิ่มโมเดลใหม่
-                ],
+                choices=SD15_MODELS,  # ใช้เฉพาะ SD1.5 models
                 value="digiplay/ChikMix_V3",
-                label="Base Model"
+                label="Base Model (SD1.5 only)"
             )
             anime_chk = gr.Checkbox(label="Use Anime ControlNet", value=True)
         
@@ -512,7 +590,8 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
     with gr.Tab("🖼️ Text-to-Image"):
         gr.Markdown("""
         ### Generate images from text descriptions
-        Describe what you want to see, and the AI will create it for you.
+        Supports both SD1.5 and SDXL models.
+        **Tip:** SDXL models produce higher quality but require more memory.
         """)
         
         with gr.Row():
@@ -525,14 +604,7 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
                 placeholder="e.g., a beautiful landscape with mountains and a lake at sunset, highly detailed, 4k"
             )
             t2i_model = gr.Dropdown(
-                choices=[
-                    "digiplay/ChikMix_V3",
-                    "digiplay/chilloutmix_NiPrunedFp16Fix",
-                    "gsdf/Counterfeit-V2.5",
-                    "stablediffusionapi/anything-v5",
-                    "digiplay/CleanLinearMix_nsfw",
-                    "Laxhar/noobai-XL-1.1"  # เพิ่มโมเดลใหม่
-                ],
+                choices=ALL_MODELS,  # ใช้ทั้ง SD1.5 และ SDXL
                 value="digiplay/ChikMix_V3",
                 label="Model"
             )
@@ -543,8 +615,9 @@ with gr.Blocks(title="🎨 Advanced Image Generation Suite", theme=gr.themes.Sof
             t2i_scale = gr.Slider(1, 20, 7.5, step=0.5, label="CFG Scale")
         
         with gr.Row():
-            w = gr.Slider(256, 1024, 512, step=64, label="Width")
-            h = gr.Slider(256, 1024, 768, step=64, label="Height")
+            # สำหรับ SDXL ขอแนะนำขนาดที่ใหญ่กว่า
+            w = gr.Slider(256, 1536, 1024, step=64, label="Width")
+            h = gr.Slider(256, 1536, 1024, step=64, label="Height")
         
         gen_btn = gr.Button("🖼️ Generate", variant="primary")
         gen_btn.click(