From 1f95aa89a4a000462b2934c43df54ae4eaa7a1b5 Mon Sep 17 00:00:00 2001 From: IDGallagher Date: Mon, 14 Jul 2025 12:28:37 -0300 Subject: [PATCH] Updated explorer to allow infinite explorations --- nodes/explorer.py | 80 ++++++++++++++++++++++++++++++++++------------- 1 file changed, 59 insertions(+), 21 deletions(-) diff --git a/nodes/explorer.py b/nodes/explorer.py index 9789df8..01822da 100644 --- a/nodes/explorer.py +++ b/nodes/explorer.py @@ -20,17 +20,20 @@ class IG_ExplorerNode: return { "required": { "folder": ("STRING", {"forceInput": True}), - "precision": ("FLOAT", {"default": 0.0001, "min": FLOAT_STEP, "max": 1, "step": FLOAT_STEP}) + "precision": ("FLOAT", {"default": 0.0001, "min": FLOAT_STEP, "max": 1, "step": FLOAT_STEP}), + "max_images": ("INT", {"default": 0, "min": 0, "step": 1}), + "num_explorations": ("INT", {"default": 0, "min": 0, "step": 1}), + "start_seed": ("INT", {"default": 0, "min": 0, "step": 1}), }, } @classmethod - def IS_CHANGED(cls, folder: str, precision, **kwargs): + def IS_CHANGED(cls, folder: str, precision, max_images, num_explorations, start_seed, **kwargs): print(f"IS CHANGED") # return random.randint(1, 100000000000) return float("NaN") - RETURN_TYPES = ("FLOAT","STRING",) + RETURN_TYPES = ("FLOAT","STRING","INT",) FUNCTION = "main" CATEGORY = TREE_EXPLORER @@ -55,10 +58,21 @@ class IG_ExplorerNode: print(f"Full file: {full_file}") return os.path.isfile(full_file) + def count_images(self): + return len([f for f in os.listdir(self.folder) if f.startswith("image_") and f.endswith("_00001_.png") and os.path.isfile(os.path.join(self.folder, f))]) + + def load_state(self, path, start_seed): + try: + with open(path, 'r') as f: + return json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + return {'current_seed': start_seed, 'explorations_done': 0, 'active_subfolder': None} + + def save_state(self, path, state): + with open(path, 'w') as f: + json.dump(state, f) + def save_queue(self, queue, filename): - # Ensure directory exists - os.makedirs(os.path.dirname(filename), exist_ok=True) - # Use the .txt extension or any other preferred plaintext format filename = os.path.splitext(filename)[0] + '.txt' @@ -106,24 +120,29 @@ class IG_ExplorerNode: print(f"{param1} {param2} - {ssim_value}") return ssim_value - def process_jobs(self, queue, precision): + def process_jobs(self, queue, precision, max_images, current_count): for job in queue: if job.difference is None: - # Generate images if not already done if not self.image_exists(job.param1): - return job.param1 + if max_images == 0 or current_count < max_images: + return job.param1 + else: + continue if not self.image_exists(job.param2): - return job.param2 - # Measure and set the difference + if max_images == 0 or current_count < max_images: + return job.param2 + else: + continue job.difference = self.measure_difference(job.param1, job.param2) - # Filter jobs that have a difference greater than precision and have a measured difference eligible_jobs = [job for job in queue if job.difference is not None and abs(job.param1 - job.param2) > precision] - # If no eligible jobs are found, return -1 if not eligible_jobs: return -1 + if max_images != 0 and current_count >= max_images: + return -1 + # Find the job with the maximum difference max_diff_job = min(eligible_jobs, key=lambda job: job.difference) print(f"Min similarity {max_diff_job.param1} {max_diff_job.param2} {max_diff_job.difference}") @@ -138,11 +157,30 @@ class IG_ExplorerNode: # Generate the new image return mid_param - def main(self, folder, precision): - self.folder = folder - queue_path = os.path.join(folder,'job_queue.json') - queue = self.load_queue(queue_path) - param = self.process_jobs(queue, precision) - self.save_queue(queue, queue_path) - assert param > -1, "🤖 Exploration finished!" - return (param, self.image_filename(param)) + def main(self, folder, precision, max_images, num_explorations, start_seed): + os.makedirs(folder, exist_ok=True) + state_path = os.path.join(folder, 'exploration_state.json') + while True: + state = self.load_state(state_path, start_seed) + if state['active_subfolder'] is None: + if num_explorations > 0 and state['explorations_done'] >= num_explorations: + raise AssertionError("🤖 All explorations finished!") + current_seed = state['current_seed'] + subfolder = os.path.join(folder, f"seed_{current_seed}") + os.makedirs(subfolder, exist_ok=True) + state['active_subfolder'] = subfolder + state['current_seed'] = current_seed + 1 + state['explorations_done'] += 1 + self.save_state(state_path, state) + self.folder = state['active_subfolder'] + queue_path = os.path.join(self.folder, 'job_queue.json') + queue = self.load_queue(queue_path) + current_count = self.count_images() + param = self.process_jobs(queue, precision, max_images, current_count) + self.save_queue(queue, queue_path) + if param > -1: + seed = int(os.path.basename(self.folder).split('_')[1]) + return (param, self.image_filename(param), seed) + else: + state['active_subfolder'] = None + self.save_state(state_path, state)