import os
import shutil
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
from huggingface_hub import hf_hub_download, HfApi

def convert_and_push(source_repo, file_name, model_type, target_repo, hf_token):
    # 檢查 Token 是否為空
    if not hf_token:
        return "❌ 錯誤：請輸入具有寫入權限的 Hugging Face Token！"
    
    # 定義暫存目錄，轉換完會刪除
    output_dir = "temp_diffusers_model"
    api = HfApi()
    
    try:
        # --- 第一步：從雲端下載檔案 ---
        gr.Info(f"正在從 {source_repo} 下載 {file_name}...")
        # hf_hub_download 會把檔案抓到 Space 的快取暫存區
        local_ckpt_path = hf_hub_download(
            repo_id=source_repo.strip(), 
            filename=file_name.strip()
        )
        
        # --- 第二步：判斷模型版本並選擇對應的轉換引擎 ---
        # SD 1.5/2.1 與 SDXL 的結構不同，所以必須切換類別
        if model_type == "SDXL":
            gr.Info("偵測到 SDXL，使用 StableDiffusionXLPipeline 轉換...")
            pipeline_class = StableDiffusionXLPipeline
        else:
            gr.Info("使用 StableDiffusionPipeline 轉換 (SD1.5/2.1)...")
            pipeline_class = StableDiffusionPipeline
            
        # --- 第三步：執行核心轉換 ---
        # from_single_file 是一個強大的方法，能把單一 .safetensors 拆解成 Diffusers 資料夾結構
        # load_safety_checker=False: 關閉安全檢查器以節省 VRAM/RAM
        # torch_dtype=torch.float32: 在免費 CPU Space 建議用 float32，最穩定
        pipe = pipeline_class.from_single_file(
            local_ckpt_path, 
            load_safety_checker=False,
            torch_dtype=torch.float32 
        )
        
        # 將轉換好的各個子元件（unet, vae, tokenizer 等）存入暫存資料夾
        pipe.save_pretrained(output_dir)
        
        # --- 第四步：推送到你的 Hugging Face 倉庫 ---
        gr.Info(f"正在推送到目標倉庫: {target_repo}...")
        
        # 1. 自動建立目標倉庫 (如果已經存在則會忽略)
        api.create_repo(
            repo_id=target_repo.strip(), 
            token=hf_token, 
            exist_ok=True, 
            repo_type="model"
        )
        
        # 2. 將整個資料夾上傳，這會自動完成 Diffusers 的標準檔案佈局
        api.upload_folder(
            folder_path=output_dir,
            repo_id=target_repo.strip(),
            token=hf_token,
            repo_type="model"
        )
        
        # --- 第五步：清理暫存，釋放硬碟空間 ---
        shutil.rmtree(output_dir)
        
        return f"✅ 成功完成！\n模型已轉換並推送到：https://huggingface.co{target_repo}"
        
    except Exception as e:
        # 發生錯誤時也嘗試清理暫存，避免硬碟爆掉
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        return f"❌ 出錯了：{str(e)}"

# --- Gradio UI 排版 ---
with gr.Blocks(theme=gr.themes.Soft(), title="SD 轉 Diffusers 工具") as demo:
    gr.Markdown("# 🚀 SD to Diffusers 雲端自動推送工具")
    gr.Markdown("本工具能將單一檔案模型 (.safetensors) 直接轉換為 Diffusers 格式並推送到你的個人 Repo。")
    
    with gr.Row():
        # 左邊區塊：來源模型資訊
        with gr.Column():
            gr.Markdown("### 📥 1. 來源模型 (Source)")
            src_repo = gr.Textbox(label="來源 Repo ID", placeholder="作者/倉庫名稱")
            src_file = gr.Textbox(label="來源檔案名稱", placeholder="包含副檔名，如 model.safetensors")
            m_type = gr.Radio(["SD1.5 / SD2.1", "SDXL"], label="模型版本", value="SD1.5 / SD2.1")
            
        # 右邊區塊：目標位置與金鑰
        with gr.Column():
            gr.Markdown("### 📤 2. 目標設定 (Target)")
            tgt_repo = gr.Textbox(label="你的目標 Repo ID", placeholder="你的名稱/新倉庫名稱")
            hf_token = gr.Textbox(label="Hugging Face Write Token", placeholder="需有 Write 權限的 Token", type="password")
            
    # 執行按鈕
    run_btn = gr.Button("🔥 開始轉換並推送", variant="primary")
    
    # 狀態回饋
    status_output = gr.Textbox(label="執行進度與狀態", lines=5)

    # 綁定按鈕動作
    run_btn.click(
        fn=convert_and_push, 
        inputs=[src_repo, src_file, m_type, tgt_repo, hf_token], 
        outputs=status_output
    )

demo.launch()