mirror of
https://github.com/vladmandic/automatic
synced 2026-09-05 12:30:46 +02:00
Vladimir Mandic
79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
import os
|
|
import time
|
|
from modules.logger import log
|
|
|
|
|
|
debug = log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
|
|
orig_http_get = None
|
|
orig_xet_get = None
|
|
orig_build_hf_headers = None
|
|
|
|
|
|
def clean_user_agent(headers):
|
|
ua = headers.get('user-agent', None)
|
|
if ua is None:
|
|
return headers
|
|
ua = ua.replace('unknown/None;', 'sdnext/0.0.0;')
|
|
if 'quant_config' in ua:
|
|
ua = ua.split('quant_config')[0]
|
|
ua = ua.rstrip(' ;')
|
|
# ua += '; telemetry/off'
|
|
headers['user-agent'] = ua
|
|
return headers
|
|
|
|
|
|
def http_get_hijack(*args, **kwargs):
|
|
from modules.shared import state, opts
|
|
if len(args) > 0 and isinstance(args[0], str) and args[0].endswith(".json"):
|
|
return orig_http_get(*args, **kwargs)
|
|
jobid = state.begin('Download')
|
|
fn = kwargs.get("displayed_filename", None)
|
|
size = kwargs.get("expected_size", None)
|
|
if 'headers' in kwargs:
|
|
kwargs['headers'] = clean_user_agent(kwargs['headers'])
|
|
if fn and not fn.endswith(".json") and size is not None and size > 10240:
|
|
log.debug(f'Download: type=http mode="{opts.hf_transfer_mode}" fn="{fn}" size={size}')
|
|
debug(f'Download start: type=http args={args} kwargs={kwargs}')
|
|
t0 = time.time()
|
|
res = orig_http_get(*args, **kwargs)
|
|
t1 = time.time()
|
|
debug(f'Download end: type=http res={res} time={t1-t0:.2f} perf={size/(t1-t0)/1024/1024:.2f} MB/s')
|
|
state.end(jobid)
|
|
return res
|
|
|
|
|
|
def xet_get_hijack(*args, **kwargs):
|
|
from modules.shared import state, opts
|
|
if len(args) > 0 and isinstance(args[0], str) and args[0].endswith(".json"):
|
|
return orig_xet_get(*args, **kwargs)
|
|
jobid = state.begin('Download')
|
|
fn = kwargs.get("displayed_filename", None)
|
|
size = kwargs.get("expected_size", None)
|
|
if 'headers' in kwargs:
|
|
kwargs['headers'] = clean_user_agent(kwargs['headers'])
|
|
if fn and not fn.endswith(".json"):
|
|
log.debug(f'Download: type=xet mode="{opts.hf_transfer_mode}" fn="{fn}" size={size}')
|
|
debug(f'Download start: type=xet args={args} kwargs={kwargs}')
|
|
res = orig_xet_get(*args, **kwargs)
|
|
debug(f'Download end: type=xet res={res}')
|
|
state.end(jobid)
|
|
return res
|
|
|
|
|
|
def build_hf_headers_hijack(*args, **kwargs):
|
|
headers = orig_build_hf_headers(*args, **kwargs)
|
|
headers = clean_user_agent(headers)
|
|
return headers
|
|
|
|
|
|
def init_hijack():
|
|
from huggingface_hub import file_download
|
|
global orig_http_get, orig_xet_get, orig_build_hf_headers # pylint: disable=global-statement
|
|
if orig_http_get is None or orig_xet_get is None:
|
|
orig_http_get = file_download.http_get
|
|
orig_xet_get = file_download.xet_get
|
|
orig_build_hf_headers = file_download.build_hf_headers # pylint: disable=protected-access
|
|
file_download.http_get = http_get_hijack
|
|
file_download.xet_get = xet_get_hijack
|
|
file_download.build_hf_headers = build_hf_headers_hijack # pylint: disable=protected-access
|