nunchaku - shut up

This commit is contained in:
Haoming 2025-08-01 12:49:46 +08:00
parent 90261d9293
commit bef4565cd2
2 changed files with 34 additions and 19 deletions

View File

@ -114,3 +114,6 @@ class SVDQFluxTransformer2DModel(nn.Module):
self._prev_timestep = timestep_float
return out
def load_state_dict(self, *args, **kwargs):
return [], []

View File

@ -1,7 +1,6 @@
import os
import sys
INITIALIZED = False
MONITOR_MODEL_MOVING = False
@ -17,7 +16,7 @@ def monitor_module_moving():
def new_to(*args, **kwargs):
traceback.print_stack()
print('Model Movement')
print("Model Movement")
return old_to(*args, **kwargs)
@ -25,6 +24,21 @@ def monitor_module_moving():
return
def fix_logging():
import logging
_original = logging.basicConfig
logging.basicConfig = lambda *args, **kwargs: None
try:
import nunchaku
except ImportError:
pass
logging.basicConfig = _original
def initialize_forge():
global INITIALIZED
@ -33,16 +47,17 @@ def initialize_forge():
INITIALIZED = True
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), 'modules_forge', 'packages'))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), "modules_forge", "packages"))
from backend.args import args
if args.gpu_device_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_device_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_device_id)
print("Set device to:", args.gpu_device_id)
if args.cuda_malloc:
from modules_forge.cuda_malloc import try_cuda_malloc
try_cuda_malloc()
from backend import memory_management
@ -55,29 +70,26 @@ def initialize_forge():
memory_management.soft_empty_cache()
from backend import stream
print('CUDA Using Stream:', stream.should_use_stream())
print("CUDA Using Stream:", stream.should_use_stream())
from modules_forge.shared import diffusers_dir
# if 'TRANSFORMERS_CACHE' not in os.environ:
# os.environ['TRANSFORMERS_CACHE'] = diffusers_dir
if "HF_HOME" not in os.environ:
os.environ["HF_HOME"] = diffusers_dir
if 'HF_HOME' not in os.environ:
os.environ['HF_HOME'] = diffusers_dir
if "HF_DATASETS_CACHE" not in os.environ:
os.environ["HF_DATASETS_CACHE"] = diffusers_dir
if 'HF_DATASETS_CACHE' not in os.environ:
os.environ['HF_DATASETS_CACHE'] = diffusers_dir
if "HUGGINGFACE_HUB_CACHE" not in os.environ:
os.environ["HUGGINGFACE_HUB_CACHE"] = diffusers_dir
if 'HUGGINGFACE_HUB_CACHE' not in os.environ:
os.environ['HUGGINGFACE_HUB_CACHE'] = diffusers_dir
if "HUGGINGFACE_ASSETS_CACHE" not in os.environ:
os.environ["HUGGINGFACE_ASSETS_CACHE"] = diffusers_dir
if 'HUGGINGFACE_ASSETS_CACHE' not in os.environ:
os.environ['HUGGINGFACE_ASSETS_CACHE'] = diffusers_dir
if 'HF_HUB_CACHE' not in os.environ:
os.environ['HF_HUB_CACHE'] = diffusers_dir
if "HF_HUB_CACHE" not in os.environ:
os.environ["HF_HUB_CACHE"] = diffusers_dir
import modules_forge.patch_basic
modules_forge.patch_basic.patch_all_basics()
return
fix_logging()