🤖 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 — 🟠 6 high · 🟡 2 medium

---

1. 🐛 Incompatible dtype on unsupported devices

Field	Details
Severity	🟠 High
Type	Bug
File	run_inference.py
Location	pipe = Ideogram4Pipeline.from_pretrained(..., dtype=torch.bfloat16, ...)
Confidence	96%

Problem:
The code forces the pipeline to use torch.bfloat16 regardless of the selected device. Many CUDA GPUs, Apple MPS, and CPUs do not support bfloat16 tensors, causing a RuntimeError at runtime and preventing inference from running.

Suggested Fix:
Select the dtype based on device capabilities, e.g., use torch.bfloat16 only if torch.backends.cuda.is_bf16_supported() or torch.backends.mps.is_bf16_supported(), otherwise fall back to torch.float16 or torch.float32.

---

2. 🔒 API key exposed via command‑line argument

Field	Details
Severity	🟠 High
Type	Security
File	run_inference.py
Location	--magic-prompt-key argument definition
Confidence	95%

Problem:
The script accepts the magic‑prompt API key through a command‑line option. Command‑line arguments are visible to other users on the system via process listings, potentially leaking the secret key.

Suggested Fix:
Remove the --magic-prompt-key argument and require the key to be provided only via environment variables (e.g., MAGIC_PROMPT_API_KEY). If a CLI option is necessary, mark it as hidden and document that it should be used with caution, or read the key from a secure file instead.

---

3. 🐛 Use of assert for runtime validation

Field	Details
Severity	🟠 High
Type	Bug
File	src/ideogram4/latent_norm.py
Location	def get_latent_norm()
Confidence	96%

Problem:
The function validates the shape of the tensors with an assert statement. In optimized Python runs (e.g., using the -O flag), assert statements are stripped, so the shape check disappears. If LATENT_SHIFT or LATENT_SCALE ever have a length different from 128, the function would return tensors of incorrect shape without any error, potentially causing downstream failures.

Suggested Fix:
Replace the assert with an explicit runtime check that raises an exception regardless of interpreter optimizations, e.g.,:

if shift.shape != (128,) or scale.shape != (128,):
    raise ValueError(f"Expected tensors of shape (128,), got {shift.shape} and {scale.shape}")

---

4. 🐛 Incorrect attention mask semantics

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

Problem:
The scaled_dot_product_attention function expects a boolean mask where True indicates positions to be ignored. The code builds attn_mask with True for positions belonging to the same segment (i.e., positions that should be attended), effectively blocking attention to the intended tokens and allowing attention across different segments. This will cause incorrect model behavior and likely training divergence.

Suggested Fix:
Invert the mask before passing it to scaled_dot_product_attention, e.g., attn_mask = ~ (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1) or construct the mask with False for same‑segment positions.

---

5. 🐛 Rotary positional embedding frequency computation drops sequence dimension

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

Problem:
The implementation computes freqs = inv_freq @ pos.unsqueeze(2) and then transposes dimensions, resulting in a tensor of shape (3, B, 1, inv_freq_size) instead of the expected (3, B, L, inv_freq_size). The matrix multiplication collapses the sequence length dimension, so each token receives the same positional frequencies, breaking the intended MRoPE behavior.

Suggested Fix:
Replace the matrix multiplication with an element‑wise multiplication that preserves the sequence dimension, e.g., freqs = inv_freq * pos.unsqueeze(-1) (after appropriate broadcasting) or compute freqs = pos[..., None] * inv_freq[None, None, :, None] and then reshape to (3, B, L, inv_freq_size). Ensure the resulting emb has the correct shape before applying cos and sin.

---

6. 🔒 Use of trust_remote_code=True enables remote code execution

Field	Details
Severity	🟠 High
Type	Security
File	src/ideogram4/pipeline_ideogram4.py
Location	_load_qwen3_vl
Confidence	99%

Problem:
The function _load_qwen3_vl loads model and config with AutoConfig.from_pretrained and AutoModel.from_pretrained using trust_remote_code=True. This allows execution of arbitrary Python code from the Hugging Face repository, which can be exploited to run malicious code on the host system if an attacker can influence the repo ID or subfolder. This is a real security risk because the code is executed without validation.

Suggested Fix:
Remove trust_remote_code=True or restrict its use to vetted, trusted repositories. If custom code is required, load it in a sandboxed environment or perform explicit verification of the downloaded files before execution.

---

7. 🐛 IndexError and StopIteration when ch_mult is empty

Field	Details
Severity	🟡 Medium
Type	Bug
File	src/ideogram4/autoencoder.py
Location	Decoder.init
Confidence	95%

Problem:
The Decoder constructor accesses ch_mult[self.num_resolutions - 1] without checking that ch_mult is non‑empty. If an empty list is passed, this raises an IndexError, and later next(self.up.parameters()) raises StopIteration, causing the model initialization to fail.

Suggested Fix:
Validate that ch_mult is non‑empty (e.g., raise a clear ValueError) or provide a default fallback. Example: if not ch_mult: raise ValueError('ch_mult must contain at least one element'); then proceed with the existing logic.

---

8. 🐛 Hex color validation rejects lowercase letters

Field	Details
Severity	🟡 Medium
Type	Bug
File	src/ideogram4/caption_verifier.py
Location	_verify_color_palette
Confidence	99%

Problem:
The _verify_color_palette method validates hex colors by checking each character against the string "0123456789ABCDEF". This only allows uppercase A-F, so valid lowercase hex colors (e.g., "#ff00aa") are incorrectly flagged as invalid, causing false warnings for correct input.

Suggested Fix:
Modify the validation to accept both uppercase and lowercase hex digits, e.g., by checking c.upper() in "0123456789ABCDEF" or using a case‑insensitive regex like re.fullmatch(r"#[0-9a-fA-F]{6}", color).

---

About this report

This report was generated using Advanced AI models.
Only findings with ≥90% confidence are included.
False positives are possible — use your own judgment.