unify hf-login

modules/model_sd3.py

@@ -1,6 +1,4 @@
-import io
 import os
-import contextlib
 import warnings
 import torch
 import diffusers
@@ -13,23 +11,9 @@ warnings.filterwarnings(action="ignore", category=FutureWarning)
 loggedin = False
 
 
-def hf_login():
-    global loggedin # pylint: disable=global-statement
-    import huggingface_hub as hf
-    from modules import shared
-    if shared.opts.huggingface_token is not None and len(shared.opts.huggingface_token) > 2 and not loggedin:
-        stdout = io.StringIO()
-        with contextlib.redirect_stdout(stdout):
-            hf.login(shared.opts.huggingface_token)
-        text = stdout.getvalue() or ''
-        line = [l for l in text.split('\n') if 'Token' in l]
-        shared.log.info(f'HF login: {line[0] if len(line) > 0 else text}')
-        loggedin = True
-
-
 def load_sd3(fn=None, cache_dir=None, config=None):
-    from modules import devices
-    hf_login()
+    from modules import devices, modelloader
+    modelloader.hf_login()
     repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
     model_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
     dtype = torch.float16
@@ -88,8 +72,8 @@ def load_sd3(fn=None, cache_dir=None, config=None):
 
 
 def load_te3(pipe, te3=None, cache_dir=None):
-    from modules import devices
-    hf_login()
+    from modules import devices, modelloader
+    modelloader.hf_login()
     repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
     if pipe is None or not hasattr(pipe, 'text_encoder_3'):
         return pipe

modules/modeldata.py

@@ -97,6 +97,8 @@ class Shared(sys.modules[__name__].__class__):
                 model_type = 'sd' # sd is compatible with sd
             elif "Kandinsky" in self.sd_model.__class__.__name__:
                 model_type = 'kandinsky'
+            elif "HunyuanDiT" in self.sd_model.__class__.__name__:
+                model_type = 'hunyuandit'
             elif "Cascade" in self.sd_model.__class__.__name__:
                 model_type = 'sc'
             else:

modules/modelloader.py

@@ -1,8 +1,10 @@
+import io
 import os
 import time
 import json
 import shutil
 import importlib
+import contextlib
 from typing import Dict
 from urllib.parse import urlparse
 from PIL import Image
@@ -12,10 +14,25 @@ from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, Upscale
 from modules.paths import script_path, models_path
 
 
+loggedin = False
 diffuser_repos = []
 debug = shared.log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
+def hf_login(token=None):
+    global loggedin # pylint: disable=global-statement
+    import huggingface_hub as hf
+    token = token or shared.opts.huggingface_token
+    if token is not None and len(token) > 2 and not loggedin:
+        stdout = io.StringIO()
+        with contextlib.redirect_stdout(stdout):
+            hf.login(shared.opts.huggingface_token)
+        text = stdout.getvalue() or ''
+        line = [l for l in text.split('\n') if 'Token' in l]
+        shared.log.info(f'HF login: {line[0] if len(line) > 0 else text}')
+        loggedin = True
+
+
 def download_civit_meta(model_path: str, model_id):
     fn = os.path.splitext(model_path)[0] + '.json'
     url = f'https://civitai.com/api/v1/models/{model_id}'
@@ -188,7 +205,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
     token = token or shared.opts.huggingface_token
     if token is not None and len(token) > 2:
         shared.log.debug(f"Diffusers authentication: {token}")
-        hf.login(token)
+        hf_login(token)
     pipeline_dir = None
 
     ok = False

modules/sd_models.py

@@ -1381,7 +1381,8 @@ def set_diffusers_attention(pipe):
         modules = [getattr(pipe, n, None) for n in module_names]
         modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attn_processor")]
         for module in modules:
-            if 'SD3Transformer2DModel' in module.__class__.__name__: # TODO SD3
+            print('HERE', module.__class__.__name__)
+            if module.__class__.__name__ in ['SD3Transformer2DModel']:
                 module.set_attn_processor(p.JointAttnProcessor2_0())
             else:
                 module.set_attn_processor(attention)