import os import random import shutil import torch import gradio as gr from PIL import Image, ImageChops from typing import List, Dict, Any from collections import defaultdict, deque import numpy as np from .base_pipeline import BasePipeline from core.settings import * from comfy_integration.nodes import * from utils.app_utils import get_value_at_index, sanitize_prompt, get_lora_path, get_embedding_path, ensure_controlnet_model_downloaded, ensure_ipadapter_models_downloaded, sanitize_filename from core.workflow_assembler import WorkflowAssembler class SdImagePipeline(BasePipeline): def get_required_models(self, model_display_name: str, **kwargs) -> List[str]: return [model_display_name] def _topological_sort(self, workflow: Dict[str, Any]) -> List[str]: graph = defaultdict(list) in_degree = {node_id: 0 for node_id in workflow} for node_id, node_info in workflow.items(): for input_value in node_info.get('inputs', {}).values(): if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str): source_node_id = input_value[0] if source_node_id in workflow: graph[source_node_id].append(node_id) in_degree[node_id] += 1 queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0]) sorted_nodes = [] while queue: current_node_id = queue.popleft() sorted_nodes.append(current_node_id) for neighbor_node_id in graph[current_node_id]: in_degree[neighbor_node_id] -= 1 if in_degree[neighbor_node_id] == 0: queue.append(neighbor_node_id) if len(sorted_nodes) != len(workflow): raise RuntimeError("Workflow contains a cycle and cannot be executed.") return sorted_nodes def _execute_workflow(self, workflow: Dict[str, Any], initial_objects: Dict[str, Any]): with torch.no_grad(): computed_outputs = initial_objects try: sorted_node_ids = self._topological_sort(workflow) print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}") except RuntimeError as e: print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---") for node_id, node_info in workflow.items(): print(f" Node {node_id} ({node_info['class_type']}):") for input_name, input_value in node_info['inputs'].items(): if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str): print(f" - {input_name} <- [{input_value[0]}, {input_value[1]}]") raise e for node_id in sorted_node_ids: if node_id in computed_outputs: continue node_info = workflow[node_id] class_type = node_info['class_type'] node_class = NODE_CLASS_MAPPINGS.get(class_type) if node_class is None: raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?") node_instance = node_class() kwargs = {} for param_name, param_value in node_info['inputs'].items(): if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str): source_node_id, output_index = param_value if source_node_id not in computed_outputs: raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.") source_output_tuple = computed_outputs[source_node_id] kwargs[param_name] = get_value_at_index(source_output_tuple, output_index) else: kwargs[param_name] = param_value function_name = getattr(node_class, 'FUNCTION') execution_method = getattr(node_instance, function_name) result = execution_method(**kwargs) computed_outputs[node_id] = result final_node_id = None for node_id in reversed(sorted_node_ids): if workflow[node_id]['class_type'] == 'SaveImage': final_node_id = node_id break if not final_node_id: raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.") save_image_inputs = workflow[final_node_id]['inputs'] image_source_node_id, image_source_index = save_image_inputs['images'] return get_value_at_index(computed_outputs[image_source_node_id], image_source_index) def _gpu_logic(self, ui_inputs: Dict, loras_string: str, required_models_for_gpu: List[str], workflow: Dict[str, Any], assembler: WorkflowAssembler, progress=gr.Progress(track_tqdm=True)): model_display_name = ui_inputs['model_display_name'] progress(0.1, desc="Moving models to GPU...") self.model_manager.move_models_to_gpu(required_models_for_gpu) progress(0.4, desc="Executing workflow...") loaded_model_tuple = self.model_manager.loaded_models[model_display_name] ckpt_loader_node_id = assembler.node_map.get("ckpt_loader") if not ckpt_loader_node_id: raise RuntimeError("Workflow is missing the 'ckpt_loader' node required for model injection.") initial_objects = { ckpt_loader_node_id: loaded_model_tuple } decoded_images_tensor = self._execute_workflow(workflow, initial_objects=initial_objects) output_images = [] start_seed = ui_inputs['seed'] if ui_inputs['seed'] != -1 else random.randint(0, 2**64 - 1) for i in range(decoded_images_tensor.shape[0]): img_tensor = decoded_images_tensor[i] pil_image = Image.fromarray((img_tensor.cpu().numpy() * 255.0).astype("uint8")) current_seed = start_seed + i width_for_meta = ui_inputs.get('width', 'N/A') height_for_meta = ui_inputs.get('height', 'N/A') params_string = f"{ui_inputs['positive_prompt']}\nNegative prompt: {ui_inputs['negative_prompt']}\n" params_string += f"Steps: {ui_inputs['num_inference_steps']}, Sampler: {ui_inputs['sampler']}, Scheduler: {ui_inputs['scheduler']}, CFG scale: {ui_inputs['guidance_scale']}, Seed: {current_seed}, Size: {width_for_meta}x{height_for_meta}, Base Model: {model_display_name}" if ui_inputs['task_type'] != 'txt2img': params_string += f", Denoise: {ui_inputs['denoise']}" if loras_string: params_string += f", {loras_string}" pil_image.info = {'parameters': params_string.strip()} output_images.append(pil_image) return output_images def run(self, ui_inputs: Dict, progress): progress(0, desc="Preparing models...") task_type = ui_inputs['task_type'] ui_inputs['positive_prompt'] = sanitize_prompt(ui_inputs.get('positive_prompt', '')) ui_inputs['negative_prompt'] = sanitize_prompt(ui_inputs.get('negative_prompt', '')) required_models = self.get_required_models(model_display_name=ui_inputs['model_display_name']) self.model_manager.ensure_models_downloaded(required_models, progress=progress) lora_data = ui_inputs.get('lora_data', []) active_loras_for_gpu, active_loras_for_meta = [], [] sources, ids, scales, files = lora_data[0::4], lora_data[1::4], lora_data[2::4], lora_data[3::4] for i, (source, lora_id, scale, _) in enumerate(zip(sources, ids, scales, files)): if scale > 0 and lora_id and lora_id.strip(): lora_filename = None if source == "File": lora_filename = sanitize_filename(lora_id) elif source == "Civitai": local_path, status = get_lora_path(source, lora_id, ui_inputs['civitai_api_key'], progress) if local_path: lora_filename = os.path.basename(local_path) else: raise gr.Error(f"Failed to prepare LoRA {lora_id}: {status}") if lora_filename: active_loras_for_gpu.append({"lora_name": lora_filename, "strength_model": scale, "strength_clip": scale}) active_loras_for_meta.append(f"{source} {lora_id}:{scale}") progress(0.1, desc="Loading models into RAM...") self.model_manager.load_managed_models(required_models, active_loras=active_loras_for_gpu, progress=progress) ui_inputs['denoise'] = 1.0 if task_type == 'img2img': ui_inputs['denoise'] = ui_inputs.get('img2img_denoise', 0.7) elif task_type == 'hires_fix': ui_inputs['denoise'] = ui_inputs.get('hires_denoise', 0.55) temp_files_to_clean = [] if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR) if task_type == 'img2img': input_image_pil = ui_inputs.get('img2img_image') if input_image_pil: temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png") input_image_pil.save(temp_file_path, "PNG") ui_inputs['input_image'] = os.path.basename(temp_file_path) temp_files_to_clean.append(temp_file_path) ui_inputs['width'] = input_image_pil.width ui_inputs['height'] = input_image_pil.height elif task_type == 'inpaint': inpaint_dict = ui_inputs.get('inpaint_image_dict') if not inpaint_dict or not inpaint_dict.get('background') or not inpaint_dict.get('layers'): raise gr.Error("Inpainting requires an input image and a drawn mask.") background_img = inpaint_dict['background'].convert("RGBA") composite_mask_pil = Image.new('L', background_img.size, 0) for layer in inpaint_dict['layers']: if layer: layer_alpha = layer.split()[-1] composite_mask_pil = ImageChops.lighter(composite_mask_pil, layer_alpha) inverted_mask_alpha = Image.fromarray(255 - np.array(composite_mask_pil), mode='L') r, g, b, _ = background_img.split() composite_image_with_mask = Image.merge('RGBA', [r, g, b, inverted_mask_alpha]) temp_file_path = os.path.join(INPUT_DIR, f"temp_inpaint_composite_{random.randint(1000, 9999)}.png") composite_image_with_mask.save(temp_file_path, "PNG") ui_inputs['inpaint_image'] = os.path.basename(temp_file_path) temp_files_to_clean.append(temp_file_path) ui_inputs.pop('inpaint_mask', None) elif task_type == 'outpaint': input_image_pil = ui_inputs.get('outpaint_image') if input_image_pil: temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png") input_image_pil.save(temp_file_path, "PNG") ui_inputs['input_image'] = os.path.basename(temp_file_path) temp_files_to_clean.append(temp_file_path) elif task_type == 'hires_fix': input_image_pil = ui_inputs.get('hires_image') if input_image_pil: temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png") input_image_pil.save(temp_file_path, "PNG") ui_inputs['input_image'] = os.path.basename(temp_file_path) temp_files_to_clean.append(temp_file_path) embedding_data = ui_inputs.get('embedding_data', []) embedding_filenames = [] if embedding_data: emb_sources, emb_ids, emb_files = embedding_data[0::3], embedding_data[1::3], embedding_data[2::3] for i, (source, emb_id, _) in enumerate(zip(emb_sources, emb_ids, emb_files)): if emb_id and emb_id.strip(): emb_filename = None if source == "File": emb_filename = sanitize_filename(emb_id) elif source == "Civitai": local_path, status = get_embedding_path(source, emb_id, ui_inputs['civitai_api_key'], progress) if local_path: emb_filename = os.path.basename(local_path) else: raise gr.Error(f"Failed to prepare Embedding {emb_id}: {status}") if emb_filename: embedding_filenames.append(emb_filename) if embedding_filenames: embedding_prompt_text = " ".join([f"embedding:{f}" for f in embedding_filenames]) if ui_inputs['positive_prompt']: ui_inputs['positive_prompt'] = f"{ui_inputs['positive_prompt']}, {embedding_prompt_text}" else: ui_inputs['positive_prompt'] = embedding_prompt_text controlnet_data = ui_inputs.get('controlnet_data', []) active_controlnets = [] (cn_images, _, _, cn_strengths, cn_filepaths) = [controlnet_data[i::5] for i in range(5)] for i in range(len(cn_images)): if cn_images[i] and cn_strengths[i] > 0 and cn_filepaths[i] and cn_filepaths[i] != "None": ensure_controlnet_model_downloaded(cn_filepaths[i], progress) if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR) cn_temp_path = os.path.join(INPUT_DIR, f"temp_cn_{i}_{random.randint(1000, 9999)}.png") cn_images[i].save(cn_temp_path, "PNG") temp_files_to_clean.append(cn_temp_path) active_controlnets.append({ "image": os.path.basename(cn_temp_path), "strength": cn_strengths[i], "start_percent": 0.0, "end_percent": 1.0, "control_net_name": cn_filepaths[i] }) ipadapter_data = ui_inputs.get('ipadapter_data', []) active_ipadapters = [] if ipadapter_data: num_ipa_units = (len(ipadapter_data) - 5) // 3 final_preset, final_weight, final_lora_strength, final_embeds_scaling, final_combine_method = ipadapter_data[-5:] ipa_images, ipa_weights, ipa_lora_strengths = [ipadapter_data[i*num_ipa_units:(i+1)*num_ipa_units] for i in range(3)] all_presets_to_download = set() for i in range(num_ipa_units): if ipa_images[i] and ipa_weights[i] > 0 and final_preset: all_presets_to_download.add(final_preset) if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR) ipa_temp_path = os.path.join(INPUT_DIR, f"temp_ipa_{i}_{random.randint(1000, 9999)}.png") ipa_images[i].save(ipa_temp_path, "PNG") temp_files_to_clean.append(ipa_temp_path) active_ipadapters.append({ "image": os.path.basename(ipa_temp_path), "preset": final_preset, "weight": ipa_weights[i], "lora_strength": ipa_lora_strengths[i] }) if active_ipadapters and final_preset: all_presets_to_download.add(final_preset) for preset in all_presets_to_download: ensure_ipadapter_models_downloaded(preset, progress) if active_ipadapters: active_ipadapters.append({ 'is_final_settings': True, 'model_type': 'sd15', 'final_preset': final_preset, 'final_weight': final_weight, 'final_lora_strength': final_lora_strength, 'final_embeds_scaling': final_embeds_scaling, 'final_combine_method': final_combine_method }) from utils.app_utils import get_vae_path vae_source = ui_inputs.get('vae_source') vae_id = ui_inputs.get('vae_id') vae_file = ui_inputs.get('vae_file') vae_name_override = None if vae_source and vae_source != "None": if vae_source == "File": vae_name_override = sanitize_filename(vae_id) elif vae_source == "Civitai" and vae_id and vae_id.strip(): local_path, status = get_vae_path(vae_source, vae_id, ui_inputs.get('civitai_api_key'), progress) if local_path: vae_name_override = os.path.basename(local_path) else: raise gr.Error(f"Failed to prepare VAE {vae_id}: {status}") if vae_name_override: ui_inputs['vae_name'] = vae_name_override conditioning_data = ui_inputs.get('conditioning_data', []) active_conditioning = [] if conditioning_data: num_units = len(conditioning_data) // 6 prompts = conditioning_data[0*num_units : 1*num_units] widths = conditioning_data[1*num_units : 2*num_units] heights = conditioning_data[2*num_units : 3*num_units] xs = conditioning_data[3*num_units : 4*num_units] ys = conditioning_data[4*num_units : 5*num_units] strengths = conditioning_data[5*num_units : 6*num_units] for i in range(num_units): if prompts[i] and prompts[i].strip(): active_conditioning.append({ "prompt": prompts[i], "width": int(widths[i]), "height": int(heights[i]), "x": int(xs[i]), "y": int(ys[i]), "strength": float(strengths[i]) }) loras_string = f"LoRAs: [{', '.join(active_loras_for_meta)}]" if active_loras_for_meta else "" progress(0.8, desc="Assembling workflow...") if ui_inputs.get('seed') == -1: ui_inputs['seed'] = random.randint(0, 2**32 - 1) dynamic_values = {'task_type': ui_inputs['task_type'], 'model_type': "sd15"} recipe_path = os.path.join(os.path.dirname(__file__), "workflow_recipes", "sd_unified_recipe.yaml") assembler = WorkflowAssembler(recipe_path, dynamic_values=dynamic_values) workflow_inputs = { "positive_prompt": ui_inputs['positive_prompt'], "negative_prompt": ui_inputs['negative_prompt'], "seed": ui_inputs['seed'], "steps": ui_inputs['num_inference_steps'], "cfg": ui_inputs['guidance_scale'], "sampler_name": ui_inputs['sampler'], "scheduler": ui_inputs['scheduler'], "batch_size": ui_inputs['batch_size'], "clip_skip": -int(ui_inputs['clip_skip']), "denoise": ui_inputs['denoise'], "input_image": ui_inputs.get('input_image'), "inpaint_image": ui_inputs.get('inpaint_image'), "inpaint_mask": ui_inputs.get('inpaint_mask'), "left": ui_inputs.get('outpaint_left'), "top": ui_inputs.get('outpaint_top'), "right": ui_inputs.get('outpaint_right'), "bottom": ui_inputs.get('outpaint_bottom'), "hires_upscaler": ui_inputs.get('hires_upscaler'), "hires_scale_by": ui_inputs.get('hires_scale_by'), "model_name": ALL_MODEL_MAP[ui_inputs['model_display_name']][1], "vae_name": ui_inputs.get('vae_name'), "controlnet_chain": active_controlnets, "ipadapter_chain": active_ipadapters, "conditioning_chain": active_conditioning, } if task_type == 'txt2img': workflow_inputs['width'] = ui_inputs['width'] workflow_inputs['height'] = ui_inputs['height'] workflow = assembler.assemble(workflow_inputs) if workflow_inputs.get("vae_name"): print("--- [Workflow Patch] VAE override provided. Adding VAELoader and rewiring connections. ---") vae_loader_id = assembler._get_unique_id() vae_loader_node = assembler._get_node_template("VAELoader") vae_loader_node['inputs']['vae_name'] = workflow_inputs["vae_name"] workflow[vae_loader_id] = vae_loader_node vae_decode_id = assembler.node_map.get("vae_decode") if vae_decode_id and vae_decode_id in workflow: workflow[vae_decode_id]['inputs']['vae'] = [vae_loader_id, 0] print(f" - Rewired 'vae_decode' (ID: {vae_decode_id}) to use new VAELoader.") vae_encode_id = assembler.node_map.get("vae_encode") if vae_encode_id and vae_encode_id in workflow: workflow[vae_encode_id]['inputs']['vae'] = [vae_loader_id, 0] print(f" - Rewired 'vae_encode' (ID: {vae_encode_id}) to use new VAELoader.") else: print("--- [Workflow Info] No VAE override. Using VAE from checkpoint. ---") progress(1.0, desc="All models ready. Requesting GPU for generation...") try: results = self._execute_gpu_logic( self._gpu_logic, duration=ui_inputs['zero_gpu_duration'], default_duration=60, task_name=f"ImageGen ({task_type})", ui_inputs=ui_inputs, loras_string=loras_string, required_models_for_gpu=required_models, workflow=workflow, assembler=assembler, progress=progress ) finally: for temp_file in temp_files_to_clean: if temp_file and os.path.exists(temp_file): os.remove(temp_file) print(f"✅ Cleaned up temp file: {temp_file}") return results