cleanup

modules/sync_models.py

@@ -6,7 +6,7 @@
 from modules.shared import syncLock
 
 FREESPACE = 20
-def check_space_s3_download(s3_client,bucket_name,s3_folder,local_folder,file,size,mode):
+def check_space_s3_download(s3_client, bucket_name, s3_folder, local_folder, file, size, mode):
     print(f"bucket_name:{bucket_name},s3_folder:{s3_folder},file:{file}")
     if file == '' or None:
         print('Debug log:file is empty, return')
@@ -38,7 +38,7 @@ def check_space_s3_download(s3_client,bucket_name,s3_folder,local_folder,file,si
     else:
         return False
 
-def free_local_disk(local_folder,size,mode):
+def free_local_disk(local_folder, size,mode):
     disk_usage = psutil.disk_usage('/tmp')
     freespace = disk_usage.free/(1024**3)
     if freespace - size >= FREESPACE:
@@ -75,7 +75,7 @@ def free_local_disk(local_folder,size,mode):
         print(f"Remove file: {oldest_file} now left space:{freespace}") 
         filename = os.path.basename(oldest_file)
 
-def list_s3_objects(s3_client,bucket_name, prefix=''):
+def list_s3_objects(s3_client, bucket_name, prefix=''):
     objects = []
     paginator = s3_client.get_paginator('list_objects_v2')
     page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix)
@@ -142,9 +142,7 @@ def initial_s3_download(s3_client, s3_folder, local_folder,cache_dir,mode):
     with open(s3_file_name, "w") as f:
         json.dump(s3_files, f)
 
-
-
-def sync_s3_folder(local_folder,cache_dir,mode):
+def sync_s3_folder(local_folder, cache_dir,mode):
     s3 = boto3.client('s3')
     def sync(mode):
         # print (f'sync:{mode}')
@@ -222,7 +220,6 @@ def sync(mode):
             elif mode == 'lora':
                 print('Nothing To do')
 
-
     # Create a thread function to keep syncing with the S3 folder
     def sync_thread(mode):  
         while True:

webui.py

@@ -308,10 +308,10 @@ def webui():
             shared.s3_folder_cn = "stable-diffusion-webui/models/ControlNet"
             shared.s3_folder_lora = "stable-diffusion-webui/models/Lora"
         #only download the cn models and the first sd model from default bucket, to accerlate the startup time
-        initial_s3_download(s3_client,shared.s3_folder_sd,sd_models_tmp_dir,cache_dir,'sd')
-        sync_s3_folder(sd_models_tmp_dir,cache_dir,'sd')
-        sync_s3_folder(cn_models_tmp_dir,cache_dir,'cn')
-        sync_s3_folder(lora_models_tmp_dir,cache_dir,'lora')
+        initial_s3_download(shared.s3_client, shared.s3_folder_sd, sd_models_tmp_dir,cache_dir,'sd')
+        sync_s3_folder(sd_models_tmp_dir, cache_dir, 'sd')
+        sync_s3_folder(cn_models_tmp_dir, cache_dir, 'cn')
+        sync_s3_folder(lora_models_tmp_dir, cache_dir, 'lora')
     initialize()
 
     while 1: