🤖 Code Audit: 8 potential issue(s) found

[asmit25805]

Code Audit Report

All findings are reviewed for confidence before posting.
Please verify each finding before acting on it.

Repository: ideogram-oss/ideogram4
Findings: 8 issue(s) found — 🟠 7 high · 🟡 1 medium

---

1. 🐛 Potential Index Out of Range Error

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/autoencoder.py
Location	Encoder.forward
Confidence	95%

Problem:
In the Encoder.forward method, the code accesses the 'down' attribute of 'self.down[i_level]' without checking if 'i_level' is within the valid range of indices for 'self.down'. This could lead to an IndexError if 'i_level' exceeds the length of 'self.down'.

Suggested Fix:
Add a check to ensure 'i_level' is within the valid range of indices for 'self.down' before accessing 'self.down[i_level]'.

---

2. 🐛 Missing error handling for file read

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/caption_verifier.py
Location	verify_file
Confidence	96%

Problem:
verify_file reads the file with Path(path).read_text(...) without catching I/O errors. If the provided path does not exist or is unreadable, an unhandled OSError propagates and crashes the program, breaking normal usage.

Suggested Fix:
Wrap the file read in a try/except block, catch OSError (e.g., FileNotFoundError, PermissionError), and append a warning to the result instead of raising. Example:

def verify_file(self, path: PathLike) -> list[str]:
    try:
        raw = Path(path).read_text(encoding="utf-8")
    except OSError as e:
        return [f"file read error: {e}"]
    return self.verify_raw(raw)

---

3. 🐛 Assert used for runtime validation may be stripped in optimized mode

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/latent_norm.py
Location	def get_latent_norm() -> tuple[torch.Tensor, torch.Tensor]:
Confidence	95%

Problem:
The function uses an assert statement to verify that the generated tensors have shape (128,). When Python is executed with the -O (optimize) flag, assert statements are removed, so the shape check is bypassed. If the LATENT_SHIFT or LATENT_SCALE tuples have an unexpected length, the function will return tensors of incorrect shape, leading to downstream runtime errors that are harder to diagnose.

Suggested Fix:
Replace the assert with an explicit conditional check and raise a descriptive exception, e.g.,

if shift.shape != (128,) or scale.shape != (128,):
    raise ValueError(f"Expected LATENT_SHIFT and LATENT_SCALE to have 128 elements, got {len(LATENT_SHIFT)} and {len(LATENT_SCALE)}")

---

4. 🐛 Potential KeyError in openrouter_chat function

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/magic_prompt.py
Location	openrouter_chat function
Confidence	95%

Problem:
The openrouter_chat function does not check if the 'choices' key exists in the response data before trying to access it. If the 'choices' key does not exist, a KeyError will be raised.

Suggested Fix:
Add a check to ensure the 'choices' key exists in the response data before trying to access it.

---

5. 🐛 Potential NaN values in velocity prediction

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/modeling_ideogram4.py
Location	Ideogram4Transformer.forward
Confidence	95%

Problem:
In the Ideogram4Transformer.forward method, the t_cond variable is calculated using the _sinusoidal_embedding function. If the input t is very large, the freq variable in _sinusoidal_embedding may become very small, potentially causing NaN values in the velocity prediction. This could lead to incorrect results or crashes.

Suggested Fix:
Add input validation to ensure t is within a reasonable range before calculating t_cond.

---

6. 🐛 Potential KeyError in _load_sharded_state_dict

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/pipeline_ideogram4.py
Location	_load_sharded_state_dict
Confidence	95%

Problem:
The function _load_sharded_state_dict assumes that the 'weight_map' key exists in the index file. However, if this key is missing, a KeyError will be raised. It would be better to add a check to ensure the key exists before trying to access it.

Suggested Fix:
Add a check to ensure 'weight_map' exists in the index file before accessing it, e.g., 'if 'weight_map' in index: weight_map = index['weight_map']'

---

7. 🐛 Invalid argument 'assign' passed to torch.nn.Module.load_state_dict

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/quantized_loading.py
Location	load_fp8_state_dict
Confidence	98%

Problem:
The function load_fp8_state_dict calls model.load_state_dict(prepared, strict=False, assign=assign). The torch.nn.Module.load_state_dict API does not accept an 'assign' keyword argument (its signature is load_state_dict(state_dict, strict=True)). Passing this unsupported argument will raise a TypeError at runtime, preventing FP8 checkpoints from being loaded.

Suggested Fix:
Remove the unsupported 'assign' argument and handle assignment manually if needed. For example:

missing, unexpected = model.load_state_dict(prepared, strict=False)
if assign:
    for name, tensor in prepared.items():
        module = model.get_submodule(name.rpartition('.')[0]) if '.' in name else model
        leaf = name.rpartition('.')[2]
        module.register_buffer(leaf, tensor, persistent=True)
``` or adjust the logic to use the standard load_state_dict behavior without the 'assign' flag.

---

### 8. 🐛 Potential KeyError in ideogram_magic_prompt function

| Field | Details |
|---|---|
| **Severity** | 🟡 Medium |
| **Type** | Bug |
| **File** | `src/ideogram4/magic_prompt.py` |
| **Location** | ideogram_magic_prompt function |
| **Confidence** | 92% |

**Problem:**
The ideogram_magic_prompt function does not check if the 'json_prompt' key exists in the response data before trying to access it. If the 'json_prompt' key does not exist, a KeyError will be raised.

**Suggested Fix:**
Add a check to ensure the 'json_prompt' key exists in the response data before trying to access it.

---

<details>
<summary>About this report</summary>

This report was generated using Llama 3.3 70B.
Only findings with ≥80% confidence are included.
False positives are possible — use your own judgment.

</details>