From c4bf19740d4457f5b1bdd7b61a9dd70f89bdf399 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Wed, 12 Aug 2026 04:29:18 +0530 Subject: [PATCH 1/2] feat: add megapixel-banded output image cost tiers --- docs/architecture/framework/model-catalog.mdx | 5 ++ docs/openapi/openapi.json | 20 ++++++ .../schemas/management/governance.yaml | 15 ++++ docs/providers/custom-pricing.mdx | 5 ++ framework/configstore/migrations.go | 44 ++++++++++++ framework/configstore/rdb.go | 5 ++ framework/configstore/rdb_test.go | 50 +++++++++++++ framework/configstore/tables/modelpricing.go | 5 ++ framework/modelcatalog/datasheet/cost.go | 72 ++++++++++++++----- framework/modelcatalog/datasheet/cost_test.go | 50 +++++++++++++ framework/modelcatalog/datasheet/overrides.go | 5 ++ .../modelcatalog/datasheet/overrides_test.go | 27 +++++++ framework/modelcatalog/datasheet/types.go | 15 ++++ .../overrides/pricingFields.test.ts | 2 +- .../custom-pricing/overrides/pricingFields.ts | 30 ++++++++ ui/lib/types/governance.ts | 5 ++ 16 files changed, 338 insertions(+), 17 deletions(-) diff --git a/docs/architecture/framework/model-catalog.mdx b/docs/architecture/framework/model-catalog.mdx index 38292f7e746..506c58e0f39 100644 --- a/docs/architecture/framework/model-catalog.mdx +++ b/docs/architecture/framework/model-catalog.mdx @@ -158,6 +158,11 @@ type PricingEntry struct { OutputCostPerImageAbove1024x1024PixelsPremium *float64 `json:"output_cost_per_image_above_1024_and_1024_pixels_and_premium_image,omitempty"` OutputCostPerImageAbove2048x2048Pixels *float64 `json:"output_cost_per_image_above_2048_and_2048_pixels,omitempty"` OutputCostPerImageAbove4096x4096Pixels *float64 `json:"output_cost_per_image_above_4096_and_4096_pixels,omitempty"` + OutputCostPerImageAbove4Megapixels *float64 `json:"output_cost_per_image_above_4_megapixels,omitempty"` + OutputCostPerImageAbove8Megapixels *float64 `json:"output_cost_per_image_above_8_megapixels,omitempty"` + OutputCostPerImageAbove16Megapixels *float64 `json:"output_cost_per_image_above_16_megapixels,omitempty"` + OutputCostPerImageAbove32Megapixels *float64 `json:"output_cost_per_image_above_32_megapixels,omitempty"` + OutputCostPerImageAbove64Megapixels *float64 `json:"output_cost_per_image_above_64_megapixels,omitempty"` OutputCostPerImageLowQuality *float64 `json:"output_cost_per_image_low_quality,omitempty"` OutputCostPerImageMediumQuality *float64 `json:"output_cost_per_image_medium_quality,omitempty"` OutputCostPerImageHighQuality *float64 `json:"output_cost_per_image_high_quality,omitempty"` diff --git a/docs/openapi/openapi.json b/docs/openapi/openapi.json index 920a0870027..8aba32c05ea 100644 --- a/docs/openapi/openapi.json +++ b/docs/openapi/openapi.json @@ -96656,6 +96656,26 @@ "type": "number", "minimum": 0 }, + "output_cost_per_image_above_4_megapixels": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_above_8_megapixels": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_above_16_megapixels": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_above_32_megapixels": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_above_64_megapixels": { + "type": "number", + "minimum": 0 + }, "output_cost_per_image_low_quality": { "type": "number", "minimum": 0 diff --git a/docs/openapi/schemas/management/governance.yaml b/docs/openapi/schemas/management/governance.yaml index 5cf59690c04..a777860ff6c 100644 --- a/docs/openapi/schemas/management/governance.yaml +++ b/docs/openapi/schemas/management/governance.yaml @@ -1861,6 +1861,21 @@ PricingPatch: output_cost_per_image_above_4096_and_4096_pixels: type: number minimum: 0 + output_cost_per_image_above_4_megapixels: + type: number + minimum: 0 + output_cost_per_image_above_8_megapixels: + type: number + minimum: 0 + output_cost_per_image_above_16_megapixels: + type: number + minimum: 0 + output_cost_per_image_above_32_megapixels: + type: number + minimum: 0 + output_cost_per_image_above_64_megapixels: + type: number + minimum: 0 output_cost_per_image_low_quality: type: number minimum: 0 diff --git a/docs/providers/custom-pricing.mdx b/docs/providers/custom-pricing.mdx index 414cf16860d..36e14bcc41a 100644 --- a/docs/providers/custom-pricing.mdx +++ b/docs/providers/custom-pricing.mdx @@ -376,6 +376,11 @@ Any field you set (including `0`) is applied as an override; omitted fields are | `output_cost_per_image_above_1024_and_1024_pixels` | Generated image at or above 1024×1024 | | `output_cost_per_image_above_2048_and_2048_pixels` | Generated image at or above 2048×2048 | | `output_cost_per_image_above_4096_and_4096_pixels` | Generated image at or above 4096×4096 | +| `output_cost_per_image_above_4_megapixels` | Generated image output at or above 4 megapixels | +| `output_cost_per_image_above_8_megapixels` | Generated image output at or above 8 megapixels | +| `output_cost_per_image_above_16_megapixels` | Generated image output at or above 16 megapixels | +| `output_cost_per_image_above_32_megapixels` | Generated image output at or above 32 megapixels | +| `output_cost_per_image_above_64_megapixels` | Generated image output at or above 64 megapixels | | `output_cost_per_image_premium_image` | Generated image - premium image | | `output_cost_per_image_above_512_and_512_pixels_and_premium_image` | Generated image at or above 512×512, premium image | | `output_cost_per_image_above_1024_and_1024_pixels_and_premium_image` | Generated image at or above 1024×1024, premium image | diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 0a159b516ea..47fc09c5020 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -469,6 +469,7 @@ var configstoreMigrationSteps = []migrationStep{ {IDs: []string{"add_bedrock_endpoints_columns"}, run: migrationAddBedrockEndpointsColumns}, {IDs: []string{"add_cost_per_request_pricing_column"}, run: migrationAddCostPerRequestPricingColumn}, {IDs: []string{"add_notifications_table"}, run: migrationAddNotificationsTable}, + {IDs: []string{"add_image_megapixel_tier_pricing_columns"}, run: migrationAddImageMegapixelTierPricingColumns}, } func migrationAddNotificationsTable(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { @@ -11936,3 +11937,46 @@ func migrationAddBedrockEndpointsColumns(ctx context.Context, db *gorm.DB, logge } return nil } + +// migrationAddImageMegapixelTierPricingColumns adds the megapixel-banded output +// image cost tier columns (output_cost_per_image_above_{4,8,16,32,64}_megapixels), +// used by providers (e.g. Replicate's upscaler models) that publish tiered +// per-image pricing by total output megapixels rather than by a squared +// width/height threshold. +func migrationAddImageMegapixelTierPricingColumns(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "add_image_megapixel_tier_pricing_columns" + logger.Info("[configstore] starting migration %s", migrationName) + defer logger.Info("[configstore] finished migration %s", migrationName) + columns := []string{ + "output_cost_per_image_above_4_megapixels", + "output_cost_per_image_above_8_megapixels", + "output_cost_per_image_above_16_megapixels", + "output_cost_per_image_above_32_megapixels", + "output_cost_per_image_above_64_megapixels", + } + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + for _, field := range columns { + if err := addColumnIfNotExists(tx, logger, &tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to add column %s: %w", field, err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + for _, field := range columns { + if err := dropColumnIfExists(tx, logger, &tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to drop column %s: %w", field, err) + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_image_megapixel_tier_pricing_columns migration: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 3f878848ac1..275847f9051 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -2783,6 +2783,11 @@ var pricingSyncUpdateColumns = []string{ "output_cost_per_image_above_1024x1024_pixels_premium", "output_cost_per_image_above_2048_and_2048_pixels", "output_cost_per_image_above_4096_and_4096_pixels", + "output_cost_per_image_above_4_megapixels", + "output_cost_per_image_above_8_megapixels", + "output_cost_per_image_above_16_megapixels", + "output_cost_per_image_above_32_megapixels", + "output_cost_per_image_above_64_megapixels", "output_cost_per_image_low_quality", "output_cost_per_image_medium_quality", "output_cost_per_image_high_quality", diff --git a/framework/configstore/rdb_test.go b/framework/configstore/rdb_test.go index ba0d779ed15..4605d30a74d 100644 --- a/framework/configstore/rdb_test.go +++ b/framework/configstore/rdb_test.go @@ -2762,6 +2762,56 @@ func TestUpsertModelPricesBatch_SQLite(t *testing.T) { assert.InDelta(t, 0.000005, *updated.InputCostPerToken, 1e-9) } +func TestUpsertModelPricesBatch_MegapixelImageTierColumns_SurviveResync(t *testing.T) { + // Regression test for pricingSyncUpdateColumns: a column present on + // TableModelPricing but missing from that explicit update-column list + // would insert fine on the first sync (Create writes every column) but + // silently revert to null on the second sync (ON CONFLICT DO UPDATE only + // touches listed columns). + s := setupRDBTestStore(t) + require.NoError(t, s.DB().AutoMigrate(&tables.TableModelPricing{})) + + ctx := context.Background() + cost := func(f float64) *float64 { return &f } + + pricing := []tables.TableModelPricing{ + { + Model: "prunaai/p-image-upscale", + Provider: "replicate", + Mode: "image_generation", + OutputCostPerImage: cost(0.005), + OutputCostPerImageAbove4Megapixels: cost(0.01), + OutputCostPerImageAbove8Megapixels: cost(0.02), + OutputCostPerImageAbove16Megapixels: cost(0.04), + OutputCostPerImageAbove32Megapixels: cost(0.06), + OutputCostPerImageAbove64Megapixels: cost(0.12), + }, + } + + require.NoError(t, s.UpsertModelPricesBatch(ctx, pricing)) + + // Re-upsert the same row (simulating the next scheduled datasheet sync) + // with a changed tier value to exercise the ON CONFLICT update path. + pricing[0].OutputCostPerImageAbove16Megapixels = cost(0.05) + require.NoError(t, s.UpsertModelPricesBatch(ctx, pricing)) + + got, err := s.GetModelPrices(ctx) + require.NoError(t, err) + require.Len(t, got, 1) + + row := got[0] + require.NotNil(t, row.OutputCostPerImageAbove4Megapixels) + require.NotNil(t, row.OutputCostPerImageAbove8Megapixels) + require.NotNil(t, row.OutputCostPerImageAbove16Megapixels) + require.NotNil(t, row.OutputCostPerImageAbove32Megapixels) + require.NotNil(t, row.OutputCostPerImageAbove64Megapixels) + assert.InDelta(t, 0.01, *row.OutputCostPerImageAbove4Megapixels, 1e-9) + assert.InDelta(t, 0.02, *row.OutputCostPerImageAbove8Megapixels, 1e-9) + assert.InDelta(t, 0.05, *row.OutputCostPerImageAbove16Megapixels, 1e-9) // survived resync with the updated value + assert.InDelta(t, 0.06, *row.OutputCostPerImageAbove32Megapixels, 1e-9) + assert.InDelta(t, 0.12, *row.OutputCostPerImageAbove64Megapixels, 1e-9) +} + func TestUpsertModelParametersBatch_SQLite(t *testing.T) { s := setupRDBTestStore(t) require.NoError(t, s.DB().AutoMigrate(&tables.TableModelParameters{})) diff --git a/framework/configstore/tables/modelpricing.go b/framework/configstore/tables/modelpricing.go index a37e62ffef3..76556fe403b 100644 --- a/framework/configstore/tables/modelpricing.go +++ b/framework/configstore/tables/modelpricing.go @@ -90,6 +90,11 @@ type TableModelPricing struct { OutputCostPerImageAbove1024x1024PixelsPremium *float64 `gorm:"default:null;column:output_cost_per_image_above_1024x1024_pixels_premium" json:"output_cost_per_image_above_1024_and_1024_pixels_and_premium_image,omitempty"` OutputCostPerImageAbove2048x2048Pixels *float64 `gorm:"default:null;column:output_cost_per_image_above_2048_and_2048_pixels" json:"output_cost_per_image_above_2048_and_2048_pixels,omitempty"` OutputCostPerImageAbove4096x4096Pixels *float64 `gorm:"default:null;column:output_cost_per_image_above_4096_and_4096_pixels" json:"output_cost_per_image_above_4096_and_4096_pixels,omitempty"` + OutputCostPerImageAbove4Megapixels *float64 `gorm:"default:null;column:output_cost_per_image_above_4_megapixels" json:"output_cost_per_image_above_4_megapixels,omitempty"` + OutputCostPerImageAbove8Megapixels *float64 `gorm:"default:null;column:output_cost_per_image_above_8_megapixels" json:"output_cost_per_image_above_8_megapixels,omitempty"` + OutputCostPerImageAbove16Megapixels *float64 `gorm:"default:null;column:output_cost_per_image_above_16_megapixels" json:"output_cost_per_image_above_16_megapixels,omitempty"` + OutputCostPerImageAbove32Megapixels *float64 `gorm:"default:null;column:output_cost_per_image_above_32_megapixels" json:"output_cost_per_image_above_32_megapixels,omitempty"` + OutputCostPerImageAbove64Megapixels *float64 `gorm:"default:null;column:output_cost_per_image_above_64_megapixels" json:"output_cost_per_image_above_64_megapixels,omitempty"` OutputCostPerImageLowQuality *float64 `gorm:"default:null;column:output_cost_per_image_low_quality" json:"output_cost_per_image_low_quality,omitempty"` OutputCostPerImageMediumQuality *float64 `gorm:"default:null;column:output_cost_per_image_medium_quality" json:"output_cost_per_image_medium_quality,omitempty"` OutputCostPerImageHighQuality *float64 `gorm:"default:null;column:output_cost_per_image_high_quality" json:"output_cost_per_image_high_quality,omitempty"` diff --git a/framework/modelcatalog/datasheet/cost.go b/framework/modelcatalog/datasheet/cost.go index 5e47e94e210..643c407bbda 100644 --- a/framework/modelcatalog/datasheet/cost.go +++ b/framework/modelcatalog/datasheet/cost.go @@ -865,22 +865,7 @@ func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageU } } if perImageRate == nil { - const pixels512x512 = 512 * 512 - const pixels1024x1024 = 1024 * 1024 - const pixels2048x2048 = 2048 * 2048 - const pixels4096x4096 = 4096 * 4096 - switch { - case pixels >= pixels4096x4096 && pricing.OutputCostPerImageAbove4096x4096Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove4096x4096Pixels - case pixels >= pixels2048x2048 && pricing.OutputCostPerImageAbove2048x2048Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove2048x2048Pixels - case pixels >= pixels1024x1024 && pricing.OutputCostPerImageAbove1024x1024Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove1024x1024Pixels - case pixels >= pixels512x512 && pricing.OutputCostPerImageAbove512x512Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove512x512Pixels - default: - perImageRate = pricing.OutputCostPerImage - } + perImageRate = selectImageSizeTierRate(pricing, pixels) } if perImageRate != nil { return float64(numOutputImages) * *perImageRate @@ -889,6 +874,61 @@ func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageU return 0 } +// imageSizeTier is one step of the per-image size-tier ladder: the pixel +// count above which the given rate applies. +type imageSizeTier struct { + pixels int + rate *float64 +} + +// selectImageSizeTierRate picks the flat per-image rate for the given output +// pixel count, from whichever size-tier fields a model's pricing row has +// populated. Two independent tier families exist because providers publish +// resolution-based pricing in different units: some tier by exact +// width×height threshold (output_cost_per_image_above_x_pixels), +// others (e.g. Replicate's upscaler models) tier by total output megapixels +// (output_cost_per_image_above__megapixels). A given model is expected to +// populate only one family; both are checked here, interleaved by their +// actual pixel threshold largest-first, so either (or in principle both) +// resolve correctly without the caller needing to know which unit a model +// uses. Falls back to pricing.OutputCostPerImage (via the nil return) when +// no tier's threshold is met or no tier fields are populated. +func selectImageSizeTierRate(pricing *configstoreTables.TableModelPricing, pixels int) *float64 { + const ( + pixels512x512 = 512 * 512 + pixels1024x1024 = 1024 * 1024 + pixels2048x2048 = 2048 * 2048 + pixels4Megapixels = 4_000_000 + pixels4096x4096 = 4096 * 4096 + pixels8Megapixels = 8_000_000 + pixels16Megapixels = 16_000_000 + pixels32Megapixels = 32_000_000 + pixels64Megapixels = 64_000_000 + ) + // Ordered by threshold, largest first, so the first match wins. + // NOTE: ordered strictly by actual pixel threshold, largest first — NOT + // by field family — since 4096x4096 (16,777,216px) falls between the 16MP + // and 32MP megapixel thresholds, and 2048x2048 (4,194,304px) falls just + // above the 4MP threshold. + tiers := []imageSizeTier{ + {pixels64Megapixels, pricing.OutputCostPerImageAbove64Megapixels}, + {pixels32Megapixels, pricing.OutputCostPerImageAbove32Megapixels}, + {pixels4096x4096, pricing.OutputCostPerImageAbove4096x4096Pixels}, + {pixels16Megapixels, pricing.OutputCostPerImageAbove16Megapixels}, + {pixels8Megapixels, pricing.OutputCostPerImageAbove8Megapixels}, + {pixels2048x2048, pricing.OutputCostPerImageAbove2048x2048Pixels}, + {pixels4Megapixels, pricing.OutputCostPerImageAbove4Megapixels}, + {pixels1024x1024, pricing.OutputCostPerImageAbove1024x1024Pixels}, + {pixels512x512, pricing.OutputCostPerImageAbove512x512Pixels}, + } + for _, t := range tiers { + if pixels >= t.pixels && t.rate != nil { + return t.rate + } + } + return pricing.OutputCostPerImage +} + // computeVideoCost handles video generation requests. // Input and output are calculated independently — tokens first, then per-second fallback. func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, videoSeconds *int, tier serviceTier) float64 { diff --git a/framework/modelcatalog/datasheet/cost_test.go b/framework/modelcatalog/datasheet/cost_test.go index 7eafd737589..15de59a2223 100644 --- a/framework/modelcatalog/datasheet/cost_test.go +++ b/framework/modelcatalog/datasheet/cost_test.go @@ -1387,6 +1387,56 @@ func TestComputeImageCost_QualityBasedRates(t *testing.T) { assert.InDelta(t, 0.05, computeImageCost(&p, usage, "", "", serviceTier{}), 1e-12) } +func TestComputeImageCost_MegapixelTier_SelectsCorrectBand(t *testing.T) { + // Mirrors replicate/prunaai/p-image-upscale's real tier structure. + p := configstoreTables.TableModelPricing{ + OutputCostPerImage: bifrost.Ptr(0.005), + OutputCostPerImageAbove4Megapixels: bifrost.Ptr(0.01), + OutputCostPerImageAbove8Megapixels: bifrost.Ptr(0.02), + OutputCostPerImageAbove16Megapixels: bifrost.Ptr(0.04), + OutputCostPerImageAbove32Megapixels: bifrost.Ptr(0.06), + OutputCostPerImageAbove64Megapixels: bifrost.Ptr(0.12), + } + usage := &schemas.ImageUsage{ + OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 1}, + } + + // 10MP output (between the 8MP and 16MP thresholds) → $0.02 tier. + cost := computeImageCost(&p, usage, "1x10000000", "", serviceTier{}) + assert.InDelta(t, 0.02, cost, 1e-12) + + // 2MP output (below the lowest 4MP threshold) → falls back to base rate. + cost = computeImageCost(&p, usage, "1x2000000", "", serviceTier{}) + assert.InDelta(t, 0.005, cost, 1e-12) + + // 70MP output (above every threshold) → top $0.12 tier. + cost = computeImageCost(&p, usage, "1x70000000", "", serviceTier{}) + assert.InDelta(t, 0.12, cost, 1e-12) +} + +func TestComputeImageCost_MegapixelAndSquaredPixelTiers_Interleave(t *testing.T) { + // 4096x4096 = 16,777,216px sits BETWEEN the 16MP and 32MP thresholds, not + // between 8MP and 16MP — the tier ladder must be ordered by actual pixel + // count, not by field family, or a model using both families would pick + // the wrong tier. + p := configstoreTables.TableModelPricing{ + OutputCostPerImage: bifrost.Ptr(0.005), + OutputCostPerImageAbove16Megapixels: bifrost.Ptr(0.04), + OutputCostPerImageAbove4096x4096Pixels: bifrost.Ptr(0.30), + } + usage := &schemas.ImageUsage{ + OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 1}, + } + + // 16.5MP: above the 16MP threshold but below 4096x4096 (16,777,216px) → $0.04. + cost := computeImageCost(&p, usage, "1x16500000", "", serviceTier{}) + assert.InDelta(t, 0.04, cost, 1e-12) + + // 17MP: above both thresholds → the larger (4096x4096) threshold wins → $0.30. + cost = computeImageCost(&p, usage, "1x17000000", "", serviceTier{}) + assert.InDelta(t, 0.30, cost, 1e-12) +} + func TestParseImagePixels(t *testing.T) { assert.Equal(t, 1048576, parseImagePixels("1024x1024")) assert.Equal(t, 262144, parseImagePixels("512x512")) diff --git a/framework/modelcatalog/datasheet/overrides.go b/framework/modelcatalog/datasheet/overrides.go index 4abd54b0402..1175d6bf832 100644 --- a/framework/modelcatalog/datasheet/overrides.go +++ b/framework/modelcatalog/datasheet/overrides.go @@ -543,6 +543,11 @@ func patchPricing(pricing configstoreTables.TableModelPricing, override Options) {dst: &patched.OutputCostPerImageAbove1024x1024PixelsPremium, src: override.OutputCostPerImageAbove1024x1024PixelsPremium}, {dst: &patched.OutputCostPerImageAbove2048x2048Pixels, src: override.OutputCostPerImageAbove2048x2048Pixels}, {dst: &patched.OutputCostPerImageAbove4096x4096Pixels, src: override.OutputCostPerImageAbove4096x4096Pixels}, + {dst: &patched.OutputCostPerImageAbove4Megapixels, src: override.OutputCostPerImageAbove4Megapixels}, + {dst: &patched.OutputCostPerImageAbove8Megapixels, src: override.OutputCostPerImageAbove8Megapixels}, + {dst: &patched.OutputCostPerImageAbove16Megapixels, src: override.OutputCostPerImageAbove16Megapixels}, + {dst: &patched.OutputCostPerImageAbove32Megapixels, src: override.OutputCostPerImageAbove32Megapixels}, + {dst: &patched.OutputCostPerImageAbove64Megapixels, src: override.OutputCostPerImageAbove64Megapixels}, {dst: &patched.CacheReadInputImageTokenCost, src: override.CacheReadInputImageTokenCost}, {dst: &patched.SearchContextCostPerQuery, src: override.SearchContextCostPerQuery}, {dst: &patched.CodeInterpreterCostPerSession, src: override.CodeInterpreterCostPerSession}, diff --git a/framework/modelcatalog/datasheet/overrides_test.go b/framework/modelcatalog/datasheet/overrides_test.go index 571f9df4f3f..b1e88ce16c4 100644 --- a/framework/modelcatalog/datasheet/overrides_test.go +++ b/framework/modelcatalog/datasheet/overrides_test.go @@ -471,6 +471,33 @@ func TestPatchPricing_CostPerRequestZero(t *testing.T) { assert.Equal(t, 0.0, *patched.CostPerRequest) } +func TestPatchPricing_MegapixelImageTiers(t *testing.T) { + base := configstoreTables.TableModelPricing{ + Model: "prunaai/p-image-upscale", + Provider: "replicate", + Mode: "image_generation", + } + + patched := patchPricing(base, Options{ + OutputCostPerImageAbove4Megapixels: bifrost.Ptr(0.01), + OutputCostPerImageAbove8Megapixels: bifrost.Ptr(0.02), + OutputCostPerImageAbove16Megapixels: bifrost.Ptr(0.04), + OutputCostPerImageAbove32Megapixels: bifrost.Ptr(0.06), + OutputCostPerImageAbove64Megapixels: bifrost.Ptr(0.12), + }) + + require.NotNil(t, patched.OutputCostPerImageAbove4Megapixels) + require.NotNil(t, patched.OutputCostPerImageAbove8Megapixels) + require.NotNil(t, patched.OutputCostPerImageAbove16Megapixels) + require.NotNil(t, patched.OutputCostPerImageAbove32Megapixels) + require.NotNil(t, patched.OutputCostPerImageAbove64Megapixels) + assert.Equal(t, 0.01, *patched.OutputCostPerImageAbove4Megapixels) + assert.Equal(t, 0.02, *patched.OutputCostPerImageAbove8Megapixels) + assert.Equal(t, 0.04, *patched.OutputCostPerImageAbove16Megapixels) + assert.Equal(t, 0.06, *patched.OutputCostPerImageAbove32Megapixels) + assert.Equal(t, 0.12, *patched.OutputCostPerImageAbove64Megapixels) +} + func TestApplyScopedOverrides_ScopePrecedence(t *testing.T) { s := newTestStore() diff --git a/framework/modelcatalog/datasheet/types.go b/framework/modelcatalog/datasheet/types.go index 5adf291b185..02f162e0e12 100644 --- a/framework/modelcatalog/datasheet/types.go +++ b/framework/modelcatalog/datasheet/types.go @@ -164,6 +164,11 @@ type Options struct { OutputCostPerImageAbove1024x1024PixelsPremium *float64 `json:"output_cost_per_image_above_1024_and_1024_pixels_and_premium_image,omitempty"` OutputCostPerImageAbove2048x2048Pixels *float64 `json:"output_cost_per_image_above_2048_and_2048_pixels,omitempty"` OutputCostPerImageAbove4096x4096Pixels *float64 `json:"output_cost_per_image_above_4096_and_4096_pixels,omitempty"` + OutputCostPerImageAbove4Megapixels *float64 `json:"output_cost_per_image_above_4_megapixels,omitempty"` + OutputCostPerImageAbove8Megapixels *float64 `json:"output_cost_per_image_above_8_megapixels,omitempty"` + OutputCostPerImageAbove16Megapixels *float64 `json:"output_cost_per_image_above_16_megapixels,omitempty"` + OutputCostPerImageAbove32Megapixels *float64 `json:"output_cost_per_image_above_32_megapixels,omitempty"` + OutputCostPerImageAbove64Megapixels *float64 `json:"output_cost_per_image_above_64_megapixels,omitempty"` OutputCostPerImageLowQuality *float64 `json:"output_cost_per_image_low_quality,omitempty"` OutputCostPerImageMediumQuality *float64 `json:"output_cost_per_image_medium_quality,omitempty"` OutputCostPerImageHighQuality *float64 `json:"output_cost_per_image_high_quality,omitempty"` @@ -667,6 +672,11 @@ func convertEntryToTablePricing(modelKey string, entry Entry) configstoreTables. OutputCostPerImageAbove1024x1024PixelsPremium: entry.OutputCostPerImageAbove1024x1024PixelsPremium, OutputCostPerImageAbove2048x2048Pixels: entry.OutputCostPerImageAbove2048x2048Pixels, OutputCostPerImageAbove4096x4096Pixels: entry.OutputCostPerImageAbove4096x4096Pixels, + OutputCostPerImageAbove4Megapixels: entry.OutputCostPerImageAbove4Megapixels, + OutputCostPerImageAbove8Megapixels: entry.OutputCostPerImageAbove8Megapixels, + OutputCostPerImageAbove16Megapixels: entry.OutputCostPerImageAbove16Megapixels, + OutputCostPerImageAbove32Megapixels: entry.OutputCostPerImageAbove32Megapixels, + OutputCostPerImageAbove64Megapixels: entry.OutputCostPerImageAbove64Megapixels, OutputCostPerImageLowQuality: entry.OutputCostPerImageLowQuality, OutputCostPerImageMediumQuality: entry.OutputCostPerImageMediumQuality, OutputCostPerImageHighQuality: entry.OutputCostPerImageHighQuality, @@ -756,6 +766,11 @@ func convertTablePricingToEntry(pricing *configstoreTables.TableModelPricing) *E OutputCostPerImageAbove1024x1024PixelsPremium: pricing.OutputCostPerImageAbove1024x1024PixelsPremium, OutputCostPerImageAbove2048x2048Pixels: pricing.OutputCostPerImageAbove2048x2048Pixels, OutputCostPerImageAbove4096x4096Pixels: pricing.OutputCostPerImageAbove4096x4096Pixels, + OutputCostPerImageAbove4Megapixels: pricing.OutputCostPerImageAbove4Megapixels, + OutputCostPerImageAbove8Megapixels: pricing.OutputCostPerImageAbove8Megapixels, + OutputCostPerImageAbove16Megapixels: pricing.OutputCostPerImageAbove16Megapixels, + OutputCostPerImageAbove32Megapixels: pricing.OutputCostPerImageAbove32Megapixels, + OutputCostPerImageAbove64Megapixels: pricing.OutputCostPerImageAbove64Megapixels, OutputCostPerImageLowQuality: pricing.OutputCostPerImageLowQuality, OutputCostPerImageMediumQuality: pricing.OutputCostPerImageMediumQuality, OutputCostPerImageHighQuality: pricing.OutputCostPerImageHighQuality, diff --git a/ui/app/workspace/custom-pricing/overrides/pricingFields.test.ts b/ui/app/workspace/custom-pricing/overrides/pricingFields.test.ts index cda794f66d7..f5962889f02 100644 --- a/ui/app/workspace/custom-pricing/overrides/pricingFields.test.ts +++ b/ui/app/workspace/custom-pricing/overrides/pricingFields.test.ts @@ -79,7 +79,7 @@ describe("pricingFieldUnit", () => { expect(byUnit[unit], `${field.key} resolved to unexpected unit ${unit}`).toBeDefined(); byUnit[unit].push(field.key); } - expect(PRICING_FIELDS).toHaveLength(77); + expect(PRICING_FIELDS).toHaveLength(82); expect(byUnit.multiplier).toEqual(["inference_geo_us_multiplier"]); expect(byUnit.character).toEqual(["input_cost_per_character"]); // Sanity: the split is real, not everything collapsing into one bucket. diff --git a/ui/app/workspace/custom-pricing/overrides/pricingFields.ts b/ui/app/workspace/custom-pricing/overrides/pricingFields.ts index efdf6a4d93d..181e3091ba2 100644 --- a/ui/app/workspace/custom-pricing/overrides/pricingFields.ts +++ b/ui/app/workspace/custom-pricing/overrides/pricingFields.ts @@ -447,6 +447,36 @@ export const PRICING_FIELDS = [ group: "image", requestTypeGroups: ["image"], }, + { + key: "output_cost_per_image_above_4_megapixels", + label: "Output / image (>4MP)", + group: "image", + requestTypeGroups: ["image"], + }, + { + key: "output_cost_per_image_above_8_megapixels", + label: "Output / image (>8MP)", + group: "image", + requestTypeGroups: ["image"], + }, + { + key: "output_cost_per_image_above_16_megapixels", + label: "Output / image (>16MP)", + group: "image", + requestTypeGroups: ["image"], + }, + { + key: "output_cost_per_image_above_32_megapixels", + label: "Output / image (>32MP)", + group: "image", + requestTypeGroups: ["image"], + }, + { + key: "output_cost_per_image_above_64_megapixels", + label: "Output / image (>64MP)", + group: "image", + requestTypeGroups: ["image"], + }, { key: "output_cost_per_image_low_quality", label: "Output / image (low quality)", diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index a5385dffd40..5763874925b 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -545,6 +545,11 @@ export interface PricingOverridePatch { output_cost_per_image_above_1024_and_1024_pixels_and_premium_image?: number; output_cost_per_image_above_2048_and_2048_pixels?: number; output_cost_per_image_above_4096_and_4096_pixels?: number; + output_cost_per_image_above_4_megapixels?: number; + output_cost_per_image_above_8_megapixels?: number; + output_cost_per_image_above_16_megapixels?: number; + output_cost_per_image_above_32_megapixels?: number; + output_cost_per_image_above_64_megapixels?: number; output_cost_per_image_low_quality?: number; output_cost_per_image_medium_quality?: number; output_cost_per_image_high_quality?: number; From 8b38d132deb5cad4561d3b83cd9f91811caf2267 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Wed, 12 Aug 2026 04:36:59 +0530 Subject: [PATCH 2/2] feat: backfill output resolution for Replicate upscale-style image models --- core/providers/replicate/images.go | 186 ++++++++++ .../images_upscale_resolution_test.go | 334 ++++++++++++++++++ core/providers/replicate/replicate.go | 23 ++ core/providers/replicate/types.go | 7 + 4 files changed, 550 insertions(+) create mode 100644 core/providers/replicate/images_upscale_resolution_test.go diff --git a/core/providers/replicate/images.go b/core/providers/replicate/images.go index b6c4c4c4e89..8b49546ec08 100644 --- a/core/providers/replicate/images.go +++ b/core/providers/replicate/images.go @@ -2,6 +2,8 @@ package replicate import ( "fmt" + "math" + "strconv" "strings" providerUtils "github.com/maximhq/bifrost/core/providers/utils" @@ -181,6 +183,190 @@ func ToBifrostImageGenerationResponse( return response, nil } +// applyUpscaleOutputResolution backfills ImageGenerationResponseParameters.Size +// on an upscale-style response (e.g. prunaai/p-image-upscale) whose output +// resolution isn't otherwise knowable: these models take an input image plus +// a "target" (desired output megapixels) or "factor" (multiplier on the input +// image's dimensions) parameter instead of a plain size string, so neither +// the request's Params.Size nor the provider's own response carries any +// resolution info by default. Without this, resolution-tiered image pricing +// silently falls back to the base per-image rate regardless of actual output +// size. No-op (leaves Size untouched) when neither signal is present. +func applyUpscaleOutputResolution(request *schemas.BifrostImageGenerationRequest, prediction *ReplicatePredictionResponse, response *schemas.BifrostImageGenerationResponse) { + if request == nil || response == nil { + return + } + setUpscaleOutputSize(resolveUpscaleOutputPixels(request, prediction), response) +} + +// applyUpscaleEditOutputResolution is applyUpscaleOutputResolution for the +// image edit path, which returns the same response shape and so feeds the same +// resolution-tiered pricing lookup. +func applyUpscaleEditOutputResolution(request *schemas.BifrostImageEditRequest, prediction *ReplicatePredictionResponse, response *schemas.BifrostImageGenerationResponse) { + if request == nil || response == nil { + return + } + setUpscaleOutputSize(resolveUpscaleEditOutputPixels(request, prediction), response) +} + +// setUpscaleOutputSize writes a resolved pixel count onto a response as a size +// string, leaving a size the provider itself reported untouched. +func setUpscaleOutputSize(pixels int, response *schemas.BifrostImageGenerationResponse) { + if pixels <= 0 { + return + } + if response.ImageGenerationResponseParameters == nil { + response.ImageGenerationResponseParameters = &schemas.ImageGenerationResponseParameters{} + } + if response.ImageGenerationResponseParameters.Size == "" { + response.ImageGenerationResponseParameters.Size = formatSquarePixelSize(pixels) + } +} + +// applyUpscaleStreamOutputResolution is setUpscaleOutputSize for a streaming +// completion chunk, which carries Size directly rather than nested under +// response parameters. Streaming callers can only supply the request-side +// signal: the SSE path never re-reads the finished prediction, so the +// metrics.resolution_target fallback that covers factor mode is unavailable +// there and a factor-mode stream is left without a size, as before. +func applyUpscaleStreamOutputResolution(pixels int, chunk *schemas.BifrostImageGenerationStreamResponse) { + if chunk == nil || pixels <= 0 || chunk.Size != "" { + return + } + chunk.Size = formatSquarePixelSize(pixels) +} + +// resolveUpscaleOutputPixels estimates the total output pixel count for an +// upscale-style request, in priority order: +// 1. "target" mode: the request declares its desired output resolution in +// megapixels directly (e.g. target: 16) — known before the call is made. +// 2. "factor" mode: output size depends on the input image's own resolution +// (unknown to Bifrost), so we fall back to the megapixel band Replicate +// itself reports post-hoc via metrics.resolution_target (e.g. "8-16MP"). +// +// Returns 0 when neither signal is present. +func resolveUpscaleOutputPixels(request *schemas.BifrostImageGenerationRequest, prediction *ReplicatePredictionResponse) int { + if request == nil || request.Params == nil || request.Params.ExtraParams == nil { + return resolveUpscaleOutputPixelsFromMetrics(prediction) + } + if pixels := resolveUpscaleTargetPixelsFromExtraParams(request.Params.ExtraParams); pixels > 0 { + return pixels + } + + return resolveUpscaleOutputPixelsFromMetrics(prediction) +} + +// resolveUpscaleTargetPixelsFromExtraParams reads the request-side "target" +// signal out of Replicate's own native parameter names, which callers can pass +// through verbatim on either the image generation or image edit path. Returns +// 0 in factor mode, where the output size follows the input image instead and +// only the provider can report it. +func resolveUpscaleTargetPixelsFromExtraParams(extraParams map[string]interface{}) int { + if extraParams == nil { + return 0 + } + upscaleMode, _ := schemas.SafeExtractString(extraParams["upscale_mode"]) + if upscaleMode != "" && upscaleMode != "target" { + return 0 + } + if targetMP, ok := schemas.SafeExtractFloat64(extraParams["target"]); ok { + return megapixelsToPixels(targetMP) + } + return 0 +} + +// resolveUpscaleEditOutputPixels is resolveUpscaleOutputPixels for the image +// edit path, where an upscale is expressed through the first-class +// TargetMegapixels / UpscaleFactor params rather than through Replicate's raw +// input names. The two are mutually exclusive: UpscaleFactor scales the input +// image, so its output size is unknowable up front and falls back to the +// band the provider reports post-hoc, exactly as factor mode does on the +// generation path. +func resolveUpscaleEditOutputPixels(request *schemas.BifrostImageEditRequest, prediction *ReplicatePredictionResponse) int { + if request == nil || request.Params == nil { + return resolveUpscaleOutputPixelsFromMetrics(prediction) + } + params := request.Params + if params.UpscaleFactor == nil && params.TargetMegapixels != nil { + // Converted as a float so an out-of-range value is rejected by the + // shared bound check rather than wrapping through the multiply. + if pixels := megapixelsToPixels(float64(*params.TargetMegapixels)); pixels > 0 { + return pixels + } + } + if params.UpscaleFactor == nil { + if pixels := resolveUpscaleTargetPixelsFromExtraParams(params.ExtraParams); pixels > 0 { + return pixels + } + } + + return resolveUpscaleOutputPixelsFromMetrics(prediction) +} + +// resolveUpscaleOutputPixelsFromMetrics parses a megapixel band string like +// "8-16MP" or "16MP" from the prediction's metrics.resolution_target field. +// Uses the upper bound of the band as the billable pixel estimate — the +// conservative choice, since underestimating post-hoc would under-bill. +func resolveUpscaleOutputPixelsFromMetrics(prediction *ReplicatePredictionResponse) int { + if prediction == nil || prediction.Metrics == nil || prediction.Metrics.ResolutionTarget == nil { + return 0 + } + band := strings.ToUpper(strings.TrimSpace(*prediction.Metrics.ResolutionTarget)) + band = strings.TrimSuffix(band, "MP") + if band == "" { + return 0 + } + // A band is either a single value ("16MP") or an ascending pair + // ("8-16MP"). Anything else is malformed, and since this value ends up + // selecting a resolution pricing tier, a malformed band must resolve to + // no size rather than to whichever number happens to parse last. + parts := strings.Split(band, "-") + if len(parts) > 2 { + return 0 + } + mp, err := strconv.ParseFloat(strings.TrimSpace(parts[len(parts)-1]), 64) + if err != nil { + return 0 + } + if len(parts) == 2 { + lower, lowerErr := strconv.ParseFloat(strings.TrimSpace(parts[0]), 64) + // Ordering is only meaningful between two finite bounds: every + // comparison against NaN is false, so a non-finite lower bound would + // pass an ordering check it never actually satisfied. + if lowerErr != nil || math.IsNaN(lower) || math.IsInf(lower, 0) || lower <= 0 || lower > mp { + return 0 + } + } + return megapixelsToPixels(mp) +} + +// megapixelsToPixels converts a megapixel figure to a total pixel count, +// rejecting anything that cannot describe a real output size. ParseFloat +// accepts "NaN" and "Inf", and a non-finite value both slips past ordinary +// range checks and saturates the conversion to int, so it has to be rejected +// by name. Rounds up for the same reason formatSquarePixelSize does: a value +// sitting on a pricing tier's threshold must never be billed one tier down. +func megapixelsToPixels(mp float64) int { + if math.IsNaN(mp) || math.IsInf(mp, 0) || mp <= 0 { + return 0 + } + pixels := math.Ceil(mp * 1_000_000) + if pixels >= math.MaxInt64 { + return 0 + } + return int(pixels) +} + +// formatSquarePixelSize formats a total pixel count as a "WxH" size string +// for ImageGenerationResponseParameters.Size, using a square approximation +// (side = ceil(sqrt(pixels))). Rounding up guarantees width*height never +// falls below the true pixel count, so a value sitting exactly on a pricing +// tier's threshold is never miscategorized into the tier below it. +func formatSquarePixelSize(pixels int) string { + side := int(math.Ceil(math.Sqrt(float64(pixels)))) + return fmt.Sprintf("%dx%d", side, side) +} + // getInputImageFieldName returns the appropriate input image field name based on the model. // Uses O(1) map lookup for high RPS performance. func getInputImageFieldName(model string) string { diff --git a/core/providers/replicate/images_upscale_resolution_test.go b/core/providers/replicate/images_upscale_resolution_test.go new file mode 100644 index 00000000000..04e6a078193 --- /dev/null +++ b/core/providers/replicate/images_upscale_resolution_test.go @@ -0,0 +1,334 @@ +package replicate + +import ( + "fmt" + "math" + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +func newUpscaleRequest(extraParams map[string]interface{}) *schemas.BifrostImageGenerationRequest { + return &schemas.BifrostImageGenerationRequest{ + Model: "prunaai/p-image-upscale", + Params: &schemas.ImageGenerationParameters{ + ExtraParams: extraParams, + }, + } +} + +func TestResolveUpscaleOutputPixels_TargetMode(t *testing.T) { + req := newUpscaleRequest(map[string]interface{}{ + "target": float64(16), + "upscale_mode": "target", + }) + got := resolveUpscaleOutputPixels(req, nil) + want := 16_000_000 + if got != want { + t.Errorf("target=16 upscale_mode=target: want %d pixels, got %d", want, got) + } +} + +func TestResolveUpscaleOutputPixels_TargetMode_DefaultsWhenModeOmitted(t *testing.T) { + // upscale_mode defaults to "target" per Replicate's docs, so a bare + // "target" param (no explicit upscale_mode) must still resolve. + req := newUpscaleRequest(map[string]interface{}{ + "target": float64(8), + }) + got := resolveUpscaleOutputPixels(req, nil) + want := 8_000_000 + if got != want { + t.Errorf("target=8 (no upscale_mode): want %d pixels, got %d", want, got) + } +} + +func TestResolveUpscaleOutputPixels_FactorMode_IgnoresTargetFallsBackToMetrics(t *testing.T) { + // A stray "target" value must not be used once upscale_mode explicitly + // says "factor" — factor mode's output size depends on the input image, + // not the target param. + req := newUpscaleRequest(map[string]interface{}{ + // Deliberately distinct from the metrics band's 16MP upper bound, so a + // regression that honours "target" in factor mode fails this test + // instead of coincidentally producing the same pixel count. + "target": float64(8), // should be ignored + "upscale_mode": "factor", + "factor": float64(4), + }) + resolutionTarget := "8-16MP" + prediction := &ReplicatePredictionResponse{ + Metrics: &ReplicateMetrics{ResolutionTarget: &resolutionTarget}, + } + got := resolveUpscaleOutputPixels(req, prediction) + want := 16_000_000 // upper bound of the "8-16MP" band, not the ignored target=8 + if got != want { + t.Errorf("factor mode: want %d pixels (from metrics band), got %d", want, got) + } +} + +func TestResolveUpscaleOutputPixels_NoSignal(t *testing.T) { + req := newUpscaleRequest(nil) + if got := resolveUpscaleOutputPixels(req, nil); got != 0 { + t.Errorf("no target, no metrics: want 0, got %d", got) + } +} + +func TestResolveUpscaleOutputPixelsFromMetrics(t *testing.T) { + cases := []struct { + name string + band *string + want int + }{ + {"range band", strPtr("8-16MP"), 16_000_000}, + {"single value band", strPtr("16MP"), 16_000_000}, + {"lowercase mp", strPtr("4-8mp"), 8_000_000}, + {"fractional bounds", strPtr("0.5-1MP"), 1_000_000}, + {"nil metrics field", nil, 0}, + // Malformed bands must not yield a billable size: this value feeds + // resolution-tiered pricing, so parsing only the trailing component + // would let junk through as a real output resolution. + {"non-numeric lower bound", strPtr("invalid-16MP"), 0}, + {"empty middle component", strPtr("8--16MP"), 0}, + {"descending bounds", strPtr("16-8MP"), 0}, + {"three components", strPtr("4-8-16MP"), 0}, + {"zero lower bound", strPtr("0-16MP"), 0}, + // Non-finite floats parse without error, and every comparison against + // NaN is false, so they slip past ordering checks unless rejected + // outright. Infinity additionally saturates the int conversion. + {"nan lower bound", strPtr("NaN-16MP"), 0}, + {"nan upper bound", strPtr("8-NaNMP"), 0}, + {"infinite upper bound", strPtr("InfMP"), 0}, + {"infinite lower bound", strPtr("Inf-16MP"), 0}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + prediction := &ReplicatePredictionResponse{ + Metrics: &ReplicateMetrics{ResolutionTarget: tc.band}, + } + got := resolveUpscaleOutputPixelsFromMetrics(prediction) + if got != tc.want { + t.Errorf("want %d, got %d", tc.want, got) + } + }) + } + // nil prediction / nil metrics must not panic. + if got := resolveUpscaleOutputPixelsFromMetrics(nil); got != 0 { + t.Errorf("nil prediction: want 0, got %d", got) + } + if got := resolveUpscaleOutputPixelsFromMetrics(&ReplicatePredictionResponse{}); got != 0 { + t.Errorf("nil metrics: want 0, got %d", got) + } +} + +func TestFormatSquarePixelSize_NeverUndercountsAtBoundary(t *testing.T) { + // 8,000,000 has no integer square root (sqrt ≈ 2828.43); ceil-rounding + // must keep width*height >= 8,000,000 so an exact-boundary target + // (e.g. target: 8) never gets miscategorized into the tier below it. + size := formatSquarePixelSize(8_000_000) + w, h := parseWxH(t, size) + if w*h < 8_000_000 { + t.Errorf("formatSquarePixelSize(8_000_000) = %q, width*height = %d, want >= 8_000_000", size, w*h) + } + + // A perfect square should round-trip exactly. + size = formatSquarePixelSize(16_000_000) + if size != "4000x4000" { + t.Errorf("formatSquarePixelSize(16_000_000) = %q, want 4000x4000", size) + } +} + +func TestApplyUpscaleOutputResolution(t *testing.T) { + req := newUpscaleRequest(map[string]interface{}{ + "target": float64(16), + "upscale_mode": "target", + }) + resp := &schemas.BifrostImageGenerationResponse{} + applyUpscaleOutputResolution(req, nil, resp) + + if resp.ImageGenerationResponseParameters == nil || resp.ImageGenerationResponseParameters.Size == "" { + t.Fatal("expected Size to be backfilled") + } + w, h := parseWxH(t, resp.ImageGenerationResponseParameters.Size) + if w*h < 16_000_000 { + t.Errorf("backfilled size %q has %d total pixels, want >= 16_000_000", resp.ImageGenerationResponseParameters.Size, w*h) + } +} + +func TestApplyUpscaleOutputResolution_DoesNotOverwriteExistingSize(t *testing.T) { + req := newUpscaleRequest(map[string]interface{}{"target": float64(16)}) + resp := &schemas.BifrostImageGenerationResponse{ + ImageGenerationResponseParameters: &schemas.ImageGenerationResponseParameters{Size: "1234x1234"}, + } + applyUpscaleOutputResolution(req, nil, resp) + + if resp.ImageGenerationResponseParameters.Size != "1234x1234" { + t.Errorf("existing Size was overwritten: got %q", resp.ImageGenerationResponseParameters.Size) + } +} + +func TestApplyUpscaleOutputResolution_NoSignalLeavesSizeEmpty(t *testing.T) { + req := newUpscaleRequest(nil) + resp := &schemas.BifrostImageGenerationResponse{} + applyUpscaleOutputResolution(req, nil, resp) + + if resp.ImageGenerationResponseParameters != nil && resp.ImageGenerationResponseParameters.Size != "" { + t.Errorf("expected no Size to be set, got %q", resp.ImageGenerationResponseParameters.Size) + } +} + +func strPtr(s string) *string { return &s } + +func parseWxH(t *testing.T, size string) (int, int) { + t.Helper() + var w, h int + n, err := fmt.Sscanf(size, "%dx%d", &w, &h) + if err != nil || n != 2 { + t.Fatalf("could not parse size %q: %v", size, err) + } + return w, h +} + +func newUpscaleEditRequest(params *schemas.ImageEditParameters) *schemas.BifrostImageEditRequest { + return &schemas.BifrostImageEditRequest{ + Model: "prunaai/p-image-upscale", + Input: &schemas.ImageEditInput{Prompt: "upscale"}, + Params: params, + } +} + +func TestResolveUpscaleEditOutputPixels_TargetMegapixels(t *testing.T) { + // The edit path has a first-class target_megapixels param, so an upscale + // routed through /images/edits must resolve without any extra params. + req := newUpscaleEditRequest(&schemas.ImageEditParameters{ + TargetMegapixels: schemas.Ptr(16), + }) + got := resolveUpscaleEditOutputPixels(req, nil) + want := 16_000_000 + if got != want { + t.Errorf("target_megapixels=16: want %d pixels, got %d", want, got) + } +} + +func TestResolveUpscaleEditOutputPixels_UpscaleFactorIgnoresTargetFallsBackToMetrics(t *testing.T) { + // upscale_factor is mutually exclusive with target_megapixels and makes the + // output size depend on the input image, so a stray target must be ignored + // in favour of the band Replicate reports back. + req := newUpscaleEditRequest(&schemas.ImageEditParameters{ + UpscaleFactor: schemas.Ptr(4), + TargetMegapixels: schemas.Ptr(8), // should be ignored + }) + resolutionTarget := "8-16MP" + prediction := &ReplicatePredictionResponse{ + Metrics: &ReplicateMetrics{ResolutionTarget: &resolutionTarget}, + } + got := resolveUpscaleEditOutputPixels(req, prediction) + want := 16_000_000 + if got != want { + t.Errorf("upscale_factor mode: want %d pixels (from metrics band), got %d", want, got) + } +} + +func TestResolveUpscaleEditOutputPixels_ExtraParamsTarget(t *testing.T) { + // Callers passing Replicate's native param names straight through must + // resolve on the edit path too, not just on image generation. + req := newUpscaleEditRequest(&schemas.ImageEditParameters{ + ExtraParams: map[string]interface{}{"target": float64(4)}, + }) + got := resolveUpscaleEditOutputPixels(req, nil) + want := 4_000_000 + if got != want { + t.Errorf("extra param target=4: want %d pixels, got %d", want, got) + } +} + +func TestResolveUpscaleEditOutputPixels_NoSignal(t *testing.T) { + if got := resolveUpscaleEditOutputPixels(newUpscaleEditRequest(nil), nil); got != 0 { + t.Errorf("no signal: want 0, got %d", got) + } +} + +func TestApplyUpscaleEditOutputResolution_SetsSize(t *testing.T) { + req := newUpscaleEditRequest(&schemas.ImageEditParameters{TargetMegapixels: schemas.Ptr(16)}) + response := &schemas.BifrostImageGenerationResponse{} + applyUpscaleEditOutputResolution(req, nil, response) + if response.ImageGenerationResponseParameters == nil { + t.Fatal("want ImageGenerationResponseParameters populated, got nil") + } + if got := response.ImageGenerationResponseParameters.Size; got != "4000x4000" { + t.Errorf("want size 4000x4000 (sqrt of 16MP), got %q", got) + } +} + +func TestApplyUpscaleEditOutputResolution_PreservesExistingSize(t *testing.T) { + req := newUpscaleEditRequest(&schemas.ImageEditParameters{TargetMegapixels: schemas.Ptr(16)}) + response := &schemas.BifrostImageGenerationResponse{ + ImageGenerationResponseParameters: &schemas.ImageGenerationResponseParameters{Size: "1024x1024"}, + } + applyUpscaleEditOutputResolution(req, nil, response) + if got := response.ImageGenerationResponseParameters.Size; got != "1024x1024" { + t.Errorf("want provider-reported size preserved, got %q", got) + } +} + +func TestApplyUpscaleStreamOutputResolution(t *testing.T) { + t.Run("sets size on completion chunk", func(t *testing.T) { + chunk := &schemas.BifrostImageGenerationStreamResponse{} + applyUpscaleStreamOutputResolution(16_000_000, chunk) + if chunk.Size != "4000x4000" { + t.Errorf("want size 4000x4000, got %q", chunk.Size) + } + }) + t.Run("no-op without a resolved pixel count", func(t *testing.T) { + chunk := &schemas.BifrostImageGenerationStreamResponse{} + applyUpscaleStreamOutputResolution(0, chunk) + if chunk.Size != "" { + t.Errorf("want size left empty, got %q", chunk.Size) + } + }) + t.Run("preserves a size already set", func(t *testing.T) { + chunk := &schemas.BifrostImageGenerationStreamResponse{Size: "1024x1024"} + applyUpscaleStreamOutputResolution(16_000_000, chunk) + if chunk.Size != "1024x1024" { + t.Errorf("want existing size preserved, got %q", chunk.Size) + } + }) +} + + +func TestResolveUpscaleOutputPixels_RejectsNonFiniteTarget(t *testing.T) { + // A non-finite target must not saturate the conversion into a + // billion-pixel output that then selects the top pricing tier. + for name, target := range map[string]float64{ + "positive infinity": math.Inf(1), + "negative infinity": math.Inf(-1), + "nan": math.NaN(), + } { + t.Run(name, func(t *testing.T) { + req := newUpscaleRequest(map[string]interface{}{"target": target}) + if got := resolveUpscaleOutputPixels(req, nil); got != 0 { + t.Errorf("want 0 pixels, got %d", got) + } + }) + } +} + +func TestResolveUpscaleOutputPixels_RoundsFractionalTargetUp(t *testing.T) { + // formatSquarePixelSize rounds up so a value on a tier boundary is never + // billed one tier down; truncating here would defeat that on the way in. + req := newUpscaleRequest(map[string]interface{}{"target": 16.0000001}) + got := resolveUpscaleOutputPixels(req, nil) + want := 16_000_001 + if got != want { + t.Errorf("target=16.0000001: want %d pixels, got %d", want, got) + } +} + +func TestResolveUpscaleEditOutputPixels_RejectsOverflowingTargetMegapixels(t *testing.T) { + // int64 multiplication wraps negative well before this, so the bound has + // to be checked before the multiply rather than after it. + req := newUpscaleEditRequest(&schemas.ImageEditParameters{ + TargetMegapixels: schemas.Ptr(9_300_000_000_000), + }) + if got := resolveUpscaleEditOutputPixels(req, nil); got != 0 { + t.Errorf("want 0 pixels for an out-of-range target, got %d", got) + } +} diff --git a/core/providers/replicate/replicate.go b/core/providers/replicate/replicate.go index 624b2981e04..52aff743b17 100644 --- a/core/providers/replicate/replicate.go +++ b/core/providers/replicate/replicate.go @@ -1833,6 +1833,11 @@ func (provider *ReplicateProvider) ImageGeneration(ctx *schemas.BifrostContext, return nil, providerUtils.EnrichError(ctx, err, jsonData, nil, provider.sendBackRawRequest, provider.sendBackRawResponse, latency) } + // Backfill output resolution for upscale-style models (target/factor + // input, no plain size param) so resolution-tiered cost calculation + // doesn't silently fall back to the base per-image rate. + applyUpscaleOutputResolution(request, prediction, bifrostResponse) + // Set extra fields bifrostResponse.ExtraFields.Latency = latency.Milliseconds() bifrostResponse.ExtraFields.ProviderResponseHeaders = providerResponseHeaders @@ -2109,6 +2114,12 @@ func (provider *ReplicateProvider) ImageGenerationStream(ctx *schemas.BifrostCon }, } + // Backfill output resolution for upscale-style models. Only the + // request-side target signal is available here: the SSE path never + // re-reads the finished prediction, so factor mode has no metrics + // band to fall back to. + applyUpscaleStreamOutputResolution(resolveUpscaleOutputPixels(request, nil), finalChunk) + // Set raw request only on final chunk if enabled if sendBackRawRequest { providerUtils.ParseAndSetRawRequest(&finalChunk.ExtraFields, jsonData) @@ -2249,6 +2260,12 @@ func (provider *ReplicateProvider) ImageEdit(ctx *schemas.BifrostContext, key sc return nil, providerUtils.EnrichError(ctx, err, jsonData, nil, provider.sendBackRawRequest, provider.sendBackRawResponse, latency) } + // Backfill output resolution for upscale-style models, which reach this path + // through the first-class target_megapixels/upscale_factor edit params, so + // resolution-tiered cost calculation doesn't silently fall back to the base + // per-image rate. + applyUpscaleEditOutputResolution(request, prediction, bifrostResponse) + // Set extra fields bifrostResponse.ExtraFields.Latency = latency.Milliseconds() bifrostResponse.ExtraFields.ProviderResponseHeaders = providerResponseHeaders @@ -2518,6 +2535,12 @@ func (provider *ReplicateProvider) ImageEditStream(ctx *schemas.BifrostContext, }, } + // Backfill output resolution for upscale-style models. Only the + // request-side target signal is available here: the SSE path never + // re-reads the finished prediction, so factor mode has no metrics + // band to fall back to. + applyUpscaleStreamOutputResolution(resolveUpscaleEditOutputPixels(request, nil), finalChunk) + if sendBackRawRequest { providerUtils.ParseAndSetRawRequest(&finalChunk.ExtraFields, jsonData) } diff --git a/core/providers/replicate/types.go b/core/providers/replicate/types.go index 3ae88c0095c..13d91a91ecd 100644 --- a/core/providers/replicate/types.go +++ b/core/providers/replicate/types.go @@ -302,6 +302,13 @@ type ReplicateMetrics struct { ImageCount *int `json:"image_count,omitempty"` // Number of images generated TimeToFirstToken *float64 `json:"time_to_first_token,omitempty"` // Time to first token (seconds) TokensPerSecond *float64 `json:"tokens_per_second,omitempty"` // Tokens generated per second + // ResolutionTarget is reported by some image models (e.g. upscalers run in + // "factor" mode, where output size depends on the input image and isn't + // known from the request alone) as a megapixel band string like "8-16MP", + // describing the actual output resolution achieved. Used for + // resolution-tiered cost calculation when the request doesn't carry an + // explicit target resolution. + ResolutionTarget *string `json:"resolution_target,omitempty"` } // ReplicatePredictionListResponse represents a paginated list of predictions