Update app.py

app.py

@@ -17,6 +17,7 @@ import time
 log_queue = queue.Queue()
 current_logs = []
 
+
 def log(msg):
     """统一的日志处理函数"""
     print(msg)
@@ -74,14 +75,11 @@ def check_system_resources(model_name):
     else:
         raise MemoryError(f"❌ 系统内存不足 (需要 {required_memory_gb:.1f}GB, 可用 {available_memory_gb:.1f}GB)")
 
-def setup_environment(model_name, hf_token):
-    """设置环境并返回设备信息""" # try to get from env
-    if not hf_token:
-        raise ValueError("请在设置HF_TOKEN")
-    login(hf_token)
+def setup_environment(api, model_name):
     
-    # 检查系统资源并决定使用什么设备
-    device, available_memory = check_system_resources(model_name)
+    # # 检查系统资源并决定使用什么设备
+    # device, available_memory = check_system_resources(model_name)
+    device = "cpu"
     return device
 
 def create_hf_repo(repo_name, hf_token, private=True):
@@ -133,6 +131,7 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
         model.save_pretrained(output_dir)
         tokenizer.save_pretrained(output_dir)
         
+        
         return output_dir
         
     except Exception as e:
@@ -198,14 +197,15 @@ 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:
+        api = HfApi(token=hf_token)
         # 清空之前的日志
         current_logs.clear()
         
         # 设置环境和检查资源
-        device = setup_environment(base_model, hf_token)
+        device = setup_environment(api, base_model)
         
         # 创建HuggingFace仓库
-        repo_url = create_hf_repo(repo_name, hf_token)
+        repo_url = create_hf_repo(api, repo_name)
         
         # 设置输出目录
         output_dir = os.path.join(".", "output", repo_name)
@@ -214,16 +214,25 @@ def process_model(base_model, lora_model, repo_name, hf_token, progress=gr.Progr
         # 下载并合并模型
         model_path = download_and_merge_model(base_model, lora_model, output_dir, device)
         
-        progress(0.4, desc="开始8位量化...")
-        # 量化并上传模型
-        quantize_and_push_model(model_path, repo_name, bits=8)
+        # 推送到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.7, desc="开始4位量化...")
-        quantize_and_push_model(model_path, repo_name, bits=4)
+        # 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="处理完成")
+        # final_message = f"全部完成！模型已上传至: https://huggingface.co/{repo_name}"
+        # log(final_message)
+        # progress(1.0, desc="处理完成")
         
         return "\n".join(current_logs)
     except Exception as e: