File size: 2,223 Bytes
6d2b369 cd8de0d 657f21b 99ed897 6d2b369 cd8de0d 6d2b369 38dacda 6d2b369 99ed897 6d2b369 99ed897 6d2b369 99ed897 6d2b369 99ed897 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import os
import json
import tempfile
import shutil
import traceback
from merge_engines.linear import merge_models
from merge_engines.evolutionary import evolutionary_merge
from uploader import upload_to_hf
from memory_manager import check_memory_safe
import asyncio
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:
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.")
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 = f"/app/data/output/{job_id}"
os.makedirs(output_dir, exist_ok=True)
method = params.get("method", "linear")
if 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:
merge_type = params.get("merge_type", "linear")
alpha = params.get("linear_alpha", 0.5)
layers_from_a = params.get("franken_layers", None)
merge_models(path_a, path_b, output_dir, method=merge_type,
alpha=alpha, layers_from_a=layers_from_a)
repo_name = params.get("output_repo_name") or f"merge-{job_id[:8]}"
upload_to_hf(output_dir, repo_name, params["hf_token"])
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
return {"repo_name": repo_name, "status": "success"}
except Exception as e:
traceback.print_exc()
raise
finally:
params["hf_token"] = None
params["civitai_key"] = None |