fix the swarm parser

scripts/iib/parsers/stable_swarm_ui.py

@@ -15,17 +15,40 @@ class StableSwarmUIParser:
     @classmethod
     def get_exif_data(clz, image: Image) -> str:
         items = image.info or {}
+        exif_bytes = b""
 
         if "exif" in items:
+            # Handle EXIF data (common in JPEGs, but can be in PNGs via eXIf chunk)
             exif = piexif.load(items["exif"])
+            # UserComment is in the Exif IFD, tag 37510
             exif_bytes = (
                 (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
             )
+        else:
+            # Fallback for PNG info, as metadata can be in text chunks.
+            # Common keys are 'parameters' (A1111-style) or 'Comment'.
+            raw_info = items.get("parameters") or items.get("Comment")
+            if isinstance(raw_info, str):
+                # Pillow reads PNG tEXt chunks as strings. To get the original bytes
+                # (which include the "UNICODE" marker), we need to re-encode it.
+                # latin-1 is a safe choice as it maps every byte value 0-255 to a character.
+                exif_bytes = raw_info.encode('latin-1')
+            elif isinstance(raw_info, bytes):
+                exif_bytes = raw_info
 
+        # Stable Swarm UI encodes metadata as UTF-16 after a "UNICODE" marker
         unicode_start = exif_bytes.find(b"UNICODE")
         if unicode_start == -1:
-            raise ValueError("'UNICODE' markup isn't found")
+            # If no UNICODE marker, try to decode as UTF-8 as a fallback
+            try:
+                decoded_str = exif_bytes.decode("utf-8", errors="ignore")
+                if '"sui_image_params"' in decoded_str:
+                    return decoded_str
+            except Exception:
+                pass # Ignore decoding errors and raise the final error
+            raise ValueError("'UNICODE' marker not found in Exif data.")
 
+        # Extract and decode the UTF-16 JSON string
         unicode_data = exif_bytes[unicode_start + len("UNICODE") + 1 :]
         geninfo = unicode_data.decode("utf-16")
         return geninfo
@@ -34,23 +57,93 @@ class StableSwarmUIParser:
     def parse(clz, img: Image, file_path):
         if not clz.test(img, file_path):
             raise Exception("The input image does not match the current parser.")
-        exif_data = json.loads(clz.get_exif_data(img))["sui_image_params"]
-        prompt = exif_data.pop("prompt")
-        negativeprompt = exif_data.pop("negativeprompt")
-        steps = exif_data.pop("steps")
-        meta_kv = [f"Steps: {steps}", "Source Identifier: StableSwarmUI"]
-        for key, value in exif_data.items():
-            value = replace_punctuation(str(value))
-            meta_kv.append(f"{key}: {value}")
-        meta = ", ".join(meta_kv)
-        info = "\n".join([prompt, f"Negative prompt: {negativeprompt}", meta])
+
+        full_exif_data = json.loads(clz.get_exif_data(img))
+        params_data = full_exif_data.get("sui_image_params", {})
+        models_data = full_exif_data.get("sui_models", [])
+
+        # --- Reconstruct Prompt with LoRAs ---
+        prompt = params_data.get("prompt", "")
+        loras = params_data.get("loras", [])
+        lora_weights = params_data.get("loraweights", [])
+
+        lora_hashes = []
+        lora_hash_map = {
+            # Normalize name by removing extension for matching
+            model_info.get("name", "").replace(".safetensors", "").replace(".pt", ""): model_info.get("hash", "").replace("0x", "")
+            for model_info in models_data if model_info.get("param") == "loras"
+        }
+
+        # Combine LoRAs and weights into the prompt string
+        for i, lora_name in enumerate(loras):
+            if i < len(lora_weights):
+                weight = lora_weights[i]
+                prompt += f" <lora:{lora_name}:{weight}>"
+
+                # Find and store the corresponding hash
+                lora_hash = lora_hash_map.get(lora_name, "N/A")
+                if lora_hash != "N/A":
+                    lora_hashes.append(f"{lora_name}: {lora_hash}")
+
+        negative_prompt = params_data.get("negativeprompt", "").strip()
+
+        # --- Build Metadata Line ---
+        meta_kv = []
+        if "steps" in params_data: meta_kv.append(f"Steps: {params_data['steps']}")
+        if "sampler" in params_data: meta_kv.append(f"Sampler: {params_data['sampler']}")
+        if "cfgscale" in params_data: meta_kv.append(f"CFG scale: {params_data['cfgscale']}")
+        if "seed" in params_data: meta_kv.append(f"Seed: {params_data['seed']}")
+        if "width" in params_data and "height" in params_data:
+            meta_kv.append(f"Size: {params_data['width']}x{params_data['height']}")
+
+        # Find model name and hash
+        model_name = params_data.get("model")
+        model_hash = ""
+        if model_name:
+            for model_info in models_data:
+                # Match by param type and ensure the name starts with the model from params
+                if model_info.get("param") == "model" and model_info.get("name", "").startswith(model_name):
+                    model_hash = model_info.get("hash", "").replace("0x", "")
+                    break
+
+        if model_hash: meta_kv.append(f"Model hash: {model_hash[:10]}")
+        if model_name: meta_kv.append(f"Model: {model_name}")
+
+        if lora_hashes:
+            meta_kv.append(f"Lora hashes: {', '.join(lora_hashes)}")
+
+        # Add remaining parameters
+        other_params = params_data.copy()
+        keys_to_ignore = [
+            "prompt", "negativeprompt", "steps", "sampler", "cfgscale", "seed",
+            "width", "height", "model", "loras", "loraweights", "lorasectionconfinement"
+        ]
+        for key in keys_to_ignore:
+            other_params.pop(key, None)
+
+        for key, value in other_params.items():
+            value_str = replace_punctuation(str(value))
+            meta_kv.append(f"{key}: {value_str}")
+
+        meta_kv.append("Source Identifier: StableSwarmUI")
+        meta_string = ", ".join(meta_kv)
+
+        # --- Final Assembly ---
+        info = f"{prompt}\nNegative prompt: {negative_prompt}\n{meta_string}"
+
         params = parse_generation_parameters(info)
+
+        final_meta = params.get("meta", {})
+        final_meta["final_width"] = img.size[0]
+        final_meta["final_height"] = img.size[1]
+        if model_hash:
+            final_meta["Model hash"] = model_hash  # Store full hash
+
         return ImageGenerationInfo(
             info,
             ImageGenerationParams(
-                meta=params["meta"]
-                | {"final_width": img.size[0], "final_height": img.size[1]},
-                pos_prompt=params["pos_prompt"],
+                meta=final_meta,
+                pos_prompt=params.get("pos_prompt", prompt),
                 extra=params,
             ),
         )
@@ -59,6 +152,7 @@ class StableSwarmUIParser:
     def test(clz, img: Image, file_path: str):
         try:
             exif = clz.get_exif_data(img)
-            return exif.find("sui_image_params") != -1
-        except Exception as e:
+            # Check for the unique key that identifies Stable Swarm UI metadata
+            return '"sui_image_params"' in exif
+        except Exception:
             return False