app.py

import gradio as gr
import torch
import random
import psutil
import os
from diffusers import Step1XEditPipelineV1P2
from RegionE import RegionEHelper
# --- BẮT ĐẦU ĐOẠN VÁ LỖI CỦA REGIONE ---
import RegionE.Step1XEditV1P2.inplace as inplace_v1p2
# Ghi đè biến False của tác giả thành None để ép nó dùng SDPA (không cần flash-attn)
inplace_v1p2.flash_attn = None

def _get_ram_mb():
    return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024

def _get_vram_mb():
    return torch.cuda.memory_allocated() / 1024 / 1024 if torch.cuda.is_available() else 0.0

def _get_vram_reserved_mb():
    return torch.cuda.memory_reserved() / 1024 / 1024 if torch.cuda.is_available() else 0.0

def _get_vram_peak_mb():
    # Mức VRAM cao nhất từng đạt được kể từ khi chạy/reset
    return torch.cuda.max_memory_allocated() / 1024 / 1024 if torch.cuda.is_available() else 0.0

def _log_resources(label: str):
    ram          = _get_ram_mb()
    vram_alloc   = _get_vram_mb()
    vram_reserved = _get_vram_reserved_mb()
    vram_peak    = _get_vram_peak_mb() # Đã sửa lỗi xóa chữ 'cls.'
    print(f"📊 [{label}] RAM: {ram:.0f} MB | VRAM alloc: {vram_alloc:.0f} MB | VRAM res: {vram_reserved:.0f} MB | Peak: {vram_peak:.0f} MB")


_log_resources("Trước khi load Model")
# 1. Khởi tạo Pipeline ở global scope để tối ưu thời gian load
model_id = "stepfun-ai/Step1X-Edit-v1p2"
print(f"Loading pipeline from {model_id}...")

pipe = Step1XEditPipelineV1P2.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

# 2. Khởi tạo RegionEHelper
regionehelper = RegionEHelper(pipe)
regionehelper.set_params() # Sử dụng hyperparameters mặc định
print("Pipeline and RegionEHelper loaded successfully!")

# 3. Hàm xử lý inference, bọc bởi @spaces.GPU để lấy GPU động trên HF
_log_resources("Sau khi load Model & RegionE")
def process_image(image, prompt, num_inference_steps, cfg_scale, seed, enable_thinking, enable_reflection):
    if image is None:
        raise gr.Error("Vui lòng tải lên một hình ảnh!")
    if not prompt:
        raise gr.Error("Vui lòng nhập câu lệnh (prompt) để chỉnh sửa!")
    _log_resources("Bắt đầu process_image (Chuẩn bị dữ liệu)")
    # Xử lý Seed
    if seed == -1 or seed is None:
        seed = random.randint(0, 2**32 - 1)
    
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    image = image.convert("RGB")

    # Bật RegionE để tối ưu tốc độ
    regionehelper.enable()

    try:
        # Chạy inference
        _log_resources("Trước khi gọi Pipeline Inference")
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()
        pipe_output = pipe(
            image=image,
            prompt=prompt,
            num_inference_steps=int(num_inference_steps),
            true_cfg_scale=cfg_scale,
            generator=generator,
            enable_thinking_mode=enable_thinking,
            enable_reflection_mode=enable_reflection,
        )
        _log_resources("Sau khi gọi Pipeline Inference")
    finally:
        # Tắt RegionE sau khi inference xong để dọn dẹp
        regionehelper.disable()

    # Thu thập kết quả
    if enable_reflection:
        final_image = pipe_output.final_images[0]
    else:
        final_image = pipe_output.images[0]

    # Thu thập log/thông tin thinking
    log_text = f"Seed sử dụng: {int(seed)}\n" + ("-" * 40) + "\n"
    
    if enable_thinking and hasattr(pipe_output, 'reformat_prompt'):
        log_text += f"Reformat Prompt:\n{pipe_output.reformat_prompt}\n"
        log_text += ("-" * 40) + "\n"
        
    if enable_reflection and hasattr(pipe_output, 'think_info'):
        log_text += f"Think Info:\n{pipe_output.think_info[0]}\n"
        log_text += ("-" * 40) + "\n"
        log_text += f"Best Info:\n{pipe_output.best_info[0]}\n"
    _log_resources("Hoàn thành process_image")
    return final_image, log_text

# 4. Xây dựng giao diện Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🪄 Step1X-Edit-v1p2 Image Editing Demo
        Demo model chỉnh sửa ảnh [Step1X-Edit-v1p2](https://huggingface.co/stepfun-ai/Step1X-Edit-v1p2) được tăng tốc bởi [RegionE](https://github.com/Peyton-Chen/RegionE).
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Hình ảnh đầu vào (Input Image)", type="pil")
            prompt = gr.Textbox(
                label="Câu lệnh chỉnh sửa (Prompt)", 
                placeholder="VD: add a ruby pendant on the girl's neck.",
                lines=3
            )
            
            with gr.Accordion("Cài đặt nâng cao (Advanced Settings)", open=False):
                num_steps = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps")
                cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, value=6.0, step=0.5, label="CFG Scale")
                seed = gr.Number(value=-1, label="Seed (-1 để random)", precision=0)
                
                enable_thinking = gr.Checkbox(value=True, label="Enable Thinking Mode")
                enable_reflection = gr.Checkbox(value=True, label="Enable Reflection Mode")
                
            submit_btn = gr.Button("Tạo ảnh (Generate)", variant="primary")
            
        with gr.Column(scale=1):
            output_image = gr.Image(label="Kết quả (Generated Image)")
            output_log = gr.Textbox(label="Logs & Thinking Process", lines=12, interactive=False)

    submit_btn.click(
        fn=process_image,
        inputs=[
            input_image, 
            prompt, 
            num_steps, 
            cfg_scale, 
            seed, 
            enable_thinking, 
            enable_reflection
        ],
        outputs=[output_image, output_log]
    )

if __name__ == "__main__":
    demo.queue().launch()