Upload convert_url_to_diffusers_flux_gr.py

convert_url_to_diffusers_flux_gr.py

@@ -10,7 +10,14 @@ import os
 import argparse
 import gradio as gr
 # also requires aria, gdown, peft, huggingface_hub, safetensors, transformers, accelerate, pytorch_lightning
+
+import subprocess
+subprocess.run('pip cache purge', shell=True)
+
 import spaces
+@spaces.GPU()
+def spaces_dummy():
+    pass
 
 flux_dev_repo = "ChuckMcSneed/FLUX.1-dev"
 flux_schnell_repo = "black-forest-labs/FLUX.1-schnell"
@@ -38,6 +45,24 @@ def is_repo_name(s):
     import re
     return re.fullmatch(r'^[^/,\s]+?/[^/,\s]+?$', s)
 
+def clear_cache():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+def clear_sd(sd: dict):
+    for k in list(sd.keys()):
+        sd.pop(k)
+    del sd
+    torch.cuda.empty_cache()
+    gc.collect()
+
+def clone_sd(sd: dict):
+    print("Cloning state dict.")
+    for k in list(sd.keys()):
+        sd[k] = sd.pop(k).detach().clone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
 def print_resource_usage():
     import psutil
     cpu_usage = psutil.cpu_percent()
@@ -363,7 +388,7 @@ def read_safetensors_metadata(path):
 
 def normalize_key(k: str):
     return k.replace("vae.", "").replace("model.diffusion_model.", "")\
-        .replace("text_encoders.clip_l.transformer.text_model.", "")\
+        .replace("text_encoders.clip_l.transformer.", "")\
         .replace("text_encoders.t5xxl.transformer.", "")
 
 def load_json_list(path: str):
@@ -465,9 +490,7 @@ with torch.no_grad():
             print(e)
             return
         finally:
-            del state_dict
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(state_dict)
         new_path = str(Path(savepath, Path(path).stem + "_fixed" + Path(path).suffix))
         metadata = read_safetensors_metadata(path)
         progress(0.5, desc=f"Saving FLUX.1 safetensors: {new_path}")
@@ -476,9 +499,7 @@ with torch.no_grad():
         save_file(new_sd, new_path, metadata={"format": "pt", **metadata})
         progress(1, desc=f"Saved FLUX.1 safetensors: {new_path}")
         print(f"Saved FLUX.1 safetensors: {new_path}")
-        del new_sd
-        torch.cuda.empty_cache()
-        gc.collect()
+        clear_sd(new_sd)
 
 with torch.no_grad():
     def extract_norm_flux_module_sd(path: str, dtype: torch.dtype = torch.bfloat16,
@@ -506,9 +527,7 @@ with torch.no_grad():
         finally:
             progress(1, desc=f"Normalized FLUX.1 {name} safetensors: {path}")
             print(f"Normalized FLUX.1 {name} safetensors: {path}")
-            del state_dict
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(state_dict)
         return new_sd
 
 with torch.no_grad():
@@ -541,9 +560,7 @@ with torch.no_grad():
                 for k, v in sharded_sd.items():
                     sharded_sd[k] = v.to(device="cpu")
                 sd = sd | sharded_sd.copy()
-                del sharded_sd
-                torch.cuda.empty_cache()
-                gc.collect()
+                clear_sd(sharded_sd)
         except Exception as e:
                 print(e)
         return sd
@@ -561,9 +578,7 @@ with torch.no_grad():
             for k, v in sd.items():
                 if k in set(keys_flux_transformer): sd[k] = v.to(device="cpu")
             save_torch_state_dict(sd, path, filename_pattern=pattern, max_shard_size=size)
-            del sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(sd)
             progress(0.25, desc=f"Saved temporary files to disk: {path}")
             print(f"Saved temporary files to disk: {path}")
             for filepath in glob.glob(f"{path}/*.safetensors"):
@@ -574,9 +589,7 @@ with torch.no_grad():
                 for k, v in sharded_sd.items():
                     sharded_sd[k] = v.to(device="cpu")
                 save_file(sharded_sd, str(filepath))
-                del sharded_sd
-                torch.cuda.empty_cache()
-                gc.collect()
+                clear_sd(sharded_sd)
             print(f"Loading temporary files from disk: {path}")
             sd = load_sharded_safetensors(path)
             print(f"Loaded temporary files from disk: {path}")
@@ -599,9 +612,7 @@ with torch.no_grad():
             for k, v in sd.items():
                 sd[k] = v.to(device="cpu")
             save_torch_state_dict(sd, path, filename_pattern=pattern, max_shard_size=size)
-            del sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(sd)
             progress(0.25, desc=f"Saved temporary files to disk: {path}")
             print(f"Saved temporary files to disk: {path}")
             for filepath in glob.glob(f"{path}/*.safetensors"):
@@ -612,9 +623,7 @@ with torch.no_grad():
                 for k, v in sharded_sd.items():
                     sharded_sd[k] = v.to(device="cpu")
                 save_file(sharded_sd, str(filepath))
-                del sharded_sd
-                torch.cuda.empty_cache()
-                gc.collect()
+                clear_sd(sharded_sd)
                 print(f"Processed temporary files: {str(filepath)}")
             print(f"Loading temporary files from disk: {path}")
             sd = load_sharded_safetensors(path)
@@ -678,8 +687,7 @@ with torch.no_grad():
                           quantization: bool = False, model_type: str = "dev", dequant: bool = False):
         save_flux_other_diffusers(savepath, model_type)
         normalize_flux_state_dict(loadpath, savepath, dtype, dequant)
-        torch.cuda.empty_cache()
-        gc.collect()
+        clear_cache()
 
 with torch.no_grad(): # Much lower memory consumption, but higher disk load
     def flux_to_diffusers_lowmem(loadpath: str, savepath: str, dtype: torch.dtype = torch.bfloat16,
@@ -698,40 +706,46 @@ with torch.no_grad(): # Much lower memory consumption, but higher disk load
         vae_sd_path = savepath.removesuffix("/") + "/vae"
         vae_sd_pattern = "diffusion_pytorch_model{suffix}.safetensors"
         vae_sd_size = "10GB"
+        print_resource_usage() #
         metadata = {"format": "pt", **read_safetensors_metadata(loadpath)}
+        clear_cache()
+        print_resource_usage() #
         if "vae" not in use_original:
             vae_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "VAE",
                                                 keys_flux_vae)
             to_safetensors_flux_module(vae_sd, vae_sd_path, vae_sd_pattern, vae_sd_size,
                                        quantization, "VAE", None)
-            del vae_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(vae_sd)
+            print_resource_usage() #
         if "text_encoder" not in use_original:
             clip_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "Text Encoder",
                                                 keys_flux_clip)
             to_safetensors_flux_module(clip_sd, clip_sd_path, clip_sd_pattern, clip_sd_size,
                                        quantization, "Text Encoder", None)
-            del clip_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(clip_sd)
+            print_resource_usage() #
         if "text_encoder_2" not in use_original:
             te_sd = extract_norm_flux_module_sd(loadpath, dtype, dequant, "Text Encoder 2",
                                                 keys_flux_t5xxl)
             to_safetensors_flux_module(te_sd, te_sd_path, te_sd_pattern, te_sd_size,
                                        quantization, "Text Encoder 2", None)
-            del te_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(te_sd)
+            print_resource_usage() #
         unet_sd = extract_norm_flux_module_sd(loadpath, dtype, dequant, "Transformer",
                                               keys_flux_transformer)
-        if not local: os.remove(loadpath)
+        clear_cache()
+        print_resource_usage() #
+        if not local:
+            os.remove(loadpath)
+            print("Deleted downloaded file.")
+        clear_cache()
+        print_resource_usage() #
         to_safetensors_flux_module(unet_sd, unet_sd_path, unet_sd_pattern, unet_sd_size,
                                         quantization, "Transformer", metadata)
-        del unet_sd
-        torch.cuda.empty_cache()
-        gc.collect()
+        clear_sd(unet_sd)
+        print_resource_usage() #
         save_flux_other_diffusers(savepath, model_type, use_original)
+        print_resource_usage() #
 
 with torch.no_grad(): # lowest memory consumption, but higheest disk load
     def flux_to_diffusers_lowmem2(loadpath: str, savepath: str, dtype: torch.dtype = torch.bfloat16,
@@ -752,47 +766,52 @@ with torch.no_grad(): # lowest memory consumption, but higheest disk load
         vae_sd_path = savepath.removesuffix("/") + "/vae"
         vae_sd_pattern = "diffusion_pytorch_model{suffix}.safetensors"
         vae_sd_size = "10GB"
+        print_resource_usage() #
         metadata = {"format": "pt", **read_safetensors_metadata(loadpath)}
+        clear_cache()
+        print_resource_usage() #
         if "vae" not in use_original:
             vae_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "VAE",
                                                 keys_flux_vae)
             to_safetensors_flux_module(vae_sd, vae_sd_path, vae_sd_pattern, vae_sd_size,
                                        quantization, "VAE", None)
-            del vae_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(vae_sd)
+            print_resource_usage() #
         if "text_encoder" not in use_original:
             clip_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "Text Encoder",
                                                 keys_flux_clip)
             to_safetensors_flux_module(clip_sd, clip_sd_path, clip_sd_pattern, clip_sd_size,
                                        quantization, "Text Encoder", None)
-            del clip_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(clip_sd)
+            print_resource_usage() #
         if "text_encoder_2" not in use_original:
             te_sd = extract_norm_flux_module_sd(loadpath, dtype, dequant, "Text Encoder 2",
                                                 keys_flux_t5xxl)
             to_safetensors_flux_module(te_sd, te_sd_path, te_sd_pattern, te_sd_size,
                                        quantization, "Text Encoder 2", None)
-            del te_sd
-            torch.cuda.empty_cache()
-            gc.collect()
+            clear_sd(te_sd)
+            print_resource_usage() #
         unet_sd = extract_normalized_flux_state_dict_sharded(loadpath, dtype, dequant,
                                         unet_temp_path, unet_sd_pattern, unet_temp_size)
+        clear_cache()
+        print_resource_usage() #
         unet_sd = convert_flux_transformer_sd_to_diffusers_sharded(unet_sd, unet_temp_path,
                                                                    unet_sd_pattern, unet_temp_size)
+        clear_cache()
+        print_resource_usage() #
         to_safetensors_flux_module(unet_sd, unet_sd_path, unet_sd_pattern, unet_sd_size,
                                         quantization, "Transformer", metadata)
-        del unet_sd
-        torch.cuda.empty_cache()
-        gc.collect()
+        clear_sd(unet_sd)
+        print_resource_usage() #
         save_flux_other_diffusers(savepath, model_type, use_original)
+        print_resource_usage() #
 
 def convert_url_to_diffusers_flux(url, civitai_key="", is_upload_sf=False, data_type="bf16",
                                   model_type="dev", dequant=False, use_original=["vae", "text_encoder"],
                                   hf_user="", hf_repo="", q=None, progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Start converting...")
     temp_dir = "."
+    print_resource_usage() #
     new_file = get_download_file(temp_dir, url, civitai_key)
     if not new_file:
         print(f"Not found: {url}")
@@ -825,6 +844,7 @@ def convert_url_to_fixed_flux_safetensors(url, civitai_key="", is_upload_sf=Fals
                                   model_type="dev", dequant=False, q=None, progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Start converting...")
     temp_dir = "."
+    print_resource_usage() #
     new_file = get_download_file(temp_dir, url, civitai_key)
     if not new_file:
         print(f"Not found: {url}")