`KeyError: 'sample'` intermittent failure on black-forest-labs/flux-2-pro predictions

[matthewwallace1515]

Bug Description

The black-forest-labs/flux-2-pro model intermittently fails with KeyError: 'sample' after successful image generation. The failure rate is approximately 30-50%.

Observed Behavior

• Image successfully generates (visible in Replicate web UI)
• Prediction status shows succeeded in the dashboard
• But API call returns KeyError: 'sample' error
• The generated image URL is accessible despite the error

To Reproduce

import replicate

# Run multiple times - approximately 30-50% will fail
output = replicate.run(
    "black-forest-labs/flux-2-pro",
    input={
        "prompt": "Abstract geometric pattern, voronoi cells",
        "aspect_ratio": "16:9",
        "resolution": "1 MP",
        "output_format": "jpeg",
        "output_quality": 100,
        "safety_tolerance": 5
    }
)

Error Details

The error occurs in the prediction result handling - the image generation itself completes successfully, but parsing the response fails with:

KeyError: 'sample'

Environment

• replicate Python package
• Python 3.12
• Observed: January 2026

Workaround

Implemented retry logic with 2-3 attempts and 10-second delays between retries. This successfully handles the intermittent failures, suggesting the issue is in result parsing rather than actual image generation.

Analysis

The bug appears to occur in Replicate's prediction result handling/wrapper code, not in the Flux model itself:

1. The image generation completes successfully (confirmed via web UI)
2. The prediction status is succeeded
3. The error occurs when extracting the result from the response

This is a production blocker for automated pipelines that rely on Flux 2 Pro.