backend/worker.py

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