download_qwen_image_models.py

#!/usr/bin/env python3
import os
from typing import Dict
from huggingface_hub import hf_hub_download

# Enable hf_transfer for faster downloads
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

DEFAULT_MODELS_DIR = os.environ.get("QWEN_IMAGE_MODELS_DIR", 
                                    "/Qwen-Image_models")
# Temporary download root (requested: use /tmp instead of /workspace)
TMP_DOWNLOAD_ROOT = os.environ.get("QWEN_IMAGE_TMP_DIR", "/tmp/qie_downloads")


def _ensure_dirs(root: str) -> None:
    os.makedirs(root, exist_ok=True)
    os.makedirs(os.path.join(root, "dit"), exist_ok=True)
    os.makedirs(os.path.join(root, "vae"), exist_ok=True)
    os.makedirs(os.path.join(root, "text_encoder"), exist_ok=True)


def _download_then_place(*, repo_id: str, filename: str, subfolder: str,
                         component: str, models_dir: str) -> str:
    """Download to /tmp and move to the final models_dir/component/filename.

    Returns absolute path to the final placed file.
    """
    # Ensure temp and final directories
    tmp_dir = os.path.join(TMP_DOWNLOAD_ROOT, component)
    os.makedirs(tmp_dir, exist_ok=True)
    final_dir = os.path.join(models_dir, component)
    os.makedirs(final_dir, exist_ok=True)

    print(f"[QIE] Downloading {component}: {repo_id}/{subfolder}/{filename}")
    tmp_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        subfolder=subfolder,
        local_dir=tmp_dir,
    )

    final_path = os.path.join(final_dir, filename)
    try:
        # If target exists, replace it
        if os.path.exists(final_path):
            try:
                os.remove(final_path)
            except FileNotFoundError:
                pass
        # Move downloaded file into place
        if os.path.abspath(tmp_path) != os.path.abspath(final_path):
            os.replace(tmp_path, final_path)
    except Exception:
        # Fallback: copy if replace fails (e.g., cross-device)
        import shutil
        shutil.copy2(tmp_path, final_path)
    return final_path


def download_all_models(models_dir: str = DEFAULT_MODELS_DIR) -> Dict[str, str]:
    """Download required Qwen-Image-Edit models into models_dir.

    Returns a dict of component -> local file path.
    """
    _ensure_dirs(models_dir)

    print(f"[QIE] Models directory: {models_dir}")

    print("[QIE] Download dir (tmp):", TMP_DOWNLOAD_ROOT)
    print("[QIE] Final models dir:", models_dir)

    # Download to /tmp then move to final path
    dit_path = _download_then_place(
        repo_id="Comfy-Org/Qwen-Image-Edit_ComfyUI",
        filename="qwen_image_edit_2509_bf16.safetensors",
        subfolder="split_files/diffusion_models",
        component="dit",
        models_dir=models_dir,
    )

    print("[QIE] Downloading VAE model(s)...")
    vae_main = _download_then_place(
        repo_id="Qwen/Qwen-Image-Edit",
        filename="diffusion_pytorch_model.safetensors",
        subfolder="vae",
        component="vae",
        models_dir=models_dir,
    )
    vae_alt = _download_then_place(
        repo_id="Comfy-Org/Qwen-Image_ComfyUI",
        filename="qwen_image_vae.safetensors",
        subfolder="split_files/vae",
        component="vae",
        models_dir=models_dir,
    )

    print("[QIE] Downloading Text Encoder...")
    te_path = _download_then_place(
        repo_id="Comfy-Org/Qwen-Image_ComfyUI",
        filename="qwen_2.5_vl_7b.safetensors",
        subfolder="split_files/text_encoders",
        component="text_encoder",
        models_dir=models_dir,
    )

    print("[QIE] All models downloaded successfully!")
    return {
        "dit": dit_path,
        "vae_main": vae_main,
        "vae_alt": vae_alt,
        "text_encoder": te_path,
        "root": models_dir,
    }


if __name__ == "__main__":
    download_all_models()