Update app.py

app.py

@@ -75,14 +75,13 @@ def check_system_resources(model_name):
     else:
         raise MemoryError(f"❌ 系统内存不足 (需要 {required_memory_gb:.1f}GB, 可用 {available_memory_gb:.1f}GB)")
 
-def setup_environment(api, model_name):
-    
+def setup_environment(model_name):
     # # 检查系统资源并决定使用什么设备
     # device, available_memory = check_system_resources(model_name)
     device = "cpu"
     return device
 
-def create_hf_repo(repo_name, hf_token, private=True):
+def create_hf_repo(repo_name, private=True):
     """创建HuggingFace仓库"""
     try:
         # check if repo already exists
@@ -90,7 +89,7 @@ def create_hf_repo(repo_name, hf_token, private=True):
         if api.repo_exists(repo_name):
             log(f"仓库已存在: {repo_name}")
             return ValueError(f"仓库已存在: {repo_name}, 请使用其他名称或删除已存在的仓库")
-        repo_url = create_repo(repo_name, private=private, token=hf_token)
+        repo_url = create_repo(repo_name, private=private)
         log(f"创建仓库成功: {repo_url}")
         return repo_url
     except Exception as e:
@@ -197,15 +196,18 @@ def quantize_and_push_model(model_path, repo_id, bits=8):
 def process_model(base_model, lora_model, repo_name, hf_token, progress=gr.Progress()):
     """处理模型的主函数，用于Gradio界面"""
     try:
+        login(hf_token)
+        os.environ["HF_TOKEN"] = hf_token
         api = HfApi(token=hf_token)
+        
         # 清空之前的日志
         current_logs.clear()
         
         # 设置环境和检查资源
-        device = setup_environment(api, base_model)
+        device = setup_environment(base_model)
         
         # 创建HuggingFace仓库
-        repo_url = create_hf_repo(api, repo_name)
+        repo_url = create_hf_repo(repo_name)
         
         # 设置输出目录
         output_dir = os.path.join(".", "output", repo_name)
@@ -216,23 +218,31 @@ def process_model(base_model, lora_model, repo_name, hf_token, progress=gr.Progr
         
         # 推送到HuggingFace
         log(f"正在将模型推送到HuggingFace...")
-        api = HfApi()
+        
         api.upload_folder(
             folder_path=model_path,
             repo_id=repo_name,
             repo_type="model"
         )
         
-        # progress(0.4, desc="开始8位量化...")
-        # # 量化并上传模型
-        # quantize_and_push_model(model_path, repo_name, bits=8)
+        progress(0.4, desc="开始8位量化...")
+        # 量化并上传模型
+        quantize_and_push_model(model_path, repo_name, bits=8)
+        
+        progress(0.7, desc="开始4位量化...")
+        quantize_and_push_model(model_path, repo_name, bits=4)
+        
+        final_message = f"全部完成！模型已上传至: https://huggingface.co/{repo_name}"
+        log(final_message)
+        progress(1.0, desc="处理完成")
         
-        # progress(0.7, desc="开始4位量化...")
-        # quantize_and_push_model(model_path, repo_name, bits=4)
+        # remove hf_token from env
+        os.environ.pop("HF_TOKEN")
+        log("HF_TOKEN已从环境变量中删除")
         
-        # final_message = f"全部完成！模型已上传至: https://huggingface.co/{repo_name}"
-        # log(final_message)
-        # progress(1.0, desc="处理完成")
+        # remove model_path
+        os.remove(model_path)
+        log(f"模型路径已删除: {model_path}")
         
         return "\n".join(current_logs)
     except Exception as e: