app.py

import spaces
import os
import sys 
import subprocess
import importlib.util
from peft import PeftModel

# --- CÀI ĐẶT ÉP BUỘC XTUNER (Bỏ qua kiểm tra xung đột) ---
if importlib.util.find_spec("xtuner") is None:
    print("Đang cài đặt xtuner từ GitHub...")
    subprocess.check_call([
        sys.executable, "-m", "pip", "install",
        "git+https://github.com/InternLM/xtuner.git@v0.2.0",
        "--no-deps"
    ])

import torch
import torch.utils._pytree as _torch_pytree

# 1. Vá lỗi PyTree cho Transformers mới
def smart_pytree_patch():
    orig_func = getattr(_torch_pytree, '_register_pytree_node', None)
    if orig_func:
        def patched_register(cls, to_iter, from_iter, serialized_type_name=None):
            return orig_func(cls, to_iter, from_iter)
        _torch_pytree.register_pytree_node = patched_register
        _torch_pytree._register_pytree_node = patched_register

smart_pytree_patch()

import torch.distributed

# 2. Vá lỗi Torch XPU
if not hasattr(torch, 'xpu'):
    class DummyXPU:
        @staticmethod
        def is_available(): return False
        @staticmethod
        def empty_cache(): pass
        @staticmethod
        def device_count(): return 0
        @staticmethod
        def current_device(): return 0
        @staticmethod
        def get_device_name(device=None): return "DummyXPU"
        @staticmethod
        def is_bf16_supported(): return False
        @staticmethod
        def synchronize(device=None): pass
        @staticmethod
        def set_device(device): pass
        @staticmethod
        def manual_seed(seed): pass
        @staticmethod
        def manual_seed_all(seed): pass
        @staticmethod
        def seed(): pass
        @staticmethod
        def seed_all(): pass
        @staticmethod
        def initial_seed(): return 0
    torch.xpu = DummyXPU()

# 3. Vá lỗi Device Mesh
if not hasattr(torch.distributed, 'device_mesh'):
    class DummyDeviceMesh: pass
    class DummyDeviceMeshModule: DeviceMesh = DummyDeviceMesh
    dummy_module = DummyDeviceMeshModule()
    sys.modules['torch.distributed.device_mesh'] = dummy_module
    torch.distributed.device_mesh = dummy_module

# 4. Vá lỗi pad_sequence() không hỗ trợ padding_side (cần PyTorch >= 2.5)
# ---------------------------------------------------------------
import torch.nn.utils.rnn as _rnn
_orig_pad_sequence = _rnn.pad_sequence

def _patched_pad_sequence(sequences, batch_first=False, padding_value=0.0, padding_side='right'):
    if padding_side == 'left':
        sequences = [seq.flip(0) for seq in sequences]
        padded = _orig_pad_sequence(sequences, batch_first=batch_first, padding_value=padding_value)
        flip_dim = 1 if batch_first else 0
        return padded.flip(flip_dim)
    else:
        return _orig_pad_sequence(sequences, batch_first=batch_first, padding_value=padding_value)

_rnn.pad_sequence = _patched_pad_sequence
torch.nn.utils.rnn.pad_sequence = _patched_pad_sequence

import importlib, sys as _sys
def _patch_pad_sequence_in_module(module_name):
    mod = _sys.modules.get(module_name)
    if mod and hasattr(mod, 'pad_sequence'):
        mod.pad_sequence = _patched_pad_sequence

import builtins as _builtins
_orig_builtins_import = _builtins.__import__

def _import_hook(name, *args, **kwargs):
    module = _orig_builtins_import(name, *args, **kwargs)
    for mod_name, mod in list(_sys.modules.items()):
        if mod_name.startswith('src.') and hasattr(mod, 'pad_sequence'):
            mod.pad_sequence = _patched_pad_sequence
    return module

_builtins.__import__ = _import_hook
print("✅ Đã patch pad_sequence() để hỗ trợ padding_side")

# 5. Vá lỗi F.interpolate bilinear nhận 3D input thay vì 4D
# ---------------------------------------------------------------
import torch.nn.functional as _F
_orig_interpolate = _F.interpolate

def _patched_interpolate(input, size=None, scale_factor=None, mode='nearest',
                          align_corners=None, recompute_scale_factor=None, antialias=False):
    squeezed = False
    if mode in ('bilinear', 'bicubic') and input.dim() == 3:
        input = input.unsqueeze(0)
        squeezed = True
    result = _orig_interpolate(
        input, size=size, scale_factor=scale_factor, mode=mode,
        align_corners=align_corners, recompute_scale_factor=recompute_scale_factor,
        antialias=antialias
    )
    if squeezed:
        result = result.squeeze(0)
    return result

_F.interpolate = _patched_interpolate
torch.nn.functional.interpolate = _patched_interpolate
print("✅ Đã patch F.interpolate() để hỗ trợ 3D input với bilinear mode")

# 6. Vá lỗi vae.encode() / vae.decode() nhận 3D input thay vì 4D
# ---------------------------------------------------------------
try:
    from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL as _AutoencoderKL
    _orig_vae_encode = _AutoencoderKL.encode
    _orig_vae_decode = _AutoencoderKL.decode

    def _patched_vae_encode(self, x, *args, **kwargs):
        if x.dim() == 3:
            x = x.unsqueeze(0)
        # FIX: Bỏ phần cắt channels - giữ nguyên để VAE decode đúng màu
        result = _orig_vae_encode(self, x, *args, **kwargs)
        return result

    def _patched_vae_decode(self, z, *args, **kwargs):
        if z.dim() == 3:
            z = z.unsqueeze(0)
        return _orig_vae_decode(self, z, *args, **kwargs)

    _AutoencoderKL.encode = _patched_vae_encode
    _AutoencoderKL.decode = _patched_vae_decode
    print("✅ Đã patch AutoencoderKL.encode/decode() - FIXED")
except Exception as _e:
    print(f"⚠️ Không thể patch AutoencoderKL: {_e}")

# 7. Vá lỗi dtype string + flash_attention_2
# ---------------------------------------------------------------
STRING_TO_TORCH_DTYPE = {
    "float16": torch.float16,
    "float32": torch.float32,
    "float64": torch.float64,
    "bfloat16": torch.bfloat16,
    "torch.bfloat16": torch.bfloat16,
    "half": torch.float16,
    "int8": torch.int8,
}

def str_to_torch_dtype(dtype):
    if isinstance(dtype, str):
        return STRING_TO_TORCH_DTYPE.get(dtype.lower(), torch.float32)
    return dtype

try:
    from transformers import PretrainedConfig
    original_pretrained_init = PretrainedConfig.__init__

    def patched_pretrained_config_init(self, *args, **kwargs):
        if 'torch_dtype' in kwargs and isinstance(kwargs['torch_dtype'], str):
            kwargs['torch_dtype'] = str_to_torch_dtype(kwargs['torch_dtype'])
        if kwargs.get('attn_implementation') == 'flash_attention_2':
            kwargs['attn_implementation'] = 'sdpa'
        original_pretrained_init(self, *args, **kwargs)
        if hasattr(self, 'torch_dtype') and isinstance(self.torch_dtype, str):
            self.torch_dtype = str_to_torch_dtype(self.torch_dtype)
        if getattr(self, '_attn_implementation', None) == 'flash_attention_2':
            self._attn_implementation = 'sdpa'
        if getattr(self, '_attn_implementation_internal', None) == 'flash_attention_2':
            self._attn_implementation_internal = 'sdpa'

    PretrainedConfig.__init__ = patched_pretrained_config_init
    print("✅ Đã patch PretrainedConfig.__init__")
except Exception as e:
    print(f"⚠️ Không thể patch PretrainedConfig: {e}")

orig_is_floating_point = torch.is_floating_point
def patched_is_floating_point(obj):
    if isinstance(obj, str):
        return obj.lower() in ["bfloat16", "float16", "float32", "float64", "half"]
    return orig_is_floating_point(obj)
torch.is_floating_point = patched_is_floating_point

# ============================================================
# IMPORTS CHÍNH
# ============================================================
import time
import psutil
import numpy as np
import gradio as gr
from PIL import Image
from einops import rearrange
from huggingface_hub import hf_hub_download

from xtuner.registry import BUILDER
from mmengine.config import Config

# ============================================================
# CONFIG PATCHES
# ============================================================
LOCAL_TO_HF_PATH = {
    "model_zoo/Qwen2.5-VL-3B-Instruct":   "Qwen/Qwen2.5-VL-3B-Instruct",
    "model_zoo/UniPic2-SD3.5M-Kontext-2B": "Skywork/UniPic2-SD3.5M-Kontext-GRPO-2B",
}

def patch_config_paths(cfg):
    cfg_text = cfg.pretty_text
    changed = False
    for local_path, hf_path in LOCAL_TO_HF_PATH.items():
        if local_path in cfg_text:
            cfg_text = cfg_text.replace(local_path, hf_path)
            print(f"  → Đã thay path: '{local_path}' → '{hf_path}'")
            changed = True
    if changed:
        return Config.fromstring(cfg_text, file_format='.py')
    return cfg

def patch_config_dtype(cfg):
    if isinstance(cfg, dict):
        for key in list(cfg.keys()):
            val = cfg[key]
            if key in ('torch_dtype', 'dtype', 'param_dtype', 'compute_dtype') and isinstance(val, str):
                cfg[key] = str_to_torch_dtype(val)
                print(f"  → Convert dtype cfg['{key}'] = '{val}' → {cfg[key]}")
            elif key == 'attn_implementation' and val == 'flash_attention_2':
                cfg[key] = 'sdpa'
                print(f"  → Thay attn_implementation: flash_attention_2 → sdpa")
            else:
                patch_config_dtype(val)
    elif isinstance(cfg, (list, tuple)):
        for item in cfg:
            patch_config_dtype(item)
    elif hasattr(cfg, '_cfg_dict'):
        patch_config_dtype(cfg._cfg_dict)
    return cfg

from xtuner.model.utils import guess_load_checkpoint

REPO_ID = "deepgenteam/DeepGen-1.0"
MODEL_WEIGHTS = {
    "Pretrain (Alignment)": "iter_200000.pth",
    "RL with MR-GRPO (Tốt nhất)": "model.pt"
}

current_loaded_method = None
model = None

def load_model(method_name):
    global current_loaded_method, model
    if current_loaded_method == method_name and model is not None:
        return model

    filename = MODEL_WEIGHTS[method_name]
    weight_path = hf_hub_download(repo_id=REPO_ID, filename=filename)
    config = Config.fromfile("configs/models/deepgen_scb.py")
    config = patch_config_paths(config)
    patch_config_dtype(config)
    new_model = BUILDER.build(config.model)

    if weight_path.endswith('.pt'):
        state_dict = torch.load(weight_path, map_location="cpu", weights_only=False)
    else:
        state_dict = guess_load_checkpoint(weight_path)

    # Load đơn giản như code gốc — KHÔNG tách lmm/non-lmm
    missing, unexpected = new_model.load_state_dict(state_dict, strict=False)
    print(f"  ✅ Loaded - Missing: {len(missing)}, Unexpected: {len(unexpected)}")
    if missing:
        print(f"     Missing sample: {missing[:5]}")

    model_dtype = new_model.dtype
    if isinstance(model_dtype, str):
        model_dtype = str_to_torch_dtype(model_dtype)
    new_model = new_model.to(model_dtype)
    new_model.eval()
    model = new_model
    current_loaded_method = method_name
    return model

def _process_image(image):
    """
    Letterbox resize: giữ nguyên tỉ lệ gốc, padding trung tính (xám 128)
    để vừa 512x512. Không crop, không stretch.
    """
    image = image.convert('RGB')
    orig_w, orig_h = image.size
    target_size = 512

    # Tính scale để cạnh dài nhất = 512
    scale = target_size / max(orig_w, orig_h)
    new_w = int(orig_w * scale)
    new_h = int(orig_h * scale)

    # Resize giữ tỉ lệ
    image_resized = image.resize((new_w, new_h), Image.LANCZOS)

    # Canvas xám trung tính (128 ≈ 0.0 sau normalize → model thấy "nền trống")
    canvas = Image.new('RGB', (target_size, target_size), (128, 128, 128))
    pad_left = (target_size - new_w) // 2
    pad_top  = (target_size - new_h) // 2
    canvas.paste(image_resized, (pad_left, pad_top))

    pixel_values = torch.from_numpy(np.array(canvas)).float()
    pixel_values = pixel_values / 255.0
    pixel_values = 2.0 * pixel_values - 1.0
    pixel_values = rearrange(pixel_values, 'h w c -> c h w')
    return pixel_values

# ============================================================
# HELPER: ĐO RAM CPU & VRAM GPU
# ============================================================
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 _log_resources(label: str):
    ram          = _get_ram_mb()
    vram_alloc   = _get_vram_mb()
    vram_reserved = _get_vram_reserved_mb()
    print(f"  📊 [{label}]"
          f"  RAM: {ram:.0f} MB"
          f"  |  VRAM alloc: {vram_alloc:.0f} MB"
          f"  |  VRAM reserved: {vram_reserved:.0f} MB")

# ============================================================
# INFERENCE
# ============================================================
@spaces.GPU(duration=1200)
def run_inference(task_type, prompt, cfg_prompt, cfg_scale, num_steps, seed, method, src_img=None):
    t_total_start = time.perf_counter()
    mode_label    = "Text-to-Image" if task_type == "t2i" else "Image-Editing"

    print(f"\n{'='*60}")
    print(f"🚀 BẮT ĐẦU [{mode_label}]  steps={num_steps}  cfg={cfg_scale}  seed={seed}")
    print(f"   Prompt: {prompt[:100]}{'...' if len(prompt) > 100 else ''}")
    print(f"{'='*60}")

    # Reset peak VRAM counter cho lần chạy này
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()

    _log_resources("Khởi đầu")

    try:
        # ── 1. Load model ──────────────────────────────────────
        t0  = time.perf_counter()
        net = load_model(method)
        net = net.to("cuda")
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        t_load = time.perf_counter() - t0
        print(f"  ⏱️  Load/move model to CUDA : {t_load:.2f}s")
        _log_resources("Sau load model")

        generator   = torch.Generator(device=net.device).manual_seed(int(seed))
        prompts     = [prompt.strip()]
        # cfg_prompts = [cfg_prompt] - debug1
        cfg_prompts = [cfg_prompt if cfg_prompt else ""]

        # ── 2. Tiền xử lý ảnh nguồn (chỉ i2i) ────────────────
        t_img = 0.0
        pixel_values_src = None
        orig_w, orig_h = 512, 512
        pad_left, pad_top, new_w, new_h = 0, 0, 512, 512
        if task_type == "i2i" and src_img is not None:
            t0 = time.perf_counter()
            orig_w, orig_h = src_img.size
            scale  = 512 / max(orig_w, orig_h)
            new_w  = int(orig_w * scale)
            new_h  = int(orig_h * scale)
            pad_left = (512 - new_w) // 2
            pad_top  = (512 - new_h) // 2
            processed_img    = _process_image(src_img).to(net.dtype).to(net.device)
            # Thêm unsqueeze(0) lần 2 để tạo chiều [Num_Images]
            pixel_values_src = processed_img.unsqueeze(0).unsqueeze(0)
            t_img = time.perf_counter() - t0
            print(f"  ⏱️  Tiền xử lý ảnh nguồn  : {t_img:.3f}s  (letterbox {new_w}x{new_h} pad {pad_left},{pad_top})")
            _log_resources("Sau tiền xử lý ảnh")
        
        # ── 3. Generate (bước nặng) ────────────────────────────
        _log_resources("Trước generate")
        t0 = time.perf_counter()
        with torch.no_grad():
            images = net.generate(
                prompt=prompts,
                cfg_prompt=cfg_prompts,
                pixel_values_src=pixel_values_src,
                cfg_scale=cfg_scale,
                num_steps=int(num_steps),
                progress_bar=False,
                generator=generator,
                height=512,
                width=512
            )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        t_gen = time.perf_counter() - t0
        print(f"  ⏱️  Generate ({num_steps} steps)     : {t_gen:.2f}s  "
              f"({t_gen / int(num_steps) * 1000:.1f} ms/step)")
        _log_resources("Sau generate")

        # Debug: kiểm tra range và distribution của output
        print(f"  🔍 images min={images.min().item():.3f} max={images.max().item():.3f} mean={images.mean().item():.3f}")
        print(f"  🔍 images shape={images.shape} dtype={images.dtype}")
        # Kiểm tra từng channel R,G,B có khác nhau không
        print(f"  🔍 R mean={images[0,0].mean().item():.3f} G mean={images[0,1].mean().item():.3f} B mean={images[0,2].mean().item():.3f}")
                
        # ── 4. Post-process ────────────────────────────────────
        t0 = time.perf_counter()
        
        if task_type == "i2i":
            images = torch.clamp(images, -1.0, 1.0)
            images = (images + 1.0) / 2.0
            images = rearrange(images, 'b c h w -> b h w c')
            images_np = torch.clamp(images * 255.0, 0, 255).to("cpu", dtype=torch.uint8).numpy()
            # Crop vùng content ra (bỏ padding letterbox), resize về kích thước gốc
            result_pil = Image.fromarray(images_np[0])
            result_pil = result_pil.crop((pad_left, pad_top, pad_left + new_w, pad_top + new_h))
            result_pil = result_pil.resize((orig_w, orig_h), Image.LANCZOS)
        else:
            images = rearrange(images, 'b c h w -> b h w c')
            images_np = torch.clamp(127.5 * images + 128.0, 0, 255).to("cpu", dtype=torch.uint8).numpy()
            result_pil = Image.fromarray(images_np[0])
        
        t_post = time.perf_counter() - t0
        print(f"  ⏱️  Post-process             : {t_post:.3f}s")

        # ── 5. Offload về CPU ──────────────────────────────────
        net = net.to("cpu")
        torch.cuda.empty_cache()

        # ── Tổng kết ──────────────────────────────────────────
        t_total    = time.perf_counter() - t_total_start
        vram_peak  = torch.cuda.max_memory_allocated() / 1024 / 1024 if torch.cuda.is_available() else 0
        vram_now   = _get_vram_mb()
        ram_now    = _get_ram_mb()

        sep = "=" * 60
        print(f"\n{sep}")
        print(f"✅ KẾT QUẢ [{mode_label}]")
        print(f"   ┌─ Thời gian ──────────────────────────────")
        print(f"   │  Tổng              : {t_total:.2f}s")
        print(f"   │  Load model        : {t_load:.2f}s")
        if task_type == "i2i" and src_img is not None:
            print(f"   │  Tiền xử lý ảnh   : {t_img:.3f}s")
        print(f"   │  Generate          : {t_gen:.2f}s  ({t_gen/int(num_steps)*1000:.1f} ms/step)")
        print(f"   │  Post-process      : {t_post:.3f}s")
        print(f"   ├─ Bộ nhớ ────────────────────────────────")
        print(f"   │  RAM CPU hiện tại  : {ram_now:.0f} MB")
        print(f"   │  VRAM peak         : {vram_peak:.0f} MB")
        print(f"   │  VRAM hiện tại     : {vram_now:.0f} MB")
        print(f"   └──────────────────────────────────────────")
        print(f"{sep}\n")

        return result_pil

    except Exception as e:
        import traceback
        t_total = time.perf_counter() - t_total_start
        print(f"\n❌ LỖI sau {t_total:.2f}s [{mode_label}]: {str(e)}")
        traceback.print_exc()
        return None

# --- GIAO DIỆN GRADIO ---
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 DeepGen-1.0 Demo trên ZeroGPU")

    method_dropdown = gr.Dropdown(
        choices=list(MODEL_WEIGHTS.keys()),
        value="RL with MR-GRPO (Tốt nhất)",
        label="Cấu hình Model / Weights"
    )

    with gr.Tabs():
        with gr.Tab("🖼️ Text-to-Image"):
            with gr.Row():
                with gr.Column():
                    t2i_prompt     = gr.Textbox(label="Prompt", value="A quiet bookstore with a sign that says 'READ'. A coffee cup on the table with the word 'MORNING'.")
                    # t2i_cfg_prompt = gr.Textbox(label="Negative / CFG Prompt", value="")
                    t2i_cfg_prompt = gr.Textbox(
                        label="CFG Prompt (để trống = unconditional guidance)", 
                        value=""
                    )
                    t2i_scale      = gr.Slider(1.0, 10.0, value=4.0, label="CFG Scale")
                    t2i_steps      = gr.Slider(10, 100, value=50, step=1, label="Steps")
                    t2i_seed       = gr.Number(value=42, label="Seed")
                    t2i_btn        = gr.Button("Tạo ảnh", variant="primary")
                with gr.Column():
                    t2i_output = gr.Image(label="Kết quả")
            t2i_btn.click(
                fn=lambda *args: run_inference("t2i", *args),
                inputs=[t2i_prompt, t2i_cfg_prompt, t2i_scale, t2i_steps, t2i_seed, method_dropdown],
                outputs=t2i_output
            )

        with gr.Tab("🎨 Image Editing"):
            with gr.Row():
                with gr.Column():
                    i2i_src        = gr.Image(label="Ảnh nguồn (Source Image)", type="pil")
                    i2i_prompt     = gr.Textbox(label="Editing Prompt", value="Transform the image to watercolor painting style")
                    # i2i_cfg_prompt = gr.Textbox(label="Negative Prompt", value="blurry, distorted faces, extra limbs, text errors, low quality, oversaturated")
                    i2i_cfg_prompt = gr.Textbox(
                        label="CFG Prompt (để trống = unconditional guidance)", 
                        value=""   # ← empty string là chuẩn
                    )
                    i2i_scale      = gr.Slider(1.0, 10.0, value=4.0, label="CFG Scale")
                    i2i_steps      = gr.Slider(10, 100, value=20, step=1, label="Steps")
                    i2i_seed       = gr.Number(value=42, label="Seed")
                    i2i_btn        = gr.Button("Chỉnh sửa ảnh", variant="primary")
                with gr.Column():
                    i2i_output = gr.Image(label="Kết quả chỉnh sửa")
            i2i_btn.click(
                fn=lambda src, p, cfg, s, st, sd, m: run_inference("i2i", p, cfg, s, st, sd, m, src),
                inputs=[i2i_src, i2i_prompt, i2i_cfg_prompt, i2i_scale, i2i_steps, i2i_seed, method_dropdown],
                outputs=i2i_output
            )

# KHÔNG XÓA DÒNG NÀY
demo.launch(theme=gr.themes.Soft())