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())