"""
Скрипт для конвертации модели Qwen из safetensors в формат diffusers
"""
import torch
import os
from diffusers import DiffusionPipeline
from safetensors.torch import load_file
from huggingface_hub import HfApi, create_repo

# Параметры
SAFETENSORS_PATH = r"D:\swarm_Comfy\SwarmUI_Model_Downloader_v82\SwarmUI\Models\unet\QwenSolobandModel\Qwen_Soloband_Model_V1-000140.safetensors"
BASE_MODEL = "Qwen/Qwen-Image"
OUTPUT_DIR = "./qwen_soloband_diffusers"
HF_REPO = "Gerchegg/Qwen-Soloband-Diffusers"
HF_TOKEN = os.environ.get("HF_TOKEN", "hf_YOUR_TOKEN_HERE")  # Замените на ваш токен

def convert_safetensors_to_diffusers():
    """
    Конвертация safetensors модели в формат diffusers
    """
    print("Загрузка базовой модели Qwen-Image...")
    # Загружаем базовую модель со всеми компонентами
    pipeline = DiffusionPipeline.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16
    )
    
    print(f"Загрузка весов из safetensors: {SAFETENSORS_PATH}")
    # Загружаем веса из safetensors
    state_dict = load_file(SAFETENSORS_PATH)
    
    print("Загрузка весов в transformer...")
    # Загружаем веса в transformer компонент
    # Qwen-Image использует DiT (Diffusion Transformer)
    if hasattr(pipeline, 'transformer'):
        # Фильтруем ключи, которые относятся к transformer
        transformer_state_dict = {}
        for key, value in state_dict.items():
            # Убираем префиксы если есть
            new_key = key.replace("model.", "").replace("diffusion_model.", "")
            transformer_state_dict[new_key] = value
        
        # Загружаем веса
        missing_keys, unexpected_keys = pipeline.transformer.load_state_dict(
            transformer_state_dict, 
            strict=False
        )
        
        if missing_keys:
            print(f"Предупреждение: отсутствующие ключи: {missing_keys[:5]}...")
        if unexpected_keys:
            print(f"Предупреждение: неожиданные ключи: {unexpected_keys[:5]}...")
    
    print(f"Сохранение модели в формате diffusers: {OUTPUT_DIR}")
    # Сохраняем в формате diffusers
    pipeline.save_pretrained(OUTPUT_DIR)
    
    print("Конвертация завершена!")
    return OUTPUT_DIR

def upload_to_huggingface(model_dir, repo_id, token):
    """
    Загрузка модели в HuggingFace Hub
    """
    print(f"Создание/проверка репозитория: {repo_id}")
    api = HfApi()
    
    try:
        # Создаем репозиторий если не существует
        create_repo(
            repo_id=repo_id,
            token=token,
            exist_ok=True,
            repo_type="model"
        )
        print("Репозиторий готов")
    except Exception as e:
        print(f"Репозиторий уже существует или ошибка: {e}")
    
    print("Загрузка файлов модели...")
    # Загружаем все файлы из директории
    api.upload_folder(
        folder_path=model_dir,
        repo_id=repo_id,
        token=token,
        repo_type="model"
    )
    
    print(f"Модель успешно загружена в {repo_id}!")
    print(f"Ссылка: https://huggingface.co/{repo_id}")

if __name__ == "__main__":
    # Конвертируем модель
    output_dir = convert_safetensors_to_diffusers()
    
    # Спрашиваем пользователя, загружать ли в HF
    upload = input("\nЗагрузить модель в HuggingFace Hub? (y/n): ")
    if upload.lower() == 'y':
        upload_to_huggingface(output_dir, HF_REPO, HF_TOKEN)
    else:
        print(f"Модель сохранена локально в: {output_dir}")