"""
Скрипт для скачивания и упаковки модели в формат diffusers
Скачивает:
1. VAE от Qwen/Qwen-Image
2. Text Encoder от Qwen/Qwen-Image  
3. Tokenizer от Qwen/Qwen-Image
4. Scheduler config от Qwen/Qwen-Image
5. Кастомный Transformer от Gerchegg/QwenSolobandModel

Упаковывает всё в diffusers формат и сохраняет в Qwen-ImageForFlo_2/
"""

from huggingface_hub import hf_hub_download, snapshot_download, HfApi
from pathlib import Path
import json
import shutil
import os

HF_TOKEN = os.environ.get("HF_TOKEN", "hf_YOUR_TOKEN_HERE")  # Замените на ваш токен

# Пути
TEMP_DIR = "temp_model_download"
SPACE_DIR = "Qwen-ImageForFlo_2"
MODEL_DIR = f"{SPACE_DIR}/model"

# Создаем директории
Path(TEMP_DIR).mkdir(exist_ok=True)
Path(MODEL_DIR).mkdir(parents=True, exist_ok=True)

print("""
╔══════════════════════════════════════════════════════════════╗
║  СКАЧИВАНИЕ И УПАКОВКА МОДЕЛИ                               ║
╚══════════════════════════════════════════════════════════════╝

📥 Будет скачано:
   1. VAE (~0.3GB) от Qwen/Qwen-Image
   2. Text Encoder (~17GB) от Qwen/Qwen-Image
   3. Tokenizer (~5MB) от Qwen/Qwen-Image
   4. Scheduler config от Qwen/Qwen-Image
   5. Кастомный Transformer (~40.9GB) от Gerchegg/QwenSolobandModel
   
📦 Упаковка в diffusers формат:
   Целевая папка: {MODEL_DIR}/
   
⏱️  Примерное время: 10-20 минут (зависит от скорости интернета)
💾 Размер: ~58GB
""")

input("Нажмите Enter для начала скачивания...")

print("\n" + "="*60)
print("ШАГ 1: Скачивание базовой модели Qwen-Image")
print("="*60)

# Скачиваем всю базовую модель
print("\n📥 Скачивание полной модели Qwen/Qwen-Image...")
print("   Это займет 10-15 минут...")

base_model_path = snapshot_download(
    repo_id="Qwen/Qwen-Image",
    local_dir=f"{TEMP_DIR}/Qwen-Image",
    token=HF_TOKEN,
    allow_patterns=["*.json", "*.txt", "*.md", "*.model", "vae/*", "text_encoder/*", "tokenizer/*", "scheduler/*"]
)

print(f"✓ Базовая модель скачана в: {base_model_path}")

print("\n" + "="*60)
print("ШАГ 2: Скачивание кастомного Transformer")
print("="*60)

print("\n📥 Скачивание Qwen_Soloband_Model_V1-000140.safetensors (40.9GB)...")
print("   Это займет 5-10 минут...")

custom_transformer_path = hf_hub_download(
    repo_id="Gerchegg/QwenSolobandModel",
    filename="Qwen_Soloband_Model_V1-000140.safetensors",
    local_dir=f"{TEMP_DIR}/custom",
    token=HF_TOKEN
)

print(f"✓ Кастомный transformer скачан: {custom_transformer_path}")

print("\n" + "="*60)
print("ШАГ 3: Упаковка в diffusers формат")
print("="*60)

print("\n📦 Создание структуры diffusers...")

# Создаем структуру папок
for subfolder in ['vae', 'text_encoder', 'tokenizer', 'scheduler', 'transformer']:
    Path(f"{MODEL_DIR}/{subfolder}").mkdir(parents=True, exist_ok=True)

# Копируем компоненты
print("  Копирование VAE...")
shutil.copytree(
    f"{TEMP_DIR}/Qwen-Image/vae",
    f"{MODEL_DIR}/vae",
    dirs_exist_ok=True
)

print("  Копирование Text Encoder...")
shutil.copytree(
    f"{TEMP_DIR}/Qwen-Image/text_encoder",
    f"{MODEL_DIR}/text_encoder",
    dirs_exist_ok=True
)

print("  Копирование Tokenizer...")
shutil.copytree(
    f"{TEMP_DIR}/Qwen-Image/tokenizer",
    f"{MODEL_DIR}/tokenizer",
    dirs_exist_ok=True
)

print("  Копирование Scheduler...")
shutil.copytree(
    f"{TEMP_DIR}/Qwen-Image/scheduler",
    f"{MODEL_DIR}/scheduler",
    dirs_exist_ok=True
)

# Копируем кастомный transformer
print("  Копирование кастомного Transformer...")
shutil.copy(
    custom_transformer_path,
    f"{MODEL_DIR}/transformer/diffusion_pytorch_model.safetensors"
)

# Копируем model_index.json
print("  Копирование конфигураций...")
if os.path.exists(f"{TEMP_DIR}/Qwen-Image/model_index.json"):
    shutil.copy(
        f"{TEMP_DIR}/Qwen-Image/model_index.json",
        f"{MODEL_DIR}/model_index.json"
    )

print("\n✓ Модель упакована в diffusers формат!")

print("\n" + "="*60)
print("ШАГ 4: Создание .gitattributes для LFS")
print("="*60)

gitattributes_content = """*.safetensors filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.gguf filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tar.gz filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
"""

with open(f"{SPACE_DIR}/.gitattributes", 'w') as f:
    f.write(gitattributes_content)

print("✓ .gitattributes создан")

print("\n" + "="*60)
print("✅ ПОДГОТОВКА ЗАВЕРШЕНА!")
print("="*60)

print(f"\n📂 Структура модели в {MODEL_DIR}/:")
print(f"   ├── vae/")
print(f"   ├── text_encoder/")
print(f"   ├── tokenizer/")
print(f"   ├── scheduler/")
print(f"   ├── transformer/")
print(f"   │   └── diffusion_pytorch_model.safetensors (40.9GB)")
print(f"   └── model_index.json")

print(f"\n📊 Размер модели:")
total_size = sum(f.stat().st_size for f in Path(MODEL_DIR).rglob('*') if f.is_file())
print(f"   Всего: {total_size / 1024**3:.1f} GB")

print(f"\n🎯 Следующие шаги:")
print(f"   1. Проверьте содержимое {MODEL_DIR}/")
print(f"   2. Запустите create_space_app.py для создания app.py")
print(f"   3. Загрузите в HF с помощью upload_to_space.py")

print(f"\n💾 Временные файлы в: {TEMP_DIR}/")
print(f"   Можно удалить после успешной загрузки")