Create worker.py

backend/worker.py

@@ -0,0 +1,50 @@
+import os
+import json
+import tempfile
+import shutil
+import traceback
+from .merge_engines.linear import linear_merge
+from .merge_engines.evolutionary import evolutionary_merge
+from .uploader import upload_to_hf
+from .memory_manager import estimate_memory, check_memory_safe
+
+async def run_job_async(job_id: str, params: dict):
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(None, run_job, job_id, params)
+
+def run_job(job_id, params):
+    try:
+        # memory pre-check
+        model_a_id = f"{params['model_a_source']}:{params['model_a_id']}"
+        model_b_id = f"{params['model_b_source']}:{params['model_b_id']}"
+        if not check_memory_safe(model_a_id, model_b_id, params.get("evo_params")):
+            raise RuntimeError("Insufficient memory for this merge. Job rejected.")
+        
+        # download models if needed
+        from .model_loader import prepare_model
+        path_a = prepare_model(params["model_a_source"], params["model_a_id"], params.get("civitai_key"))
+        path_b = prepare_model(params["model_b_source"], params["model_b_id"], params.get("civitai_key"))
+        
+        # output dir
+        output_dir = f"/data/output/{job_id}"
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # merge
+        if params["method"] == "linear":
+            alpha = params.get("linear_alpha", 0.5)
+            linear_merge(path_a, path_b, alpha, output_dir)
+        elif params["method"] == "evolutionary":
+            evo_params = params.get("evo_params", {})
+            dataset_path = params.get("dataset")
+            evolutionary_merge(path_a, path_b, dataset_path, output_dir, evo_params)
+        else:
+            raise ValueError(f"Unknown method: {params['method']}")
+        
+        # upload
+        repo_name = params.get("output_repo_name") or f"merge-{job_id[:8]}"
+        upload_to_hf(output_dir, repo_name, params["hf_token"])
+        
+        return {"repo_name": repo_name, "status": "success"}
+    except Exception as e:
+        traceback.print_exc()
+        raise