Upload convert_url_to_diffusers_flux_gr.py

convert_url_to_diffusers_flux_gr.py

@@ -607,7 +607,8 @@ with torch.no_grad():
             for filepath in glob.glob(f"{path}/*.safetensors"):
                 progress(0.25, desc=f"Processing temporary files: {str(filepath)}")
                 print(f"Processing temporary files: {str(filepath)}")
-                sharded_sd = extract_normalized_flux_state_dict_unet(str(filepath), dtype, dequant)
+                sharded_sd = extract_norm_flux_module_sd(str(filepath), dtype, dequant,
+                                                         "Transformer", keys_flux_transformer)
                 for k, v in sharded_sd.items():
                     sharded_sd[k] = v.to(device="cpu")
                 save_file(sharded_sd, str(filepath))
@@ -622,55 +623,42 @@ with torch.no_grad():
             print(e)
         return sd
 
-def download_repo(repo_name, path, download_sf=False, progress=gr.Progress(track_tqdm=True)):
+def download_repo(repo_name, path, use_original=["vae", "text_encoder"], progress=gr.Progress(track_tqdm=True)):
     from huggingface_hub import snapshot_download
     print(f"Downloading {repo_name}.")
     try:
-        if download_sf:
+        if "text_encoder_2" not in use_original:
             snapshot_download(repo_id=repo_name, local_dir=path, ignore_patterns=["transformer/", "*.sft", ".*", "README*", "*.md", "*.index", "*.jpg", "*.png", "*.webp"])
         else:
             snapshot_download(repo_id=repo_name, local_dir=path, ignore_patterns=["transformer/", "text_encoder_2/", "*.sft", ".*", "README*", "*.md", "*.index", "*.jpg", "*.png", "*.webp"])
     except Exception as e:
         print(e)
 
-def copy_nontensor_files(from_path, to_path, copy_sf=False):
+def copy_nontensor_files(from_path, to_path, use_original=["vae", "text_encoder"]):
     import shutil
-    if copy_sf:
-        print(f"Copying non-tensor files {from_path} to {to_path}")
-        shutil.copytree(from_path, to_path, ignore=shutil.ignore_patterns("*.safetensors", "*.bin", "*.sft", ".*", "README*", "*.md", "*.index", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
+    if "text_encoder_2" in use_original:
         te_from = str(Path(from_path, "text_encoder_2"))
         te_to = str(Path(to_path, "text_encoder_2"))
         print(f"Copying Text Encoder 2 files {te_from} to {te_to}")
         shutil.copytree(te_from, te_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
+    if "text_encoder" in use_original:
         te1_from = str(Path(from_path, "text_encoder"))
         te1_to = str(Path(to_path, "text_encoder"))
         print(f"Copying Text Encoder 1 files {te1_from} to {te1_to}")
         shutil.copytree(te1_from, te1_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-        tn2_from = str(Path(from_path, "tokenizer_2"))
-        tn2_to = str(Path(to_path, "tokenizer_2"))
-        print(f"Copying Tokenizer 2 files {tn2_from} to {tn2_to}")
-        shutil.copytree(tn2_from, tn2_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
+    if "vae" in use_original:
         vae_from = str(Path(from_path, "vae"))
         vae_to = str(Path(to_path, "vae"))
         print(f"Copying VAE files {vae_from} to {vae_to}")
         shutil.copytree(vae_from, vae_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-    else:
-        print(f"Copying non-tensor files {from_path} to {to_path}")
-        shutil.copytree(from_path, to_path, ignore=shutil.ignore_patterns("*.safetensors", "*.bin", "*.sft", ".*", "README*", "*.md", "*.index", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-        te1_from = str(Path(from_path, "text_encoder"))
-        te1_to = str(Path(to_path, "text_encoder"))
-        print(f"Copying Text Encoder 1 files {te1_from} to {te1_to}")
-        shutil.copytree(te1_from, te1_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-        tn2_from = str(Path(from_path, "tokenizer_2"))
-        tn2_to = str(Path(to_path, "tokenizer_2"))
-        print(f"Copying Tokenizer 2 files {tn2_from} to {tn2_to}")
-        shutil.copytree(tn2_from, tn2_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-        vae_from = str(Path(from_path, "vae"))
-        vae_to = str(Path(to_path, "vae"))
-        print(f"Copying VAE files {vae_from} to {vae_to}")
-        shutil.copytree(vae_from, vae_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
-  
-def save_flux_other_diffusers(path: str, model_type: str = "dev", copy_te: bool = False, progress=gr.Progress(track_tqdm=True)):
+    tn2_from = str(Path(from_path, "tokenizer_2"))
+    tn2_to = str(Path(to_path, "tokenizer_2"))
+    print(f"Copying Tokenizer 2 files {tn2_from} to {tn2_to}")
+    shutil.copytree(tn2_from, tn2_to, ignore=shutil.ignore_patterns(".*", "README*", "*.md", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
+    print(f"Copying non-tensor files {from_path} to {to_path}")
+    shutil.copytree(from_path, to_path, ignore=shutil.ignore_patterns("*.safetensors", "*.bin", "*.sft", ".*", "README*", "*.md", "*.index", "*.jpg", "*.png", "*.webp"), dirs_exist_ok=True)
+    
+def save_flux_other_diffusers(path: str, model_type: str = "dev", use_original: list = ["vae", "text_encoder"], progress=gr.Progress(track_tqdm=True)):
     import shutil
     progress(0, desc="Loading FLUX.1 Components.")
     print("Loading FLUX.1 Components.")
@@ -679,10 +667,10 @@ def save_flux_other_diffusers(path: str, model_type: str = "dev", copy_te: bool
     else: repo = flux_dev_repo
     os.makedirs(temppath, exist_ok=True)
     os.makedirs(path, exist_ok=True)
-    download_repo(repo, temppath, copy_te)
+    download_repo(repo, temppath, use_original)
     progress(0.5, desc="Saving FLUX.1 Components.")
     print("Saving FLUX.1 Components.")
-    copy_nontensor_files(temppath, path, copy_te)
+    copy_nontensor_files(temppath, path, use_original)
     shutil.rmtree(temppath)
 
 with torch.no_grad():
@@ -711,7 +699,6 @@ with torch.no_grad(): # Much lower memory consumption, but higher disk load
         vae_sd_pattern = "diffusion_pytorch_model{suffix}.safetensors"
         vae_sd_size = "10GB"
         metadata = {"format": "pt", **read_safetensors_metadata(loadpath)}
-        save_flux_other_diffusers(savepath, model_type, "text_encoder_2" in use_original)
         if "vae" not in use_original:
             vae_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "VAE",
                                                 keys_flux_vae)
@@ -744,6 +731,7 @@ with torch.no_grad(): # Much lower memory consumption, but higher disk load
         del unet_sd
         torch.cuda.empty_cache()
         gc.collect()
+        save_flux_other_diffusers(savepath, model_type, use_original)
 
 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,
@@ -765,7 +753,6 @@ with torch.no_grad(): # lowest memory consumption, but higheest disk load
         vae_sd_pattern = "diffusion_pytorch_model{suffix}.safetensors"
         vae_sd_size = "10GB"
         metadata = {"format": "pt", **read_safetensors_metadata(loadpath)}
-        save_flux_other_diffusers(savepath, model_type, "text_encoder_2" in use_original)
         if "vae" not in use_original:
             vae_sd = extract_norm_flux_module_sd(loadpath, torch.bfloat16, dequant, "VAE",
                                                 keys_flux_vae)
@@ -799,6 +786,7 @@ with torch.no_grad(): # lowest memory consumption, but higheest disk load
         del unet_sd
         torch.cuda.empty_cache()
         gc.collect()
+        save_flux_other_diffusers(savepath, model_type, use_original)
 
 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"],