app.py

import os
import gradio as gr
import numpy as np
import torch
import random
from PIL import Image
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# กำหนดธีม Steel Blue
colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(self, **kwargs):
        super().__init__(
            primary_hue=colors.gray,
            secondary_hue=colors.steel_blue,
            neutral_hue=colors.slate,
            text_size=sizes.text_lg,
            font=(fonts.GoogleFont("Outfit"), "Arial", "sans-serif"),
            font_mono=(fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace"),
        )

steel_blue_theme = SteelBlueTheme()

# ตั้งค่า device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print("=" * 50)
print("🎨 Image Editor - Text Instruction")
print("=" * 50)
print(f"Using device: {device}")
print(f"Using dtype: {dtype}")
print("=" * 50)

from diffusers import StableDiffusionInstructPix2PixPipeline

# โหลดโมเดล
try:
    print("🔄 กำลังโหลดโมเดล InstructPix2Pix...")
    
    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        "timbrooks/instruct-pix2pix",
        torch_dtype=dtype,
        safety_checker=None,
        use_safetensors=True
    )
    
    if device.type == "cuda":
        pipe = pipe.to(device)
        pipe.enable_attention_slicing()  # ลดการใช้ VRAM
    
    print("✅ โหลดโมเดลสำเร็จ!")
    print("=" * 50)
    
except Exception as e:
    print(f"❌ Error: {e}")
    pipe = None

MAX_SEED = np.iinfo(np.int32).max

def resize_image(image, max_size=512):
    """ปรับขนาดภาพให้เหมาะสม"""
    width, height = image.size
    
    if width > max_size or height > max_size:
        if width > height:
            new_width = max_size
            new_height = int(height * (max_size / width))
        else:
            new_height = max_size
            new_width = int(width * (max_size / height))
    else:
        new_width = width
        new_height = height
    
    # ทำให้ขนาดหารด้วย 8 ลงตัว (สำคัญสำหรับ Stable Diffusion)
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return image.resize((new_width, new_height))

def edit_image(
    input_image,
    instruction,
    seed,
    randomize_seed,
    text_guidance_scale,
    image_guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True)
):
    if input_image is None:
        raise gr.Error("กรุณาอัพโหลดภาพ")
    
    if pipe is None:
        raise gr.Error("โมเดลยังไม่พร้อมใช้งาน")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    
    # แปลงภาพเป็น RGB และปรับขนาด
    original_image = input_image.convert("RGB")
    original_image = resize_image(original_image, max_size=512)
    
    print(f"📏 Image size: {original_image.size}")
    print(f"🎯 Instruction: {instruction}")
    print(f"⚙️ Settings: steps={steps}, text_scale={text_guidance_scale}, image_scale={image_guidance_scale}")
    
    try:
        # แก้ไขภาพ
        result = pipe(
            prompt=instruction,
            image=original_image,
            num_inference_steps=int(steps),
            image_guidance_scale=image_guidance_scale,
            guidance_scale=text_guidance_scale,
            generator=generator,
        ).images[0]
        
        return result, seed
    
    except torch.cuda.OutOfMemoryError:
        # หาก RAM/VRAM ไม่พอ
        gr.Warning("หน่วยความจำไม่พอ! ลดขนาดภาพลง...")
        original_image = resize_image(original_image, max_size=384)
        
        result = pipe(
            prompt=instruction,
            image=original_image,
            num_inference_steps=int(steps),
            image_guidance_scale=image_guidance_scale,
            guidance_scale=text_guidance_scale,
            generator=generator,
        ).images[0]
        
        return result, seed
    
    except Exception as e:
        raise gr.Error(f"เกิดข้อผิดพลาด: {str(e)}")

# CSS สำหรับปรับแต่ง UI
css = """
#col-container {
    margin: 0 auto;
    max-width: 1000px;
}
#main-title h1 {
    font-size: 2.2em !important;
    text-align: center;
}
.comparison-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    border-radius: 12px;
    padding: 20px;
    margin: 15px 0;
}
.instruction-examples {
    background: #f5f5f5;
    border-left: 4px solid #4682B4;
    padding: 15px;
    margin: 10px 0;
    border-radius: 4px;
}
.gallery-container {
    display: flex;
    justify-content: center;
    gap: 20px;
    margin: 20px 0;
}
.image-box {
    text-align: center;
    padding: 10px;
    border-radius: 8px;
    background: #f8f9fa;
}
"""

# สร้าง Gradio Interface
with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""# 🎨 **Image Editor - Text Instructions**
        ### 💬 แบบ Qwen | 💾 ขนาดเล็ก (5.5GB) | ✅ เหมาะสำหรับ RAM 16GB""", 
        elem_id="main-title")
        
        # Comparison box
        gr.HTML("""
        <div class="comparison-box">
            <h3>⚡ InstructPix2Pix - ทางเลือกที่ดีสำหรับ RAM 16GB</h3>
            <p>ใช้วิธีเดียวกับ Qwen: เขียนคำสั่ง → ภาพถูกแก้ไขตามคำสั่ง</p>
        </div>
        """)
        
        with gr.Row():
            # คอลัมน์ซ้าย: อินพุต
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="📤 อัพโหลดภาพต้นฉบับ", 
                    type="pil", 
                    height=350,
                    elem_id="input-image"
                )
                
                instruction = gr.Textbox(
                    label="💬 คำสั่งแก้ไขภาพ (ภาษาอังกฤษ)",
                    placeholder="ตัวอย่าง: 'colorize this manga', 'make it sunset', 'turn into anime style'",
                    lines=3,
                    value="colorize this image"
                )
                
                gr.HTML("""
                <div class="instruction-examples">
                    <strong>📝 ตัวอย่างคำสั่งที่ได้ผลดี:</strong><br><br>
                    • <strong>"colorize this black and white manga"</strong><br>
                    • <strong>"make it look like a painting"</strong><br>
                    • <strong>"add vibrant colors"</strong><br>
                    • <strong>"turn day into night"</strong><br>
                    • <strong>"make it look cyberpunk"</strong><br>
                    • <strong>"add magical sparkles"</strong>
                </div>
                """)
                
                run_button = gr.Button(
                    "✨ แก้ไขภาพ", 
                    variant="primary", 
                    size="lg",
                    elem_id="run-button"
                )
            
            # คอลัมน์ขวา: เอาต์พุตและตั้งค่า
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="✨ ภาพผลลัพธ์", 
                    type="pil", 
                    height=450,
                    elem_id="output-image"
                )
                
                with gr.Accordion("⚙️ การตั้งค่าขั้นสูง", open=True):
                    seed = gr.Slider(
                        label="🎲 Seed", 
                        minimum=0, 
                        maximum=MAX_SEED, 
                        step=1, 
                        value=0,
                        info="ค่า seed เดียวกัน → ผลลัพธ์เดียวกัน"
                    )
                    randomize_seed = gr.Checkbox(
                        label="🔀 สุ่ม Seed อัตโนมัติ", 
                        value=True
                    )
                    
                    text_guidance_scale = gr.Slider(
                        label="💬 Text Guidance Scale", 
                        minimum=1.0, 
                        maximum=20.0, 
                        step=0.5, 
                        value=7.5,
                        info="ยิ่งสูง ยิ่งทำตามคำสั่งมาก"
                    )
                    
                    image_guidance_scale = gr.Slider(
                        label="🖼️ Image Guidance Scale", 
                        minimum=1.0, 
                        maximum=3.0, 
                        step=0.1, 
                        value=1.5,
                        info="ยิ่งสูง ยิ่งรักษาโครงสร้างเดิมมาก (แนะนำ: 1.2-1.8)"
                    )
                    
                    steps = gr.Slider(
                        label="🔢 จำนวน Steps", 
                        minimum=10, 
                        maximum=100, 
                        step=5, 
                        value=30,
                        info="ยิ่งมาก ยิ่งละเอียด แต่ใช้เวลานาน (แนะนำ: 20-40)"
                    )
        
        # ส่วนของคำแนะนำ
        gr.Markdown("""
        ---
        ### 📚 คู่มือการใช้งาน
        
        #### 🎯 **วิธีใช้งาน:**
        1. **อัพโหลดภาพ** - รองรับ PNG, JPG, JPEG
        2. **พิมพ์คำสั่ง** - เป็นภาษาอังกฤษ สั้นๆ ชัดเจน
        3. **ปรับตั้งค่า** (ถ้าต้องการ) - โดยเฉพาะ Image Guidance Scale
        4. **กดปุ่ม "แก้ไขภาพ"** - รอประมาณ 30-60 วินาที
        
        #### ⚙️ **การปรับตั้งค่าที่แนะนำ:**
        
        **สำหรับลงสีมังงะ:**
        - Text Guidance: 7-9
        - Image Guidance: 1.4-1.6
        - Steps: 25-35
        
        **สำหรับเปลี่ยนสไตล์:**
        - Text Guidance: 8-12
        - Image Guidance: 1.2-1.5
        - Steps: 30-40
        
        #### 💡 **เคล็ดลับ:**
        - ใช้คำสั่งภาษาอังกฤษที่ชัดเจน
        - ถ้าผลลัพธ์ไม่ดี ลองเปลี่ยนคำสั่งหรือ seed
        - ภาพขนาด 512x512 pixels จะทำงานได้ดีที่สุด
        - หากหน่วยความจำไม่พอ จะลดขนาดภาพอัตโนมัติ
        
        ---
        
        ⚠️ **หมายเหตุ:** โมเดลนี้ทำงานได้ดีกับภาพทั่วไป แต่สำหรับมังงะอาจต้องทดลองหลายครั้ง
        ⏱️ **เวลาประมวลผล:** 30-90 วินาที (ขึ้นกับจำนวน steps)
        """)
    
    # กำหนด event handler
    run_button.click(
        fn=edit_image,
        inputs=[
            input_image, 
            instruction, 
            seed, 
            randomize_seed, 
            text_guidance_scale, 
            image_guidance_scale, 
            steps
        ],
        outputs=[output_image, seed]
    )

# เรียกใช้งาน
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )