Fix PiD Injector.

chain_injectors/pid_injector.py

@@ -121,8 +121,13 @@ def inject(assembler, chain_definition, chain_items):
     if active_model_file:
         try:
             model_config = load_model_config()
-            checkpoints = model_config.get("Checkpoints", {})
+            checkpoints = {}
+            for k, v in model_config.items():
+                if isinstance(v, dict):
+                    checkpoints.update(v)
             for arch_name, arch_data in checkpoints.items():
+                if not isinstance(arch_data, dict):
+                    continue
                 models_list = arch_data.get("models", [])
                 for model_entry in models_list:
                     if model_entry.get('path') == active_model_file:
@@ -138,7 +143,16 @@ def inject(assembler, chain_definition, chain_items):
             print(f"Error looking up model architecture in PiD injector: {e}")
 
         if architecture:
-            architecture = architecture.lower().replace(" ", "-").replace(".", "")
+            try:
+                from core.settings import ARCHITECTURES_CONFIG
+                arch_info = ARCHITECTURES_CONFIG.get('architectures', {}).get(architecture, {})
+                mapped_arch = arch_info.get("model_type")
+                if mapped_arch:
+                    architecture = mapped_arch
+                else:
+                    architecture = architecture.lower().replace(" ", "-").replace(".", "")
+            except Exception:
+                architecture = architecture.lower().replace(" ", "-").replace(".", "")
         else:
             file_lower = active_model_file.lower().replace("-", "").replace("_", "").replace(".", "")
             for arch in sorted(architectures_settings.keys(), key=len, reverse=True):
@@ -149,7 +163,13 @@ def inject(assembler, chain_definition, chain_items):
                     candidates.append(arch.replace("-i1", ""))
                 if "-kv" in arch:
                     candidates.append(arch.replace("-kv", ""))
-                
+                if arch == "sdxl":
+                    candidates.append("xl")
+                if arch == "sd35":
+                    candidates.append("sd3")
+                if arch == "flux1":
+                    candidates.append("flux")
+
                 matched = False
                 for cand in candidates:
                     if cand.replace("-", "").replace(".", "") in file_lower:

core/pipelines/sd_image_pipeline.py

@@ -80,7 +80,7 @@ class SdImagePipeline(BasePipeline):
         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"
+            pid_unet_name = "pid_1.5_flux1_1024_to_4096_4step_int8_convrot.safetensors"
             try:
                 with open(pid_config_path, 'r', encoding='utf-8') as f:
                     pid_config = yaml.safe_load(f) or {}