fix: handle input images in image gen for replicate by TejasGhatte · Pull Request #5454 · maximhq/bifrost

TejasGhatte · 2026-07-22T06:05:38Z

Summary

ToReplicateImageGenerationInput previously ignored InputImages entirely and could not surface validation errors to callers. This PR adds input image support to the image generation path (mirroring what already existed for image edits) and propagates URL sanitization errors instead of silently dropping them.

Changes

Changed ToReplicateImageGenerationInput to return (*ReplicatePredictionRequest, error) so URL validation errors can be surfaced to callers.
Added InputImages handling in the generation path: each image URL is sanitized via schemas.SanitizeImageURL, and the resulting slice is routed to the correct model-specific field using the new shared helper.
Extracted the model-to-field dispatch logic (image_prompt, input_image, image, input_images) into a setInputImageField helper, eliminating the duplicated switch block that previously existed only in the edit path.
Updated both ImageGeneration and ImageGenerationStream call sites to handle the new error return.
Removed stale line-number references from the Replicate provider docs.

Type of change

Affected areas

How to test

go test ./core/providers/replicate/...

New test cases cover:

InputImages_SingleImageField — verifies that a kontext-pro model receives the first image in input_image and no other image fields are set.
InputImages_ArrayFieldWithBase64Normalization — verifies that a generic model receives all images in input_images and that bare base64 strings are prefixed with the data:image/png;base64, URI scheme.
InputImages_InvalidURL — verifies that a file:// URI causes an error return and a nil result.

Breaking changes

Yes
No

ToReplicateImageGenerationInput now returns (*ReplicatePredictionRequest, error) instead of *ReplicatePredictionRequest. Any external callers must be updated to handle the additional return value.

Related issues

Security considerations

Input images are now validated through schemas.SanitizeImageURL before being forwarded to Replicate. This prevents schemes such as file:// from being passed through to the provider.

Checklist

I read docs/contributing/README.md and followed the guidelines
I added/updated tests where appropriate
I updated documentation where needed
I verified builds succeed (Go and UI)
I verified the CI pipeline passes locally if applicable

coderabbitai · 2026-07-22T06:05:45Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 914cddeb-5a9f-4cb4-a882-b880fecc2632

📥 Commits

Reviewing files that changed from the base of the PR and between 081f2d8 and 92d2c39.

📒 Files selected for processing (4)

core/providers/replicate/images.go
core/providers/replicate/replicate.go
core/providers/replicate/replicate_test.go
docs/providers/supported-providers/replicate.mdx

💤 Files with no reviewable changes (1)

docs/providers/supported-providers/replicate.mdx

🚧 Files skipped from review as they are similar to previous changes (3)

core/providers/replicate/replicate.go
core/providers/replicate/replicate_test.go
core/providers/replicate/images.go

📝 Walkthrough

Summary by CodeRabbit

New Features
- Improved image generation and editing support across Replicate models.
- Multiple input images are now mapped to the appropriate model fields.
- Base64 image inputs are normalized for compatibility.
Bug Fixes
- Invalid image URLs are rejected with clear errors instead of producing malformed requests.
- Image generation requests now stop safely when input conversion fails.
Documentation
- Simplified Replicate provider caveats by removing internal code references while preserving behavioral details.

Walkthrough

Replicate image generation now sanitizes InputImages, maps them to model-specific fields through shared logic, and returns conversion errors. Streaming and non-streaming callers propagate those errors. Tests cover mapping, base64 normalization, and invalid URLs; related documentation references are removed.

Changes

Replicate image input handling

Layer / File(s)	Summary
Image conversion and field mapping `core/providers/replicate/images.go`	Generation input conversion sanitizes image URLs and assigns them to `ImagePrompt`, `InputImage`, `Image`, or `InputImages`; edit conversion reuses the shared assignment helper.
Generation error propagation `core/providers/replicate/replicate.go`	Image generation and streaming now handle and return conversion errors from `ToReplicateImageGenerationInput`.
Conversion tests and documentation `core/providers/replicate/replicate_test.go`, `docs/providers/supported-providers/replicate.mdx`	Tests cover input-image mapping, base64 normalization, invalid URLs, and error returns; Replicate caveat code references are removed.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant BifrostImageGenerationRequest
  participant ToReplicateImageGenerationInput
  participant SanitizeImageURL
  participant ReplicatePredictionRequest
  BifrostImageGenerationRequest->>ToReplicateImageGenerationInput: provide Params.InputImages
  ToReplicateImageGenerationInput->>SanitizeImageURL: sanitize each image URL
  SanitizeImageURL-->>ToReplicateImageGenerationInput: sanitized URL or error
  ToReplicateImageGenerationInput->>ReplicatePredictionRequest: assign images to mapped input field

Suggested reviewers: pratham-mishra04, akshaydeo

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and matches the main change: adding input image handling for Replicate image generation.
Description check	✅ Passed	The description follows the template well and covers the key sections with concrete details about changes, testing, and breaking impact.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch 07-21-fix_handle_input_images_in_image_gen_for_replicate

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

_{Comment @coderabbitai help to get the list of available commands.}

CLAassistant · 2026-07-22T06:05:48Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

tejas ghatte seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

TejasGhatte · 2026-07-22T06:05:52Z

fix: gemini tool schema integer constraint #5466
fix: handle input images in image gen for replicate #5454 👈 (View in Graphite)
dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

coderabbitai

🧹 Nitpick comments (1)

core/providers/replicate/replicate_test.go (1)
609-662: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cover the remaining model-specific fields.

These cases cover input_image and fallback input_images, but not the new image_prompt and image branches. Add one case for each to lock the Flux model mappings.

As per coding guidelines, “table-driven coverage for behavior changes.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/providers/replicate/replicate_test.go` around lines 609 - 662, Add
table-driven test cases covering the `image_prompt` and `image` branches in the
existing image-generation test table, using the relevant Flux model mappings and
asserting each field is populated with the expected value while the alternative
image fields remain unset. Follow the existing `InputImages_*` case structure
and validation style.
Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@core/providers/replicate/replicate_test.go`:
- Around line 609-662: Add table-driven test cases covering the `image_prompt`
and `image` branches in the existing image-generation test table, using the
relevant Flux model mappings and asserting each field is populated with the
expected value while the alternative image fields remain unset. Follow the
existing `InputImages_*` case structure and validation style.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: b8c9c116-86a6-49cb-b05c-63b5e8f06d4c

📥 Commits

Reviewing files that changed from the base of the PR and between 9b02739 and 081f2d8.

📒 Files selected for processing (4)

core/providers/replicate/images.go
core/providers/replicate/replicate.go
core/providers/replicate/replicate_test.go
docs/providers/supported-providers/replicate.mdx

💤 Files with no reviewable changes (1)

docs/providers/supported-providers/replicate.mdx

akshaydeo · 2026-07-22T15:18:42Z

Merge activity

Jul 22, 3:18 PM UTC: A user started a stack merge that includes this pull request via Graphite.
Jul 22, 3:19 PM UTC: @akshaydeo merged this pull request with Graphite.

## Summary `ToReplicateImageGenerationInput` previously ignored `InputImages` entirely and could not surface validation errors to callers. This PR adds input image support to the image generation path (mirroring what already existed for image edits) and propagates URL sanitization errors instead of silently dropping them. ## Changes - Changed `ToReplicateImageGenerationInput` to return `(*ReplicatePredictionRequest, error)` so URL validation errors can be surfaced to callers. - Added `InputImages` handling in the generation path: each image URL is sanitized via `schemas.SanitizeImageURL`, and the resulting slice is routed to the correct model-specific field using the new shared helper. - Extracted the model-to-field dispatch logic (`image_prompt`, `input_image`, `image`, `input_images`) into a `setInputImageField` helper, eliminating the duplicated switch block that previously existed only in the edit path. - Updated both `ImageGeneration` and `ImageGenerationStream` call sites to handle the new error return. - Removed stale line-number references from the Replicate provider docs. ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [x] Docs ## How to test ```sh go test ./core/providers/replicate/... ``` New test cases cover: - `InputImages_SingleImageField` — verifies that a kontext-pro model receives the first image in `input_image` and no other image fields are set. - `InputImages_ArrayFieldWithBase64Normalization` — verifies that a generic model receives all images in `input_images` and that bare base64 strings are prefixed with the `data:image/png;base64,` URI scheme. - `InputImages_InvalidURL` — verifies that a `file://` URI causes an error return and a `nil` result. ## Breaking changes - [x] Yes - [ ] No `ToReplicateImageGenerationInput` now returns `(*ReplicatePredictionRequest, error)` instead of `*ReplicatePredictionRequest`. Any external callers must be updated to handle the additional return value. ## Related issues ## Security considerations Input images are now validated through `schemas.SanitizeImageURL` before being forwarded to Replicate. This prevents schemes such as `file://` from being passed through to the provider. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [x] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable

TejasGhatte marked this pull request as ready for review July 22, 2026 06:05

coderabbitai Bot requested review from Pratham-Mishra04 and akshaydeo July 22, 2026 06:06

coderabbitai Bot reviewed Jul 22, 2026

View reviewed changes

coderabbitai Bot approved these changes Jul 22, 2026

View reviewed changes

fix: handle input images in image gen for replicate

92d2c39

TejasGhatte force-pushed the 07-21-fix_handle_input_images_in_image_gen_for_replicate branch from 081f2d8 to 92d2c39 Compare July 22, 2026 12:16

TejasGhatte mentioned this pull request Jul 22, 2026

fix: gemini tool schema integer constraint #5466

Merged

18 tasks

akshaydeo approved these changes Jul 22, 2026

View reviewed changes

akshaydeo merged commit 37e15b0 into dev Jul 22, 2026
13 of 14 checks passed

akshaydeo deleted the 07-21-fix_handle_input_images_in_image_gen_for_replicate branch July 22, 2026 15:19

This was referenced Jul 29, 2026

[Bug]: defaultImageDetail fix (#5352) doesn't cover images inside function_call_output (tool results), Responses API to self-host vLLM still fails #5649

Open

[fix]: WeightedRandom key selector - fractional weights below 0.01 no longer starve #5490

Open

Navigation Menu

fix: handle input images in image gen for replicate - #5454

fix: handle input images in image gen for replicate#5454
akshaydeo merged 1 commit into
devfrom
07-21-fix_handle_input_images_in_image_gen_for_replicate

TejasGhatte commented Jul 22, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

CLAassistant commented Jul 22, 2026

Uh oh!

TejasGhatte commented Jul 22, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

akshaydeo commented Jul 22, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

TejasGhatte commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Type of change

Affected areas

How to test

Breaking changes

Related issues

Security considerations

Checklist

Uh oh!

coderabbitai Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

CLAassistant commented Jul 22, 2026

Uh oh!

TejasGhatte commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

akshaydeo commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge activity

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TejasGhatte commented Jul 22, 2026 •

edited

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading

TejasGhatte commented Jul 22, 2026 •

edited

Loading

akshaydeo commented Jul 22, 2026 •

edited

Loading