File size: 14,200 Bytes
c21ae5c b4cd781 c21ae5c e0b7ec8 088ae2c b4cd781 26d09bf 23fbbb0 81d2307 c21ae5c b4cd781 c21ae5c 59036a6 c21ae5c 59036a6 088ae2c 59036a6 26d09bf 59036a6 81d2307 c21ae5c d0dbbdd | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 | import os
import random
import shutil
import torch
import gradio as gr
from PIL import Image
from typing import List, Dict, Any
from .base_pipeline import BasePipeline
from core.settings import *
from utils.app_utils import sanitize_prompt
from core.workflow_assembler import WorkflowAssembler
from .workflow_executor import WorkflowExecutor
from .pipeline_input_processor import process_pipeline_inputs
class SdImagePipeline(BasePipeline):
def get_required_models(self, model_display_name: str, **kwargs) -> List[str]:
model_info = ALL_MODEL_MAP.get(model_display_name)
if not model_info:
return [model_display_name]
path_or_components = model_info[1]
if isinstance(path_or_components, dict):
return [v for v in path_or_components.values() if v and v != "pixel_space"]
else:
return [model_display_name]
def _gpu_logic(self, ui_inputs: Dict, loras_string: str, workflow: Dict[str, Any], assembler: WorkflowAssembler, progress=gr.Progress(track_tqdm=True)):
model_display_name = ui_inputs['model_display_name']
progress(0.4, desc="Executing workflow...")
initial_objects = {}
decoded_images_tensor = WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)
output_images = []
raw_seed = ui_inputs.get('seed')
start_seed = int(raw_seed) if (raw_seed is not None and raw_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 ui_inputs.get('clip_skip') and ui_inputs['clip_skip'] != 1: params_string += f", Clip skip: {abs(ui_inputs['clip_skip'])}"
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']
model_display_name = ui_inputs['model_display_name']
model_type = MODEL_TYPE_MAP.get(model_display_name, 'sdxl')
architectures_dict = ARCHITECTURES_CONFIG.get('architectures', {})
workflow_model_type = architectures_dict.get(model_type, {}).get("model_type", model_type.lower().replace(" ", "").replace(".", ""))
ui_inputs['positive_prompt'] = sanitize_prompt(ui_inputs.get('positive_prompt', ''))
ui_inputs['negative_prompt'] = sanitize_prompt(ui_inputs.get('negative_prompt', ''))
if 'clip_skip' in ui_inputs and ui_inputs['clip_skip'] is not None:
ui_inputs['clip_skip'] = -int(ui_inputs['clip_skip'])
else:
ui_inputs['clip_skip'] = -1
required_models = self.get_required_models(model_display_name=model_display_name)
is_pid_enabled = (ui_inputs.get('pid_settings', 'OFF') == 'ON' and task_type == 'txt2img')
if is_pid_enabled:
import yaml
pid_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'yaml', 'pid.yaml')
pid_unet_name = "pid_flux1_1024_to_4096_4step_mxfp8.safetensors"
try:
with open(pid_config_path, 'r', encoding='utf-8') as f:
pid_config = yaml.safe_load(f) or {}
pid_items = pid_config.get("PiD", [])
for item in pid_items:
archs = item.get("architectures", [])
if workflow_model_type in archs:
pid_unet_name = item.get("filepath")
break
except Exception as e:
print(f"Error loading PiD config for download: {e}")
if pid_unet_name not in required_models:
required_models.append(pid_unet_name)
if "gemma_2_2b_it_elm_fp8_scaled.safetensors" not in required_models:
required_models.append("gemma_2_2b_it_elm_fp8_scaled.safetensors")
self.model_manager.ensure_models_downloaded(required_models, progress=progress)
temp_files_to_clean = []
try:
processed = process_pipeline_inputs(ui_inputs, progress, workflow_model_type)
temp_files_to_clean.extend(processed["temp_files_to_clean"])
active_loras_for_gpu = processed["active_loras_for_gpu"]
active_loras_for_meta = processed["active_loras_for_meta"]
active_controlnets = processed["active_controlnets"]
active_anima_controlnets = processed["active_anima_controlnets"]
active_diffsynth_controlnets = processed["active_diffsynth_controlnets"]
active_krea2_controlnets = processed.get("active_krea2_controlnets", [])
active_ipadapters = processed["active_ipadapters"]
active_flux1_ipadapters = processed["active_flux1_ipadapters"]
active_sd3_ipadapters = processed["active_sd3_ipadapters"]
active_styles = processed["active_styles"]
active_reference_latents = processed["active_reference_latents"]
active_hidream_o1_reference = processed["active_hidream_o1_reference"]
active_joyai_reference = processed.get("active_joyai_reference", [])
active_krea2_identity_edit = processed.get("active_krea2_identity_edit", [])
active_krea2_reference_edit = processed.get("active_krea2_reference_edit", [])
active_qwen_image_edit = processed.get("active_qwen_image_edit", [])
active_boogu_edit = processed.get("active_boogu_edit", [])
active_reference_images = processed.get("active_reference_images", [])
active_conditioning = processed["active_conditioning"]
loras_string = f"LoRAs: [{', '.join(active_loras_for_meta)}]" if active_loras_for_meta else ""
progress(0.8, desc="Assembling workflow...")
seed_val = ui_inputs.get('seed')
if seed_val is None or seed_val == -1:
ui_inputs['seed'] = random.randint(0, 2**32 - 1)
model_info = ALL_MODEL_MAP[model_display_name]
path_or_components = model_info[1]
latent_type = model_info[3] if len(model_info) > 3 and model_info[3] else 'latent'
latent_generator_template = "EmptyLatentImage"
if latent_type == 'sd3_latent':
latent_generator_template = "EmptySD3LatentImage"
elif latent_type == 'chroma_radiance_latent':
latent_generator_template = "EmptyChromaRadianceLatentImage"
elif latent_type == 'hunyuan_latent':
latent_generator_template = "EmptyHunyuanImageLatent"
dynamic_values = {
'task_type': ui_inputs['task_type'],
'model_type': workflow_model_type,
'latent_type': latent_type,
'latent_generator_template': latent_generator_template
}
recipe_path = os.path.join(os.path.dirname(__file__), "workflow_recipes", "unified_recipe.yaml")
assembler = WorkflowAssembler(recipe_path, dynamic_values=dynamic_values)
hidream_o1_smoothing_data = []
if workflow_model_type == 'hidream-o1' and model_display_name == "HiDream-O1-Image":
hidream_o1_smoothing_data.append({})
workflow_inputs = {
**ui_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": ui_inputs['clip_skip'],
"denoise": ui_inputs['denoise'],
"vae_name": ui_inputs.get('vae_name'),
"guidance": ui_inputs.get('guidance', 3.5),
"lora_chain": active_loras_for_gpu,
"controlnet_chain": active_controlnets if not active_anima_controlnets else [],
"anima_controlnet_lllite_chain": active_anima_controlnets,
"diffsynth_controlnet_chain": active_diffsynth_controlnets,
"krea2_controlnet_chain": active_krea2_controlnets,
"ipadapter_chain": active_ipadapters,
"flux1_ipadapter_chain": active_flux1_ipadapters,
"sd3_ipadapter_chain": active_sd3_ipadapters,
"style_chain": active_styles,
"conditioning_chain": active_conditioning,
"reference_latent_chain": active_reference_latents,
"hidream_o1_reference_chain": active_hidream_o1_reference,
"joyai_image_chain": active_joyai_reference,
"krea2_identity_edit_chain": active_krea2_identity_edit,
"krea2_style_reference_chain": active_krea2_reference_edit,
"qwen_image_edit_chain": active_qwen_image_edit,
"boogu_image_edit_chain": active_boogu_edit,
"reference_image_chain": active_reference_images,
"vae_chain": [ui_inputs.get('vae_name')] if ui_inputs.get('vae_name') else [],
"hidream_o1_smoothing_chain": hidream_o1_smoothing_data,
"pid_chain": [ui_inputs.get('pid_settings', 'OFF')] if is_pid_enabled else [],
"scheduler_width": ui_inputs.get('width', 1024),
"scheduler_height": ui_inputs.get('height', 1024),
}
if isinstance(path_or_components, dict):
workflow_inputs.update({
'unet_name': path_or_components.get('unet'),
'unet_uncond_name': path_or_components.get('unet_uncond'),
'vae_name': ui_inputs.get('vae_name') or path_or_components.get('vae'),
'clip_name': path_or_components.get('clip'),
'clip1_name': path_or_components.get('clip1'),
'clip2_name': path_or_components.get('clip2'),
'clip3_name': path_or_components.get('clip3'),
'clip4_name': path_or_components.get('clip4'),
'lora_name': path_or_components.get('lora'),
})
else:
workflow_inputs['model_name'] = path_or_components
if task_type == 'txt2img':
workflow_inputs['width'] = ui_inputs['width']
workflow_inputs['height'] = ui_inputs['height']
workflow = assembler.assemble(workflow_inputs)
progress(1.0, desc="All models ready. Requesting GPU for generation...")
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,
workflow=workflow,
assembler=assembler,
progress=progress
)
import json
import glob
from PIL import PngImagePlugin
prompt_json = json.dumps(workflow)
out_dir = os.path.abspath(OUTPUT_DIR)
os.makedirs(out_dir, exist_ok=True)
try:
existing_files = glob.glob(os.path.join(out_dir, "gen_*.png"))
existing_files.sort(key=os.path.getmtime)
while len(existing_files) > 50:
os.remove(existing_files.pop(0))
except Exception as e:
print(f"Warning: Failed to cleanup output dir: {e}")
final_results = []
for img in results:
if not isinstance(img, Image.Image):
final_results.append(img)
continue
metadata = PngImagePlugin.PngInfo()
params_string = img.info.get("parameters", "")
if params_string:
metadata.add_text("parameters", params_string)
metadata.add_text("prompt", prompt_json)
filename = f"gen_{random.randint(1000000, 9999999)}.png"
filepath = os.path.join(out_dir, filename)
img.save(filepath, "PNG", pnginfo=metadata)
final_results.append(filepath)
results = final_results
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 |