load with connector.

.gitignore

@@ -6,6 +6,7 @@ __pycache__
 /SwinIR/*
 /repositories
 /venv
+/venv310
 /tmp
 /model.ckpt
 /models/**/*
@@ -37,3 +38,5 @@ notification.mp3
 /node_modules
 /package-lock.json
 /.coverage*
+/output.log
+/webui-user.sh

modules/api/api.py

@@ -546,7 +546,9 @@ def unloadapi(self):
         return {}
 
     def reloadapi(self):
+        print("start reload api")
         reload_model_weights()
+        print("end reload api")
 
         return {}
 

modules/hashes.py

@@ -30,6 +30,14 @@ def calculate_remote_sha256(filename):
 
     return hash_sha256.hexdigest()
 
+# def calculate_remote_sha256(filename):
+#     blksize = 1024 * 1024
+
+#     buf = read_remote_model(filename, start = 0, size=blksize)
+#     hash_object = hashlib.sha256(buf)
+
+#     return hash_object.hexdigest()
+
 def sha256_from_cache(filename, title, use_addnet_hash=False, remote_model = False):
     hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes")
     ondisk_mtime = os.path.getmtime(filename) if not remote_model else get_remote_model_mmtime(filename)

modules/initialize.py

@@ -152,7 +152,7 @@ def load_model():
         from modules import devices
         devices.first_time_calculation()
 
-    Thread(target=load_model).start()
+    # Thread(target=load_model).start()
 
     from modules import shared_items
     shared_items.reload_hypernetworks()

modules/launch_utils.py

@@ -351,7 +351,11 @@ def prepare_environment():
     if args.reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):
         run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)
         startup_timer.record("install torch")
-
+
+    if not is_installed("cython"):
+        run_pip("install cython", "cython")
+        startup_timer.record("install cython")
+
     if not args.skip_torch_cuda_test and not check_run_python("import torch; assert torch.cuda.is_available()"):
         raise RuntimeError(
             'Torch is not able to use GPU; '

modules/sd_models.py

@@ -76,21 +76,26 @@ def read_metadata():
                 self.metadata = cache.cached_data_for_file('safetensors-metadata', "checkpoint/" + name + ('[remote]' if self.remote_model else '' ), filename, read_metadata, remote_model)
             except Exception as e:
                 errors.display(e, f"reading metadata for {filename}")
-
+        print("CheckpointInfo start")
         self.name = name
         self.name_for_extra = os.path.splitext(os.path.basename(filename))[0]
         self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]
         self.hash = model_hash(filename) if not remote_model else remote_model_hash(filename)
 
         self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{name + ('[remote]' if self.remote_model else '' )}", remote_model=remote_model)
         self.shorthash = self.sha256[0:10] if self.sha256 else None
-
+        print("sha256: %s" % self.sha256)
+        print("shorthash: %s" % self.shorthash)
+
         self.title = name + ('[remote]' if self.remote_model else '' )+ ('' if self.shorthash is None else f'[{self.shorthash}]')
         self.short_title = self.name_for_extra + ('[remote]' if self.remote_model else '') + ('' if self.shorthash is None else f'[{self.shorthash}]')
+        print("title: %s" % self.title)
+        print("short_title: %s" % self.short_title)
 
         self.ids = [self.hash, self.model_name, self.title, name, self.name_for_extra, f'{name} [{self.hash}]']
         if self.shorthash:
             self.ids += [self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]', f'{self.name_for_extra} [{self.shorthash}]']
+        print("CheckpointInfo end")
 
     def register(self):
         checkpoints_list[self.title] = self
@@ -170,6 +175,8 @@ def list_models():
         for filename in remote_models:
             checkpoint_info = CheckpointInfo(filename, remote_model=True)
             checkpoint_info.register()
+            print ("list_model: %s " % filename)
+
 re_strip_checksum = re.compile(r"\s*\[[^]]+]\s*$")
 
 

modules/sd_remote_models.py

@@ -4,8 +4,14 @@
 import threading
 from io import BytesIO
 from modules import shared
+from osstorchconnector import OssCheckpoint
+import torch
 
-
+def __check_bucket_opts():
+    if shared.opts.bucket_name and shared.opts.bucket_endpoint:
+        return True
+    print("Bucket opts not specified.")
+    return False
 
 def __bucket__():
     auth = oss2.Auth(os.environ.get('ACCESS_KEY_ID'), os.environ.get('ACCESS_KEY_SECRET'))
@@ -19,8 +25,10 @@ def get_remote_model_mmtime(model_name):
     return  __bucket__().head_object(model_name).last_modified
 
 def list_remote_models(ext_filter):
-    dir = shared.opts.bucket_model_ckpt_dir if shared.opts.bucket_model_ckpt_dir.endswith('/') else shared.opts.bucket_model_ckpt_dir + '/'
+    if not __check_bucket_opts():
+        return []
     output = []
+    dir = shared.opts.bucket_model_ckpt_dir if shared.opts.bucket_model_ckpt_dir.endswith('/') else shared.opts.bucket_model_ckpt_dir + '/'
     for obj in oss2.ObjectIteratorV2(__bucket__(), prefix = dir, delimiter = '/', start_after=dir, fetch_owner=False):
         if obj.is_prefix():
             print('directory: ', obj.key)
@@ -69,6 +77,21 @@ def read_remote_model(checkpoint_file, start=0, size=-1):
     buffer.seek(0)
     return buffer
 
+
+
+def load_remote_model_ckpt(checkpoint_file, map_location) -> bytes:
+    if not __check_bucket_opts():
+        return bytes()
+
+    checkpoint = OssCheckpoint(endpoint=shared.opts.bucket_endpoint)
+    CHECKPOINT_URI = "oss://%s/%s" % (shared.opts.bucket_name, checkpoint_file)
+    print("load %s state.." % CHECKPOINT_URI)
+    state_dict = None
+    with checkpoint.reader(CHECKPOINT_URI) as reader:
+        state_dict = torch.load(reader, map_location = map_location, weights_only = True)
+        print("type:", type(state_dict))
+    return state_dict
+
 def __range_get(object_name, buffer, offset, start, end, read_chunk_size):
     chunk_size = int(read_chunk_size)
     with __bucket__().get_object(object_name, byte_range=(start, end))as object_stream:

modules/shared_options.py

@@ -264,7 +264,7 @@
     "hires_fix_show_sampler": OptionInfo(False, "Hires fix: show hires checkpoint and sampler selection").needs_reload_ui(),
     "hires_fix_show_prompts": OptionInfo(False, "Hires fix: show hires prompt and negative prompt").needs_reload_ui(),
     "disable_token_counters": OptionInfo(False, "Disable prompt token counters").needs_reload_ui(),
-    "load_remote_ckpt": OptionInfo(False, "Load ckpt models from remote object storage").needs_reload_ui(),
+    "load_remote_ckpt": OptionInfo(True, "Load ckpt models from remote object storage").needs_reload_ui(),
     'bucket_name': OptionInfo("", "Bucket name to download ckpt model"),
     'bucket_endpoint': OptionInfo("", "Bucket endpoint to download ckpt model"),
     'bucket_model_ckpt_dir': OptionInfo("", "Ckpt model directory in bucket"),