import gradio as gr
import torch
from diffusers import DiffusionPipeline
from peft import PeftModel
import shutil
import os

# --- Cấu hình mô hình ---
MODEL_NAME = "Qwen/Qwen-Image"
OUTPUT_DIR = "./merged_checkpoint"

def merge_loras(lora_a_id, lora_b_id):
    # Xóa thư mục cũ nếu tồn tại
    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print(f"✅ Tải mô hình gốc Qwen-Image từ Hugging Face...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

    try:
        pipe = DiffusionPipeline.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch_dtype,
            use_safetensors=True,
            safety_checker=None,
            feature_extractor=None,
            trust_remote_code=True
        ).to(device)
        print("✅ Tải mô hình gốc thành công!")
    except Exception as e:
        return f"❌ Lỗi khi tải mô hình gốc: {str(e)}"

    # --- Hợp nhất LoRA A ---
    print(f"🔧 Hợp nhất LoRA A: {lora_a_id}")
    try:
        pipe.unet = PeftModel.from_pretrained(pipe.unet, lora_a_id)
        pipe.unet.merge_and_unload()
        print("✅ Hợp nhất LoRA A thành công")
    except Exception as e:
        return f"❌ Lỗi khi hợp nhất LoRA A ({lora_a_id}): {str(e)}"

    # --- Hợp nhất LoRA B ---
    print(f"🔧 Hợp nhất LoRA B: {lora_b_id}")
    try:
        pipe.unet = PeftModel.from_pretrained(pipe.unet, lora_b_id)
        pipe.unet.merge_and_unload()
        print("✅ Hợp nhất LoRA B thành công")
    except Exception as e:
        return f"❌ Lỗi khi hợp nhất LoRA B ({lora_b_id}): {str(e)}"

    # --- Lưu checkpoint hoàn chỉnh ---
    try:
        pipe.save_pretrained(OUTPUT_DIR, safe_serialization=True, max_shard_size="2GB")
        print(f"✅ Lưu checkpoint vào: {OUTPUT_DIR}")
    except Exception as e:
        return f"❌ Lỗi khi lưu checkpoint: {str(e)}"

    # --- Nén để download ---
    shutil.make_archive(OUTPUT_DIR, 'zip', OUTPUT_DIR)
    output_zip = OUTPUT_DIR + ".zip"

    return output_zip

# --- Giao diện Gradio ---
with gr.Blocks(title="🎨 Qwen-Image: Insw + NSFW LoRA Merger") as demo:
    gr.Markdown("""
    # 🎨 Qwen-Image: Merge Insw + NSFW LoRA (Online Tool)
    
    Hợp nhất hai LoRA nổi bật trên Hugging Face:
    - **`rorito/Insw`**: Phong cách chân dung Trung Hoa cổ điển, tinh tế, ánh sáng dịu.
    - **`StuffedPumpkins/nsfw_qwen_image`**: Thêm yếu tố NSFW (gợi cảm, thân thể con người) vào Qwen-Image.
    
    ✅ Kết quả: **Chân dung Trung Hoa cổ điển mang phong cách gợi cảm**, vẫn giữ khả năng **render văn bản siêu chuẩn** (chữ Hán, π≈3.14, biển hiệu...).
    
    🔹 Không cần download gì cả  
    🔹 Chạy hoàn toàn trên cloud (miễn phí)  
    🔹 Xuất ra file `.zip` dùng ngay trong ComfyUI / Automatic1111  
    🔹 Tối ưu cho prompt chứa văn bản phức tạp
    
    ---
    ### 📥 Cách dùng:
    1. Nhấn “🚀 Hợp nhất LoRA” (mặc định đã điền sẵn 2 LoRA bên dưới).
    2. Sau vài phút, hệ thống sẽ tạo ra file `.zip`.
    3. Tải về → giải nén → copy `merged_checkpoint/unet/diffusion_pytorch_model.safetensors` vào thư mục model của WebUI.
    4. Dùng prompt ví dụ:  
       `"A beautiful Chinese woman in traditional hanfu, standing beside a neon sign reading '通义千问', soft candlelight, cinematic composition, ultra HD, 4K"`  
       → Văn bản Trung và phong cách gợi cảm đều hiện ra chính xác!
    
    > Model gốc: [Qwen-Image on Hugging Face](https://huggingface.co/Qwen/Qwen-Image)  
    > LoRA A: [rorito/Insw](https://huggingface.co/rorito/Insw)  
    > LoRA B: [StuffedPumpkins/nsfw_qwen_image](https://huggingface.co/StuffedPumpkins/nsfw_qwen_image)
    """)

    with gr.Row():
        with gr.Column():
            lora_a_input = gr.Textbox(label="📌 LoRA A (Chinese Portrait)", value="rorito/Insw", interactive=True)
            lora_b_input = gr.Textbox(label="📌 LoRA B (NSFW)", value="StuffedPumpkins/nsfw_qwen_image", interactive=True)
            merge_btn = gr.Button("🚀 Hợp nhất LoRA", variant="primary")
        
        with gr.Column():
            output_file = gr.File(label="📥 Kết quả: Checkpoint đã hợp nhất (.zip)")

    merge_btn.click(fn=merge_loras, inputs=[lora_a_input, lora_b_input], outputs=output_file)

demo.launch()