app.py

import gradio as gr
import os
import subprocess
import requests
from huggingface_hub import snapshot_download, login

# Ссылка на официальный скрипт конвертации SDXL от Diffusers
SCRIPT_URL = "https://raw.githubusercontent.com/huggingface/diffusers/main/scripts/convert_diffusers_to_original_sdxl.py"
SCRIPT_NAME = "convert_script.py"

def download_script():
    if not os.path.exists(SCRIPT_NAME):
        print("Скачивание скрипта конвертации...")
        response = requests.get(SCRIPT_URL)
        with open(SCRIPT_NAME, "wb") as f:
            f.write(response.content)

def convert_model(repo_id, token, fp16_save):
    output_file = "model.safetensors"
    
    # 0. Подготовка
    download_script()
    if token.strip():
        login(token=token.strip())
        
    status_msg = ""
    
    try:
        # 1. Скачивание репозитория (diffusers format) в локальную папку
        print(f"Скачивание файлов модели {repo_id}...")
        local_dir = snapshot_download(
            repo_id=repo_id,
            token=token if token else None,
            ignore_patterns=["*.msgpack", "*.bin", "*.h5"] # Игнорируем ненужное для экономии места
        )
        print(f"Модель скачана в: {local_dir}")

        # 2. Формирование команды для скрипта
        # python convert_script.py --model_path "путь" --checkpoint_path "файл" --use_safetensors
        cmd = [
            "python", SCRIPT_NAME,
            "--model_path", local_dir,
            "--checkpoint_path", output_file,
            "--use_safetensors"
        ]
        
        if fp16_save:
            cmd.append("--half") # Флаг для конвертации в FP16
            
        print("Запуск конвертации...", cmd)
        
        # 3. Запуск конвертации
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            return None, f"Ошибка скрипта:\n{result.stderr}"
            
        # Проверка результата
        if os.path.exists(output_file):
            size_gb = os.path.getsize(output_file) / (1024**3)
            return output_file, f"Успешно! Файл создан ({size_gb:.2f} GB)."
        else:
            return None, "Файл не был создан по неизвестной причине."

    except Exception as e:
        return None, f"Критическая ошибка: {str(e)}"

# Интерфейс
with gr.Blocks() as demo:
    gr.Markdown("# SDXL Diffusers -> Safetensors (Script Method)")
    gr.Markdown("Этот метод использует официальный скрипт конвертации Diffusers. Он надежнее и требует меньше памяти.")
    
    with gr.Row():
        repo_input = gr.Textbox(label="Repo ID", placeholder="rockeycoss/SPO-SDXL_4k-p_10ep", value="rockeycoss/SPO-SDXL_4k-p_10ep")
        token_input = gr.Textbox(label="HF Token (опционально)", type="password")
    
    is_fp16 = gr.Checkbox(label="Конвертировать в FP16", value=True, info="Рекомендуется для SDXL (размер ~6.5GB вместо 13GB).")
    
    btn = gr.Button("Скачать и конвертировать", variant="primary")
    out_file = gr.File(label="Готовый файл")
    out_log = gr.Textbox(label="Лог", lines=5)
    
    btn.click(convert_model, inputs=[repo_input, token_input, is_fp16], outputs=[out_file, out_log])

demo.queue().launch()