diff --git a/.gitignore b/.gitignore index 2bf37729..71ce1f52 100644 --- a/.gitignore +++ b/.gitignore @@ -210,4 +210,4 @@ model_display_names.json # DreamLayer logs and test logs logs -test-logs \ No newline at end of file +test-logsdream_layer_backend/matrix_runner_state.json diff --git a/dream_layer_backend/dream_layer.py b/dream_layer_backend/dream_layer.py index 70d11481..94da0c10 100644 --- a/dream_layer_backend/dream_layer.py +++ b/dream_layer_backend/dream_layer.py @@ -10,7 +10,9 @@ import json import subprocess from dream_layer_backend_utils.random_prompt_generator import fetch_positive_prompt, fetch_negative_prompt +from dream_layer_backend_utils.task_runner import MatrixRunner from dream_layer_backend_utils.fetch_advanced_models import get_lora_models, get_settings, is_valid_directory, get_upscaler_models, get_controlnet_models + # Add ComfyUI directory to Python path current_dir = os.path.dirname(os.path.abspath(__file__)) parent_dir = os.path.dirname(current_dir) @@ -120,6 +122,10 @@ def import_comfyui_main(): } }) +# Persist state next to backend code so it survives restarts/page refresh +MATRIX_STATE_FILE = os.path.join(os.path.dirname(__file__), "matrix_runner_state.json") +runner = MatrixRunner(state_file=MATRIX_STATE_FILE) + COMFY_API_URL = "http://127.0.0.1:8188" def get_available_models(): @@ -432,6 +438,55 @@ def send_to_img2img(): except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 +@app.route('/api/matrix-runner/start', methods=['POST']) +def matrix_runner_start(): + try: + data = request.json or {} + # accept lists/ranges for seeds, steps, samplers, etc. + # keep only list-valued keys for the sweep + param_dict = {k: v for k, v in data.items() if isinstance(v, list)} + runner.generate(param_dict) # deterministic expansion + return jsonify({"status":"success","total_jobs": len(runner.jobs)}) + except Exception as e: + return jsonify({"status":"error","message": str(e)}), 500 + +@app.route('/api/matrix-runner/pause', methods=['POST']) +def matrix_runner_pause(): + runner.pause() + return jsonify({"status":"paused"}) + +@app.route('/api/matrix-runner/resume', methods=['POST']) +def matrix_runner_resume(): + runner.resume() + return jsonify({"status":"resumed"}) + +@app.route('/api/matrix-runner/status', methods=['GET']) +def matrix_runner_status(): + jobs = runner.jobs + return jsonify({ + "status":"ok", + "total_jobs": len(jobs), + "pending": sum(j["status"]=="pending" for j in jobs), + "running": sum(j["status"]=="running" for j in jobs), + "done": sum(j["status"]=="done" for j in jobs), + "paused": runner.paused + }) + +@app.route('/api/matrix-runner/next', methods=['POST']) +def matrix_runner_next(): + """Return the next job (and mark it running) so the executor can process it.""" + job = runner.next_job() + if not job: + return jsonify({"status":"empty"}) + return jsonify({"status":"ok","job": job}) + +@app.route('/api/matrix-runner/complete', methods=['POST']) +def matrix_runner_complete(): + data = request.json or {} + job_id = data.get("job_id") + runner.complete_job(job_id) + return jsonify({"status":"ok"}) + @app.route('/api/send-to-extras', methods=['POST', 'OPTIONS']) def send_to_extras(): """Send image to extras tab""" diff --git a/dream_layer_backend/dream_layer_backend_utils/task_runner.py b/dream_layer_backend/dream_layer_backend_utils/task_runner.py new file mode 100644 index 00000000..efcadd58 --- /dev/null +++ b/dream_layer_backend/dream_layer_backend_utils/task_runner.py @@ -0,0 +1,89 @@ +""" +# dream_layer_backend/task_runner.py +# Author: Natalia Rodriguez Figueroa +# Email: natalia_rodriguezuc@berkeley.edu + +# Task 2: Implement a matrix/grid image generation workflow. +""" + +import os +import json +import random +import logging +import itertools +from flask import Flask, request, jsonify +from flask_cors import CORS +from PIL import Image +import numpy as np + +logger = logging.getLogger(__name__) + +import itertools +import json +import os + +class MatrixRunner: + def __init__(self, state_file="matrix_jobs.json"): + self.state_file = state_file + self.jobs = [] + self.index = 0 + self.paused = False + self.load() + + # --- Core Job Management --- + def generate(self, param_dict): + """Create all jobs from parameter ranges and reset state""" + keys = list(param_dict.keys()) + combos = list(itertools.product(*param_dict.values())) + self.jobs = [ + {"id": i, **dict(zip(keys, combo)), "status": "pending"} + for i, combo in enumerate(combos) + ] + self.index = 0 + self.paused = False + self.save() + + def next_job(self): + """Get the next pending job and mark it as running""" + if self.paused: + return None + while self.index < len(self.jobs): + job = self.jobs[self.index] + self.index += 1 + if job["status"] == "pending": + job["status"] = "running" + self.save() + return job + return None # No more jobs + + def complete_job(self, job_id): + """Mark job as done""" + for job in self.jobs: + if job["id"] == job_id: + job["status"] = "done" + break + self.save() + + # --- Pause / Resume --- + def pause(self): + self.paused = True + self.save() + + def resume(self): + self.paused = False + self.save() + + # --- Persistence --- + def save(self): + with open(self.state_file, "w") as f: + json.dump( + {"jobs": self.jobs, "index": self.index, "paused": self.paused}, f, indent=2 + ) + + def load(self): + if os.path.exists(self.state_file): + with open(self.state_file, "r") as f: + state = json.load(f) + self.jobs = state.get("jobs", []) + self.index = state.get("index", 0) + self.paused = state.get("paused", False) diff --git a/dream_layer_backend/matrix_runner_state.json b/dream_layer_backend/matrix_runner_state.json new file mode 100644 index 00000000..f3525001 --- /dev/null +++ b/dream_layer_backend/matrix_runner_state.json @@ -0,0 +1,48 @@ +{ + "jobs": [ + { + "id": 0, + "seeds": 1, + "steps": 10, + "Samplers": "euler", + "status": "done" + }, + { + "id": 1, + "seeds": 1, + "steps": 20, + "Samplers": "euler", + "status": "running" + }, + { + "id": 2, + "seeds": 2, + "steps": 10, + "Samplers": "euler", + "status": "pending" + }, + { + "id": 3, + "seeds": 2, + "steps": 20, + "Samplers": "euler", + "status": "pending" + }, + { + "id": 4, + "seeds": 3, + "steps": 10, + "Samplers": "euler", + "status": "pending" + }, + { + "id": 5, + "seeds": 3, + "steps": 20, + "Samplers": "euler", + "status": "pending" + } + ], + "index": 2, + "paused": false +} \ No newline at end of file