Update A.py

A.py

@@ -1,8 +1,20 @@
+import os
+import requests
 from safetensors.torch import load_file, save_file
 import torch
 torch.cuda.empty_cache()
 import torch.nn.functional as F
-from tqdm import tqdm  # Ensure tqdm is installed
+from tqdm import tqdm
+
+def download_file(url, dest_path):
+    print(f"Downloading {url} to {dest_path}")
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        with open(dest_path, 'wb') as f:
+            for chunk in response.iter_content(1024):
+                f.write(chunk)
+    else:
+        raise Exception(f"Failed to download file from {url}")
 
 def load_model(file_path):
     return load_file(file_path)
@@ -38,19 +50,38 @@ def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.5):
 
     return merged
 
+def cleanup_files(*file_paths):
+    for file_path in file_paths:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            print(f"Deleted {file_path}")
+
 if __name__ == "__main__":
     try:
+        model1_url = "https://huggingface.co/multimodalart/FLUX.1-dev2pro-full/resolve/main/flux1-dev.safetensors"
+        model2_url = "https://huggingface.co/datasets/John6666/flux1-backup-202409/resolve/main/brainflux_v10.safetensors"
+        
         model1_path = "flux1-dev.safetensors"
         model2_path = "brainflux_v10.safetensors"
         blend_ratio = 0.4
         output_file = "output_checkpoint.safetensors"
 
+        # Downloading files
+        download_file(model1_url, model1_path)
+        download_file(model2_url, model2_path)
+
+        # Loading models
         model1 = load_model(model1_path)
         model2 = load_model(model2_path)
 
+        # Merging models
         merged_model = merge_checkpoints(model1, model2, blend_ratio)
 
+        # Saving merged model
         save_model(merged_model, output_file)
+
+        # Cleaning up downloaded files
+        cleanup_files(model1_path, model2_path)
         
     except Exception as e:
         print(f"An error occurred: {e}")