# app.py (ไฟล์นี้จะถูกอัปโหลดไปยัง Hugging Face Space)
import gradio as gr
from PIL import Image
import torch
from diffusers import DiffusionPipeline
import transformers # จำเป็นสำหรับ diffusers
import accelerate # จำเป็นสำหรับ diffusers

# --- ตรวจสอบและตั้งค่า Device ---
# Hugging Face Spaces มักจะตั้งค่า CUDA_VISIBLE_DEVICES ให้ถูกต้องแล้ว
# แต่โค้ดนี้ก็ยังใช้ได้ดี
if torch.cuda.is_available():
    device = "cuda"
    # ใช้ float16 เพื่อประหยัด VRAM บน GPU
    # สำหรับ SDXL แนะนำให้ใช้ variant="fp16" ด้วย
    torch_dtype = torch.float16
    print(f"ใช้ GPU: {torch.cuda.get_device_name(0)}")
else:
    device = "cpu"
    # CPU ไม่รองรับ float16 โดยตรงสำหรับบาง operations
    torch_dtype = torch.float32
    print("ไม่พบ GPU, จะใช้ CPU (อาจช้ามากสำหรับ Stable Diffusion)")

# --- โหลดโมเดล Stable Diffusion ---
# โหลดโมเดลนอกฟังก์ชัน generate เพื่อให้โหลดเพียงครั้งเดียวเมื่อแอปเริ่มทำงาน
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = None # กำหนดค่าเริ่มต้นเป็น None

try:
    print(f"กำลังโหลดโมเดล: {model_id}")
    pipe = DiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        use_safetensors=True,
        # ใช้ variant fp16 ถ้ามี GPU เพื่อลดขนาดโมเดลและ VRAM
        variant="fp16" if device == "cuda" else None
    )
    pipe = pipe.to(device)
    print(f"โหลดโมเดล {model_id} สำเร็จ และย้ายไปที่ {device}")

    # --- Optimizations สำหรับ GPU ---
    if device == "cuda":
        # เปิดใช้งาน memory efficient attention ด้วย xformers (ถ้าติดตั้ง)
        # ช่วยลด VRAM ได้มาก
        try:
            import xformers.ops
            pipe.enable_xformers_memory_efficient_attention()
            print("เปิดใช้งาน xformers memory efficient attention")
        except ImportError:
            print("xformers ไม่ได้ติดตั้งหรือไม่พร้อมใช้งาน, ไม่สามารถเปิดใช้งาน memory efficient attention ได้")
        except Exception as e:
             print(f"เกิดข้อผิดพลาดในการเปิดใช้งาน xformers: {e}")

        # Optional: เปิดใช้งาน attention slicing (อาจช่วยลด VRAM ได้อีก แต่ช้าลงเล็กน้อย)
        # pipe.enable_attention_slicing()
        # print("เปิดใช้งาน attention slicing")

        # Optional: เปิดใช้งาน sequential CPU offload (ลด VRAM ได้มาก แต่ช้าลงมาก)
        # pipe.enable_sequential_cpu_offload()
        # print("เปิดใช้งาน sequential CPU offload")


except Exception as e:
    print(f"เกิดข้อผิดพลาดในการโหลดโมเดล: {e}")
    print("อาจเกิดจาก VRAM ไม่เพียงพอ หรือการเชื่อมต่ออินเทอร์เน็ตมีปัญหา")
    pipe = None # ตั้งค่า pipe เป็น None หากโหลดไม่สำเร็จ

# --- ฟังก์ชันสร้างภาพสำหรับ Gradio ---
def generate_image_for_gradio(prompt: str):
    # ตรวจสอบว่าโมเดลโหลดสำเร็จหรือไม่
    if pipe is None:
        error_message = "โมเดล Stable Diffusion ยังไม่ได้โหลด หรือเกิดข้อผิดพลาดในการโหลด"
        print(error_message)
        # ส่งคืนภาพ placeholder หรือข้อความแสดงข้อผิดพลาด
        # สร้างภาพสีแดงพร้อมข้อความแสดงข้อผิดพลาด
        img = Image.new('RGB', (256, 256), color = (255, 0, 0))
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(img)
        try:
            # ลองโหลดฟอนต์มาตรฐาน
            font = ImageFont.load_default()
        except:
            font = None # ใช้ค่าเริ่มต้นถ้าโหลดฟอนต์ไม่ได้

        text_to_display = "Error: Model not loaded"
        if font:
             # ใช้ textbbox สำหรับ PIL เวอร์ชันใหม่
             try:
                 bbox = draw.textbbox((0, 0), text_to_display, font=font)
                 text_width = bbox[2] - bbox[0]
                 text_height = bbox[3] - bbox[1]
             except:
                  # Fallback สำหรับ PIL เวอร์ชันเก่า
                  text_width, text_height = draw.textsize(text_to_display, font)

             text_x = (256 - text_width) // 2
             text_y = (256 - text_height) // 2
             draw.text((text_x, text_y), text_to_display, fill=(255,255,255), font=font)
        else:
             draw.text((10, 10), text_to_display, fill=(255,255,255))

        return img # ส่งคืนภาพแสดงข้อผิดพลาด

    print(f"\nกำลังสร้างภาพสำหรับ Prompt: \"{prompt}\"")
    print(f"Device: {device}")

    # กำหนดค่าพารามิเตอร์การสร้างภาพ
    # ค่าเหล่านี้ปรับได้ตามต้องการ แต่ค่าเริ่มต้นนี้เหมาะสำหรับ Free GPU
    num_inference_steps = 25 # จำนวนขั้นตอน (20-50)
    guidance_scale = 7.0     # ความเข้มข้นของ prompt (5-15)
    # ลดขนาดภาพลงเพื่อประหยัด VRAM
    height = 768
    width = 768
    # Negative prompt เพื่อควบคุมสิ่งที่ไม่ต้องการ
    negative_prompt = "blurry, low quality, cartoon, watermark, text, ugly, deformed, extra limbs, missing limbs"

    try:
        # เรียก pipeline เพื่อสร้างภาพ
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width
        ).images[0] # .images จะคืนค่าเป็น list, เราต้องการภาพแรก

        print("สร้างภาพสำเร็จ!")
        return image

    except torch.cuda.OutOfMemoryError:
        oom_message = "เกิดข้อผิดพลาด: VRAM ไม่เพียงพอ ลองลดขนาดภาพ หรือใช้โมเดลที่เล็กกว่า"
        print(oom_message)
        # ส่งคืนภาพสีส้มพร้อมข้อความแสดง VRAM Error
        img = Image.new('RGB', (256, 256), color = (255, 165, 0))
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.load_default()
        except:
            font = None
        text_to_display = "Error: Out of VRAM"
        if font:
             try:
                 bbox = draw.textbbox((0, 0), text_to_display, font=font)
                 text_width = bbox[2] - bbox[0]
                 text_height = bbox[3] - bbox[1]
             except:
                  text_width, text_height = draw.textsize(text_to_display, font)

             text_x = (256 - text_width) // 2
             text_y = (256 - text_height) // 2
             draw.text((text_x, text_y), text_to_display, fill=(0,0,0), font=font)
        else:
             draw.text((10, 10), text_to_display, fill=(0,0,0))
        return img

    except Exception as e:
        print(f"เกิดข้อผิดพลาดในการสร้างภาพ: {e}")
        # ส่งคืนภาพสีแดงพร้อมข้อความแสดงข้อผิดพลาดทั่วไป
        img = Image.new('RGB', (256, 256), color = (255, 0, 0))
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.load_default()
        except:
            font = None
        # แสดงข้อความ error บางส่วน
        text_to_display = "Error: " + str(e)[:50] + "..."
        if font:
             try:
                 bbox = draw.textbbox((0, 0), text_to_display, font=font)
                 text_width = bbox[2] - bbox[0]
                 text_height = bbox[3] - bbox[1]
             except:
                  text_width, text_height = draw.textsize(text_to_display, font)

             text_x = (256 - text_width) // 2
             text_y = (256 - text_height) // 2
             draw.text((text_x, text_y), text_to_display, fill=(255,255,255), font=font)
        else:
             draw.text((10, 10), text_to_display, fill=(255,255,255))
        return img


# --- สร้าง Gradio Interface ---
iface = gr.Interface(
    fn=generate_image_for_gradio,
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Stable Diffusion XL Text-to-Image (Hugging Face Free GPU Demo)",
    description="Generate images using Stable Diffusion XL 1.0 on a free Hugging Face GPU Space. May be slow or fail due to resource limits like VRAM.",
    examples=[
        ["A cute cat wearing a hat, cinematic lighting"],
        ["A beautiful sunset over mountains, digital art"],
        ["An epic fantasy landscape, highly detailed, trending on artstation"]
    ]
)

# --- รัน Interface (สำหรับทดสอบ local หรือเมื่อ deploy บน Hugging Face Spaces) ---
if __name__ == "__main__":
    # บน Hugging Face Spaces ไม่ต้องใส่ share=True
    # ถ้าทดสอบ local และต้องการ share ให้ใส่ share=True
    iface.launch()