Upload 2 files

download_utils.py

@@ -0,0 +1,341 @@
+import os
+import re
+import shutil
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from urllib.parse import unquote, urlparse
+
+import fal
+from fal.toolkit.utils.download_utils import (
+    FAL_MODEL_WEIGHTS_DIR,
+    DownloadError,
+    _hash_url,
+)
+
+FAL_VERSION = getattr(fal, "__version__", "<1.0.0")
+_REQUEST_HEADERS = {"User-Agent": f"fal-client ({FAL_VERSION}/python)"}
+
+
+def get_civitai_headers() -> dict[str, str]:
+    headers: dict[str, str] = {}
+
+    civitai_token = os.getenv("CIVITAI_TOKEN", None)
+
+    if not civitai_token:
+        print("CIVITAI_TOKEN is not set in the environment variables.")
+        return headers
+
+    headers["Authorization"] = f"Bearer {civitai_token}"
+
+    return headers
+
+
+def get_huggingface_headers() -> dict[str, str]:
+    headers: dict[str, str] = {}
+
+    hf_token = os.getenv("HF_TOKEN", None)
+
+    if not hf_token:
+        print("HF_TOKEN is not set in the environment variables.")
+        return headers
+
+    headers["Authorization"] = f"Bearer {hf_token}"
+
+    return headers
+
+
+def get_local_file_content_length(file_path: Path) -> int:
+    return file_path.stat().st_size
+
+
+def download_url_to_file(
+    url: str,
+    dst: str | Path,
+    progress: bool = True,
+    headers: dict[str, str] = None,
+    chunk_size_in_mb=16,
+    file_integrity_check_callback=None,
+) -> Path:
+    """Download object at the given URL to a local path.
+
+    Args:
+        url (str): URL of the object to download
+        dst (str): Full path where object will be saved, e.g. ``/tmp/temporary_file``
+        progress (bool, optional): whether or not to display a progress bar to stderr
+            Default: True
+        headers (dict, optional): HTTP headers to include with the request
+            Default: None
+        chunk_size_in_mb (int, optional): size of each chunk in MB
+            Default: 16
+        file_integrity_check_callback (callable, optional): callback function to check file integrity
+            Default: None
+
+    """
+    from tqdm import tqdm
+
+    file_size = None
+
+    request_headers = {
+        **_REQUEST_HEADERS,
+        **(headers or {}),
+    }
+
+    url = url.strip()
+
+    if url.startswith("data:"):
+        return _download_data_url_to_file(url, dst)
+
+    import requests
+
+    req = requests.get(url, headers=request_headers, stream=True, allow_redirects=True)
+    req.raise_for_status()
+
+    headers = req.headers  # type: ignore
+    content_length = headers.get("Content-Length", None)  # type: ignore
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        file_path = temp_file.name
+        try:
+            with tqdm(
+                total=file_size,
+                disable=not progress,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as pbar, open(file_path, "wb") as f:
+                for chunk in req.iter_content(
+                    chunk_size=chunk_size_in_mb * 1024 * 1024
+                ):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+
+            # Move the file when the file is downloaded completely. Since the
+            # file used is temporary, in a case of an interruption, the downloaded
+            # content will be lost. So, it is safe to redownload the file in such cases.
+            shutil.move(file_path, dst)
+        except Exception as error:
+            raise error
+        finally:
+            Path(temp_file.name).unlink(missing_ok=True)
+
+    if file_integrity_check_callback:
+        file_integrity_check_callback(dst)
+
+    return Path(dst)
+
+
+def _download_data_url_to_file(url: str, dst: str | Path):
+    import base64
+
+    data = url.split(",")[1]
+    data = base64.b64decode(data)
+
+    with open(dst, "wb") as fp:
+        fp.write(data)
+
+    return Path(dst)
+
+
+def download_model_weights(url: str, force: bool = False) -> Path:
+    parsed_url = urlparse(url)
+
+    headers = {}
+    if parsed_url.netloc == "civitai.com":
+        headers.update(get_civitai_headers())
+    elif parsed_url.netloc == "huggingface.co":
+        headers.update(get_huggingface_headers())
+
+    return download_model_weights_fal(url, request_headers=headers, force=force)
+
+def url_without_query(url)->str:
+    # 找到 '?' 的位置
+    query_index = url.find('?')
+
+    # 如果找到 '?', 则截取到 '?' 之前的部分
+    if query_index != -1:
+        url_without_query = url[:query_index]
+    else:
+        url_without_query = url  # 如果没有查询参数，保持原 URL
+    return url_without_query
+
+def download_model_weights_fal(
+    url: str, force: bool = False, request_headers: dict[str, str] | None = None
+) -> Path:
+    without_query = os.environ.get("CKPT_DOWNLOAD_WITHOUT_QUERY", "false") == "true"
+    if without_query:
+        url = url_without_query(url)
+        
+    weights_dir = Path(FAL_MODEL_WEIGHTS_DIR / _hash_url(url))
+
+    if weights_dir.exists() and not force:
+        try:
+            weights_path = next(weights_dir.glob("*"))
+            is_safetensors_file(weights_path)
+            return weights_path
+
+        # The model weights directory is empty, so we need to download the weights
+        except StopIteration:
+            pass
+
+    try:
+        file_name, file_content_length = _get_remote_file_properties(
+            url, request_headers=request_headers
+        )
+    except Exception as e:
+        print(e)
+        raise DownloadError(f"Failed to get remote file properties for {url}")    
+
+    target_path = weights_dir / file_name
+
+    if (
+        target_path.exists()
+        and get_local_file_content_length(target_path) == file_content_length
+        and not force
+    ):
+        is_safetensors_file(target_path)
+        return target_path
+
+    # Make sure the parent directory exists
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    # download from network-volume
+    ckpt_download_dir = os.environ.get("CKPT_DOWNLOAD_DIR", None)
+    if ckpt_download_dir:
+        src_path = os.path.join(ckpt_download_dir, file_name)
+        if not os.path.exists(src_path):
+            src_path = os.path.join(ckpt_download_dir, _hash_url(url), file_name)
+        if not os.path.exists(src_path):
+            src_path = os.path.join(ckpt_download_dir, _hash_url(url_without_query(url)), file_name)
+            
+        # 如果文件存在则复制到路径，并返回
+        if os.path.exists(src_path):
+            src_file_content_length = get_local_file_content_length(Path(src_path))
+            if src_file_content_length == file_content_length:
+                print(f"copy start from:{src_path} to:{target_path}")
+                shutil.copy(src_path, target_path)
+                print(f"copy done  from:{src_path} to:{target_path}")
+                return target_path
+            else:
+                print(f"cannot copy file length not same src_path:{src_path} len:{src_file_content_length} target_path:{target_path} len:{file_content_length}")
+
+    try:
+        download_url_to_file(
+            url,
+            target_path,
+            progress=True,
+            headers=request_headers,
+            file_integrity_check_callback=is_safetensors_file,
+        )
+    except Exception as e:
+        print(e)
+        raise DownloadError(f"Failed to download {url}")
+
+    return target_path
+
+
+def _get_filename_from_content_disposition(cd: str | None) -> str | None:
+    if not cd:
+        return None
+
+    filenames = re.findall('filename="(.+)"', cd)
+
+    if len(filenames) == 0:
+        filenames = re.findall("filename=(.+)", cd)
+
+    if len(filenames) == 0:
+        return None
+
+    return unquote(filenames[0])
+
+
+def _parse_filename(url: str, cd: str | None) -> str:
+    url = url.strip()
+    file_name = _get_filename_from_content_disposition(cd)
+
+    if not file_name:
+        parsed_url = urlparse(url)
+
+        if parsed_url.scheme == "data":
+            file_name = _hash_url(url)
+        else:
+            url_path = parsed_url.path
+            file_name = Path(url_path).name or _hash_url(url)
+
+    if url.startswith("data:"):
+        import mimetypes
+
+        mime_type = url.split(",")[0].split(":")[1].split(";")[0]
+        extension = mimetypes.guess_extension(mime_type)
+        if extension:
+            file_name += extension
+
+    return file_name  # type: ignore
+
+
+def _get_remote_file_properties(
+    url: str, request_headers: dict[str, str] = None
+) -> tuple[str, int]:
+    import requests
+
+    headers = {
+        **_REQUEST_HEADERS,
+        **(request_headers or {}),
+    }
+
+    req = requests.get(
+        url, headers=headers, stream=True, allow_redirects=True, verify=False
+    )
+    req.raise_for_status()
+
+    headers = req.headers  # type: ignore
+    content_disposition = headers.get("Content-Disposition", None)
+    file_name = _parse_filename(url, content_disposition)
+    content_length = int(headers.get("Content-Length", -1))
+
+    return file_name, content_length
+
+
+def is_safetensors_file(path: str | Path):
+    from safetensors import safe_open
+
+    path = str(path)
+
+    if not path.endswith(".safetensors"):
+        raise ValueError(f"File {path} is not a .safetensors file")
+
+    try:
+        with safe_open(path, framework="pt"):
+            pass
+    except Exception as e:
+        print(e)
+        error_mesage = e.args[0]
+        if error_mesage == "Error while deserializing header: HeaderTooLarge":
+            raise ValueError(f"File {path} is not a .safetensors file")
+        else:
+            raise e
+
+
+@contextmanager
+def download_file_temp(
+    url: str,
+    progress: bool = True,
+    headers: dict[str, str] = None,
+    chunk_size_in_mb=16,
+    file_integrity_check_callback=None,
+):
+    file_name = _parse_filename(url, None)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        file_path = download_url_to_file(
+            url,
+            f"{temp_dir}/{file_name}",
+            progress=progress,
+            headers=headers,
+            chunk_size_in_mb=chunk_size_in_mb,
+            file_integrity_check_callback=file_integrity_check_callback,
+        )
+        yield file_path

start.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+
+# Use libtcmalloc for better memory management
+TCMALLOC="$(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1)"
+export LD_PRELOAD="${TCMALLOC}"
+
+echo "Worker Initiated"
+
+# Serve the API and don't shutdown the container
+if [ "$SERVE_API_LOCALLY" == "true" ]; then
+    echo "runpod-worker-comfy: Starting ComfyUI"
+    python3 /comfyui/main.py --disable-auto-launch --disable-metadata --listen &
+
+    echo "runpod-worker-comfy: Starting RunPod Handler"
+    python3 -u /rp_handler.py --rp_serve_api --rp_api_host=0.0.0.0
+else
+    echo "runpod-worker-comfy: Starting ComfyUI"
+    python3 /comfyui/main.py --disable-auto-launch --disable-metadata &
+
+    echo "runpod-worker-comfy: Starting RunPod Handler"
+    python3 -u /rp_handler.py
+fi