Update app.py

app.py

@@ -8,6 +8,8 @@ from huggingface_hub import login, create_repo, HfApi
 import gradio as gr
 import queue
 import time
+import shutil
+
 
 # 全局日志
 log_queue = queue.Queue()
@@ -46,9 +48,7 @@ def check_system_resources(model_name):
     log("Checking system resources...")
     system_memory = psutil.virtual_memory()
     total_memory_gb = system_memory.total / (1024 ** 3)
-    available_memory_gb = system_memory.available / (1024 ** 3)
     log(f"Total system memory: {total_memory_gb:.1f}GB")
-    log(f"Available memory: {available_memory_gb:.1f}GB")
     
     model_size_gb = get_model_size_in_gb(model_name)
     required_memory_gb = model_size_gb * 2.5  # 预留额外内存
@@ -66,9 +66,9 @@ def check_system_resources(model_name):
     else:
         log("❌ No GPU detected.")
         
-    if available_memory_gb >= required_memory_gb:
+    if total_memory_gb >= required_memory_gb:
         log("✅ Sufficient CPU memory available; using CPU.")
-        return "cpu", available_memory_gb
+        return "cpu", total_memory_gb
     else:
         raise MemoryError(f"❌ Insufficient system memory (requires {required_memory_gb:.1f}GB, available {available_memory_gb:.1f}GB).")
 
@@ -117,10 +117,9 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
     model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True)
     log("Loading adapter tokenizer...")
     adapter_tokenizer = AutoTokenizer.from_pretrained(lora_model_name)
-    if adapter_tokenizer.pad_token != model.config.pad_token:
-        log("Resizing token embeddings...")
-        added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
-        model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
+    log("Resizing token embeddings...")
+    added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
+    model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
     log("Loading LoRA adapter...")
     peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True)
     log("Merging and unloading model...")
@@ -157,7 +156,9 @@ def clone_llamacpp_and_download_build():
     log("llama.cpp build completed.")
     # 返回到原始目录
     os.chdir(os.path.dirname(llamacpp_dir))
-    
+
+def remove_illegal_chars_in_path(text):
+    return text.replace(".", "_").replace(":", "_").replace("/", "_")
 
 @timeit
 def quantize(model_path, repo_id, quant_method=None):
@@ -180,12 +181,13 @@ def quantize(model_path, repo_id, quant_method=None):
     os.makedirs(model_output_dir, exist_ok=True)
 
     # 中间文件保存在 model_output 目录下
-    temp_gguf_path = os.path.join(model_output_dir, f"{repo_id}-f16.gguf")
+    guff_16 = os.path.join(model_output_dir, f"{repo_id}-f16.gguf")
     
-    if not os.path.exists(temp_gguf_path):
+    if not os.path.exists(guff_16):
         log(f"正在将模型转换为GGML格式")
         convert_script = os.path.join(llamacpp_dir, "convert_hf_to_gguf.py")
-        convert_cmd = f"python {convert_script} {model_path} --outfile {temp_gguf_path}"
+        convert_cmd = f"python {convert_script} {model_path} --outfile {guff_16}"
+        print(f"syscall:[{convert_cmd}]")
         os.system(convert_cmd)
     else:
         log(f"GGML中间文件已存在，跳过转换")
@@ -194,7 +196,8 @@ def quantize(model_path, repo_id, quant_method=None):
     final_path = os.path.join(model_output_dir, f"{repo_id}-{quant_method}.gguf")
     log(f"正在进行{quant_method}量化")
     quantize_bin = os.path.join(llamacpp_dir, "build", "bin", "llama-quantize") 
-    quant_cmd = f"{quantize_bin} {temp_gguf_path} {final_path} {quant_method}"
+    quant_cmd = f"{quantize_bin} {guff_16} {final_path} {quant_method}"
+    print(f"syscall:[{quant_cmd}]")
     
     if not os.path.exists(final_path):
         os.system(quant_cmd)
@@ -226,7 +229,8 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
             adapter_config = PeftConfig.from_pretrained(lora_model_name)
             base_model_name = adapter_config.base_model_name_or_path
         if repo_name.strip().lower() == "auto":
-            repo_name = f"{username}/{base_model_name.split('/')[-1]}_{lora_model_name.split('/')[-1]}"
+            repo_name = f"{base_model_name.split('/')[-1]}_{lora_model_name.split('/')[-1]}"
+            repo_name = remove_illegal_chars_in_path(repo_name)
         
         device = setup_environment(base_model_name)
         repo_name = create_hf_repo(repo_name)
@@ -238,15 +242,20 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
         
         # 量化模型
         for quant_method in quant_methods:
-            quantize(f"{output_dir}/model.gguf", repo_name, bits=8, quant_method=quant_method)
+            quantize(output_dir, repo_name, quant_method=quant_method)
         
         # 上传合并后的模型和量化模型
         api.upload_large_folder(
             folder_path=model_path,
             repo_id=repo_name,
             repo_type="model",
-            num_workers=4,
+            num_workers=os.cpu_count() if os.cpu_count() > 4 else 4,
+            print_report_every=10,
         )
+        
+        # rm -rf model_path
+        shutil.rmtree(model_path)
+        
         return "\n".join(current_logs)
     except Exception as e:
         error_message = f"Error during processing: {e}"