From d82073fe73968c967e00c6b35911f1b54cbc7014 Mon Sep 17 00:00:00 2001 From: Phil Winder Date: Sat, 23 Nov 2024 22:39:40 +0000 Subject: [PATCH] try that? --- runner/helix-diffusers/pull.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/runner/helix-diffusers/pull.py b/runner/helix-diffusers/pull.py index c3d0d7c1..0303f9fe 100644 --- a/runner/helix-diffusers/pull.py +++ b/runner/helix-diffusers/pull.py @@ -1,7 +1,7 @@ import argparse import os -from diffusers import FluxPipeline, StableDiffusionPipeline +from huggingface_hub import snapshot_download def download_model(model_name: str, save_path: str, pipeline_type: str = "sd"): @@ -13,20 +13,16 @@ def download_model(model_name: str, save_path: str, pipeline_type: str = "sd"): save_path (str): Local directory path to save the model pipeline_type (str): Type of pipeline to use ('sd' for StableDiffusion or 'flux' for Flux) """ - print(f"Downloading model: {model_name} using {pipeline_type} pipeline") + print(f"Downloading model: {model_name}") os.makedirs(save_path, exist_ok=True) - # Select appropriate pipeline based on type - pipeline_class = FluxPipeline if pipeline_type == "flux" else StableDiffusionPipeline - - pipeline_class.from_pretrained( - model_name, + # Download all model files directly without pipeline initialization + snapshot_download( + repo_id=model_name, cache_dir=save_path, - safety_checker=None, - feature_extractor=None, - requires_safety_checker=False, - low_cpu_mem_usage=True, + local_dir=save_path, + ignore_patterns=["*.md", "*.txt"], ) print(f"Model successfully downloaded to: {save_path}")