app.py

import torch, random, gc
import numpy as np
import gradio as gr
from diffusers import FluxPipeline
from PIL import Image

# --- 1. Setup For CPU Only (คุม RAM 16GB) ---
device = "cpu"
model_id = "black-forest-labs/FLUX.1-schnell" 
# บน CPU ต้องใช้ float32 หรือ bfloat16 (ถ้า CPU รองรับ) เพื่อประหยัด RAM
torch_dtype = torch.bfloat16 

if 'pipe' not in locals():
    print("🛡️ กำลังโหลดโหมด CPU Low-RAM... (เป้าหมาย 13-14.5GB)")
    
    # โหลดแบบไม่เอาโมเดลขึ้น RAM ทีเดียวทั้งหมด
    pipe = FluxPipeline.from_pretrained(
        model_id, 
        torch_dtype=torch_dtype,
        # ห้ามใช้ device_map="auto" บน CPU-Only Space เพราะจะ Error
    )

    # --- 🛑 หัวใจสำคัญของการคุม RAM 13-14GB ---
    # บังคับให้โหลดเฉพาะส่วนที่ใช้งาน ถ้าส่วนไหนไม่ใช้ให้โยนทิ้งไป
    pipe.enable_model_cpu_offload() 
    pipe.enable_vae_tiling()
    pipe.enable_vae_slicing()
    
    gc.collect()

# --- 2. Logic Function (คุมเข้มการรัน) ---
def generate_safe(prompt, seed, steps, width, height):
    # เคลียร์ขยะก่อนรันเสมอ
    gc.collect()
    
    actual_seed = int(seed) if int(seed) != -1 else random.randint(0, 2**32 - 1)
    generator = torch.Generator(device="cpu").manual_seed(actual_seed)
    
    print(f"🚀 เริ่มเจนภาพ: {width}x{height} | Steps: {steps}")
    
    image = pipe(
        prompt=prompt,
        guidance_scale=0.0, 
        num_inference_steps=int(steps), 
        width=int(width),
        height=int(height),
        max_sequence_length=128, # สำคัญมาก! ห้ามเพิ่ม เพื่อคุม RAM
        generator=generator
    ).images[0]
    
    return image, actual_seed

# --- 3. UI ---
with gr.Blocks() as demo:
    gr.Markdown(f"### 🛡️ FLUX CPU-Only (RAM Safe Mode: 13-14.5GB Target)")
    with gr.Row():
        p_in = gr.Textbox(label="Prompt", value="a simple cat")
        btn = gr.Button("🚀 Generate")
    out = gr.Image()
    
    # ส่งค่าคงที่เพื่อคุมโหลดของ RAM
    btn.click(
        fn=generate_safe, 
        inputs=[
            p_in, 
            gr.State(-1), # Seed
            gr.State(4),  # Steps (Schnell ใช้แค่ 4 พอ)
            gr.State(512),# Width
            gr.State(512) # Height
        ], 
        outputs=[out, gr.State(0)]
    )

demo.launch()