Update app.py

app.py

@@ -10,6 +10,7 @@ dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 # Lazy import (to avoid long startup if unused)
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionPipeline
+from diffusers import StableDiffusionInstructPix2PixPipeline
 from controlnet_aux import LineartDetector, LineartAnimeDetector
 
 # ===== Model & Config =====
@@ -18,6 +19,8 @@ LINEART_DETECTOR = None
 LINEART_ANIME_DETECTOR = None
 CURRENT_T2I_PIPE = None
 CURRENT_T2I_MODEL = None
+CURRENT_PIX2PIX_PIPE = None
+CURRENT_PIX2PIX_MODEL = None
 
 def get_pipeline(model_name: str, anime_model: bool = False):
     """Get or create a ControlNet pipeline for the given model and anime flag"""
@@ -28,60 +31,108 @@ def get_pipeline(model_name: str, anime_model: bool = False):
     
     print(f"Loading ControlNet pipeline for model: {model_name}, anime: {anime_model}")
     
-    # โหลด ControlNet ที่เหมาะสม
-    if anime_model:
-        controlnet = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15s2_lineart_anime", 
-            torch_dtype=dtype
-        ).to(device)
-    else:
-        controlnet = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_lineart", 
-            torch_dtype=dtype
+    try:
+        # โหลด ControlNet ที่เหมาะสม
+        if anime_model:
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/control_v11p_sd15s2_lineart_anime", 
+                torch_dtype=dtype
+            ).to(device)
+        else:
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/control_v11p_sd15_lineart", 
+                torch_dtype=dtype
+            ).to(device)
+        
+        # สร้าง pipeline
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            model_name, 
+            controlnet=controlnet, 
+            torch_dtype=dtype,
+            safety_checker=None, 
+            requires_safety_checker=False
         ).to(device)
-    
-    # สร้าง pipeline
-    pipe = StableDiffusionControlNetPipeline.from_pretrained(
-        model_name, 
-        controlnet=controlnet, 
-        torch_dtype=dtype,
-        safety_checker=None, 
-        requires_safety_checker=False
-    ).to(device)
-    
-    pipe.enable_attention_slicing()
-    if device.type == "cuda":
-        pipe.enable_model_cpu_offload()
-    
-    # เก็บ pipeline ไว้ใน cache
-    PIPELINES[key] = pipe
-    return pipe
+        
+        pipe.enable_attention_slicing()
+        if device.type == "cuda":
+            pipe.enable_model_cpu_offload()
+        
+        # เก็บ pipeline ไว้ใน cache
+        PIPELINES[key] = pipe
+        return pipe
+        
+    except Exception as e:
+        print(f"Error loading ControlNet pipeline: {e}")
+        raise
 
 def load_lineart_detectors():
     """Load lineart detectors if not already loaded"""
     global LINEART_DETECTOR, LINEART_ANIME_DETECTOR
     if LINEART_DETECTOR is None:
         print("Loading lineart detectors...")
-        LINEART_DETECTOR = LineartDetector.from_pretrained("lllyasviel/Annotators")
-        LINEART_ANIME_DETECTOR = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
+        try:
+            LINEART_DETECTOR = LineartDetector.from_pretrained("lllyasviel/Annotators")
+            LINEART_ANIME_DETECTOR = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators")
+        except Exception as e:
+            print(f"Error loading lineart detectors: {e}")
+            raise
 
 def load_t2i_model(model_name: str):
     global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL
-    if CURRENT_T2I_MODEL == model_name and CURRENT_T2I_PIPE is not None:
-        return
-    if CURRENT_T2I_PIPE is not None:
-        del CURRENT_T2I_PIPE
-        gc.collect()
-        if torch.cuda.is_available():
-            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
-    ).to(device)
-    CURRENT_T2I_PIPE.enable_attention_slicing()
-    if device.type == "cuda":
-        CURRENT_T2I_PIPE.enable_model_cpu_offload()
-    CURRENT_T2I_MODEL = model_name
+    try:
+        if CURRENT_T2I_MODEL == model_name and CURRENT_T2I_PIPE is not None:
+            return
+        if CURRENT_T2I_PIPE is not None:
+            del CURRENT_T2I_PIPE
+            CURRENT_T2I_PIPE = None
+            gc.collect()
+            if torch.cuda.is_available():
+                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
+        ).to(device)
+        CURRENT_T2I_PIPE.enable_attention_slicing()
+        if device.type == "cuda":
+            CURRENT_T2I_PIPE.enable_model_cpu_offload()
+        CURRENT_T2I_MODEL = model_name
+        
+    except Exception as e:
+        print(f"Error loading T2I model {model_name}: {e}")
+        # รีเซ็ตตัวแปรเมื่อโหลดไม่สำเร็จ
+        CURRENT_T2I_PIPE = None
+        CURRENT_T2I_MODEL = None
+        raise
+
+def load_pix2pix_model():
+    """Load Instruct-Pix2Pix model for image editing"""
+    global CURRENT_PIX2PIX_PIPE, CURRENT_PIX2PIX_MODEL
+    
+    if CURRENT_PIX2PIX_PIPE is not None:
+        return CURRENT_PIX2PIX_PIPE
+    
+    try:
+        print("Loading Instruct-Pix2Pix model...")
+        CURRENT_PIX2PIX_PIPE = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+            "timbrooks/instruct-pix2pix",
+            torch_dtype=dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(device)
+        
+        CURRENT_PIX2PIX_PIPE.enable_attention_slicing()
+        if device.type == "cuda":
+            CURRENT_PIX2PIX_PIPE.enable_model_cpu_offload()
+        
+        CURRENT_PIX2PIX_MODEL = "timbrooks/instruct-pix2pix"
+        return CURRENT_PIX2PIX_PIPE
+        
+    except Exception as e:
+        print(f"Error loading Instruct-Pix2Pix model: {e}")
+        CURRENT_PIX2PIX_PIPE = None
+        CURRENT_PIX2PIX_MODEL = None
+        raise
 
 # ===== Utils =====
 def is_lineart(img: Image.Image) -> bool:
@@ -97,76 +148,158 @@ def extract_lineart(img, anime: bool = False):
     out = detector(img, detect_resolution=512, image_resolution=512)
     return Image.fromarray(out) if isinstance(out, np.ndarray) else out
 
+def resize_image(image, max_size=512):
+    """Resize image while maintaining aspect ratio"""
+    width, height = image.size
+    if max(width, height) > max_size:
+        ratio = max_size / max(width, height)
+        new_width = int(width * ratio)
+        new_height = int(height * ratio)
+        return image.resize((new_width, new_height), Image.LANCZOS)
+    return image
+
 # ===== Functions =====
 def colorize(sketch, base_model, anime_model, prompt, seed, steps, scale, cn_weight):
-    # โหลด pipeline ที่เหมาะสม
-    pipe = get_pipeline(base_model, anime_model)
-    
-    # สกัด lineart
-    lineart = extract_lineart(sketch, anime_model)
-    
-    # สร้างภาพ
-    gen = torch.Generator(device=device).manual_seed(int(seed))
-    out = pipe(
-        prompt, 
-        image=lineart, 
-        num_inference_steps=int(steps),
-        guidance_scale=float(scale), 
-        controlnet_conditioning_scale=float(cn_weight),
-        generator=gen
-    ).images[0]
-    
-    return out, lineart
+    try:
+        # โหลด pipeline ที่เหมาะสม
+        pipe = get_pipeline(base_model, anime_model)
+        
+        # สกัด lineart
+        lineart = extract_lineart(sketch, anime_model)
+        
+        # สร้างภาพ
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        out = pipe(
+            prompt, 
+            image=lineart, 
+            num_inference_steps=int(steps),
+            guidance_scale=float(scale), 
+            controlnet_conditioning_scale=float(cn_weight),
+            generator=gen
+        ).images[0]
+        
+        return out, lineart
+    except Exception as e:
+        print(f"Error in colorize: {e}")
+        # ส่งกลับรูปภาพว่างพร้อมแสดง error
+        error_img = Image.new('RGB', (512, 512), color='red')
+        error_text = f"Error: {str(e)[:50]}..."
+        return error_img, Image.new('RGB', (512, 512), color='gray')
 
 def t2i(prompt, model, seed, steps, scale, w, h):
-    load_t2i_model(model)
-    gen = torch.Generator(device=device).manual_seed(int(seed))
-    return CURRENT_T2I_PIPE(
-        prompt, 
-        width=int(w), 
-        height=int(h),
-        num_inference_steps=int(steps), 
-        guidance_scale=float(scale),
-        generator=gen
-    ).images[0]
+    try:
+        load_t2i_model(model)
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        return CURRENT_T2I_PIPE(
+            prompt, 
+            width=int(w), 
+            height=int(h),
+            num_inference_steps=int(steps), 
+            guidance_scale=float(scale),
+            generator=gen
+        ).images[0]
+    except Exception as e:
+        print(f"Error in t2i: {e}")
+        # ส่งกลับรูปภาพว่างพร้อมแสดง error
+        error_img = Image.new('RGB', (int(w), int(h)), color='red')
+        return error_img
+
+def pix2pix_edit(image, instruction, seed, steps, scale, image_scale):
+    """Edit image using Instruct-Pix2Pix"""
+    try:
+        # โหลดโมเดล
+        pipe = load_pix2pix_model()
+        
+        # ปรับขนาดภาพ
+        image = resize_image(image, max_size=768)
+        
+        # สร้าง generator
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        
+        # แก��ไขภาพ
+        result = pipe(
+            instruction,
+            image=image,
+            num_inference_steps=int(steps),
+            guidance_scale=float(scale),
+            image_guidance_scale=float(image_scale),
+            generator=gen
+        ).images[0]
+        
+        return result
+        
+    except Exception as e:
+        print(f"Error in pix2pix_edit: {e}")
+        # ส่งกลับรูปภาพ error
+        if image:
+            error_img = Image.new('RGB', image.size, color='red')
+        else:
+            error_img = Image.new('RGB', (512, 512), color='red')
+        return error_img
 
 # ===== Function to unload all models =====
 def unload_all_models():
     global PIPELINES, LINEART_DETECTOR, LINEART_ANIME_DETECTOR
     global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL
+    global CURRENT_PIX2PIX_PIPE, CURRENT_PIX2PIX_MODEL
     
     print("Unloading all models from memory...")
     
     # Unload ControlNet pipelines
     for key, pipe in list(PIPELINES.items()):
-        del pipe
+        try:
+            del pipe
+        except:
+            pass
     PIPELINES.clear()
     
     # Unload lineart detectors
-    if LINEART_DETECTOR is not None:
-        del LINEART_DETECTOR
-        LINEART_DETECTOR = None
-    if LINEART_ANIME_DETECTOR is not None:
-        del LINEART_ANIME_DETECTOR
-        LINEART_ANIME_DETECTOR = None
+    try:
+        if LINEART_DETECTOR is not None:
+            del LINEART_DETECTOR
+            LINEART_DETECTOR = None
+    except:
+        pass
+    
+    try:
+        if LINEART_ANIME_DETECTOR is not None:
+            del LINEART_ANIME_DETECTOR
+            LINEART_ANIME_DETECTOR = None
+    except:
+        pass
     
     # Unload T2I model
-    if CURRENT_T2I_PIPE is not None:
-        del CURRENT_T2I_PIPE
-        CURRENT_T2I_PIPE = None
+    try:
+        if CURRENT_T2I_PIPE is not None:
+            del CURRENT_T2I_PIPE
+            CURRENT_T2I_PIPE = None
+    except:
+        pass
+    
     CURRENT_T2I_MODEL = None
     
+    # Unload Pix2Pix model
+    try:
+        if CURRENT_PIX2PIX_PIPE is not None:
+            del CURRENT_PIX2PIX_PIPE
+            CURRENT_PIX2PIX_PIPE = None
+    except:
+        pass
+    
+    CURRENT_PIX2PIX_MODEL = None
+    
     # Force garbage collection
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        print(f"GPU memory cleared. Allocated: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        print(f"GPU memory cleared. Allocated: {allocated:.2f} GB")
     
     return "✅ All models unloaded from memory!"
 
 # ===== Gradio UI =====
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Minimal Style2Paints")
+    gr.Markdown("# 🎨 Advanced Image Generation & Editing Suite")
     
     # Add unload button at the top
     with gr.Row():
@@ -174,7 +307,7 @@ with gr.Blocks() as demo:
         status_text = gr.Textbox(label="Status", interactive=False)
     unload_btn.click(unload_all_models, outputs=status_text)
 
-    with gr.Tab("🎨 Colorize"):
+    with gr.Tab("🎨 Colorize Sketch"):
         with gr.Row():
             inp = gr.Image(label="Input Sketch/Image", type="pil")
             out = gr.Image(label="Colored Output")
@@ -183,18 +316,16 @@ with gr.Blocks() as demo:
             sketch_out = gr.Image(label="Detected Lineart", type="pil")
         
         with gr.Row():
-            # Dropdown สำหรับเลือกโมเดลฐาน
             base_model = gr.Dropdown(
                 choices=[
-                    "admruul/anything-v3.0",  # เปลี่ยนเป็น admruul/anything-v3.0
+                    "admruul/anything-v3.0",
                     "digiplay/ChikMix_V3",
                     "digiplay/chilloutmix_NiPrunedFp16Fix",
                     "runwayml/stable-diffusion-v1-5",
                     "stabilityai/stable-diffusion-2-1",
-                    "andite/anything-v4.0",  # เพิ่มอีกตัวเลือก
-                    "gsdf/Counterfeit-V2.5"  # เพิ่มอีกตัวเลือก
+                    "gsdf/Counterfeit-V2.5"
                 ],
-                value="admruul/anything-v3.0",  # เปลี่ยนค่าเริ่มต้นเป็น admruul/anything-v3.0
+                value="admruul/anything-v3.0",
                 label="Base Model"
             )
             anime_chk = gr.Checkbox(label="Use Anime ControlNet", value=True)
@@ -227,15 +358,14 @@ with gr.Blocks() as demo:
             t2i_prompt = gr.Textbox(label="Prompt", lines=3)
             t2i_model = gr.Dropdown(
                 choices=[
-                    "admruul/anything-v3.0",  # เปลี่ยนเป็น admruul/anything-v3.0
+                    "admruul/anything-v3.0",
                     "digiplay/ChikMix_V3",
                     "digiplay/chilloutmix_NiPrunedFp16Fix",
                     "runwayml/stable-diffusion-v1-5",
                     "stabilityai/stable-diffusion-2-1",
-                    "andite/anything-v4.0",  # เพิ่มอีกตัวเลือก
-                    "gsdf/Counterfeit-V2.5"  # เพิ่มอีกตัวเลือก
+                    "gsdf/Counterfeit-V2.5"
                 ],
-                value="admruul/anything-v3.0",  # เปลี่ยนค่าเริ่มต้นเป็น admruul/anything-v3.0
+                value="admruul/anything-v3.0",
                 label="Model"
             )
         
@@ -254,5 +384,56 @@ with gr.Blocks() as demo:
             [t2i_prompt, t2i_model, t2i_seed, t2i_steps, t2i_scale, w, h], 
             t2i_out
         )
+    
+    with gr.Tab("🔄 Instruct-Pix2Pix"):
+        gr.Markdown("### Edit Images with Text Instructions")
+        gr.Markdown("ตัวอย่างคำสั่ง: 'make it winter', 'turn day into night', 'add sunglasses', 'make it look like a painting'")
+        
+        with gr.Row():
+            with gr.Column():
+                pix2pix_input = gr.Image(label="Input Image", type="pil")
+                pix2pix_instruction = gr.Textbox(
+                    label="Edit Instruction",
+                    placeholder="e.g., make it winter, turn day into night, add sunglasses...",
+                    lines=2
+                )
+                
+                with gr.Row():
+                    pix2pix_seed = gr.Number(value=42, label="Seed")
+                    pix2pix_steps = gr.Slider(10, 100, 50, step=5, label="Steps")
+                
+                with gr.Row():
+                    pix2pix_scale = gr.Slider(1, 20, 7.5, step=0.5, label="Text Guidance Scale")
+                    pix2pix_image_scale = gr.Slider(1, 5, 1.5, step=0.1, label="Image Guidance Scale")
+                
+                pix2pix_btn = gr.Button("🔄 Edit Image", variant="primary")
+            
+            with gr.Column():
+                pix2pix_output = gr.Image(label="Edited Image", type="pil")
+        
+        # ตัวอย่างคำสั่งที่พบบ่อย
+        with gr.Row():
+            gr.Examples(
+                examples=[
+                    ["make it winter", 42, 50, 7.5, 1.5],
+                    ["turn day into night", 42, 50, 7.5, 1.5],
+                    ["make it look like a painting", 42, 50, 7.5, 1.5],
+                    ["add sunglasses", 42, 50, 7.5, 1.5],
+                    ["make it cyberpunk style", 42, 50, 7.5, 1.5],
+                    ["change hair color to blue", 42, 50, 7.5, 1.5],
+                ],
+                inputs=[pix2pix_instruction, pix2pix_seed, pix2pix_steps, pix2pix_scale, pix2pix_image_scale],
+                label="Quick Examples"
+            )
+        
+        pix2pix_btn.click(
+            pix2pix_edit,
+            [pix2pix_input, pix2pix_instruction, pix2pix_seed, pix2pix_steps, pix2pix_scale, pix2pix_image_scale],
+            pix2pix_output
+        )
 
-demo.launch()
+# เพิ่ม error handling ในการ launch
+try:
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+except Exception as e:
+    print(f"Error launching Gradio app: {e}")