Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions convex/_generated/api.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import type * as follows from "../follows.js";
import type * as generatedImages from "../generatedImages.js";
import type * as http from "../http.js";
import type * as lib_crypto from "../lib/crypto.js";
import type * as lib_dirtberryCrop from "../lib/dirtberryCrop.js";
import type * as lib_groq from "../lib/groq.js";
import type * as lib_index from "../lib/index.js";
import type * as lib_nsfwDetection from "../lib/nsfwDetection.js";
Expand Down Expand Up @@ -73,6 +74,7 @@ declare const fullApi: ApiFromModules<{
generatedImages: typeof generatedImages;
http: typeof http;
"lib/crypto": typeof lib_crypto;
"lib/dirtberryCrop": typeof lib_dirtberryCrop;
"lib/groq": typeof lib_groq;
"lib/index": typeof lib_index;
"lib/nsfwDetection": typeof lib_nsfwDetection;
Expand Down
60 changes: 51 additions & 9 deletions convex/batchProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ import { internalAction } from "./_generated/server"
import {
buildPollinationsUrl,
classifyApiError,
cropDirtberryImageBuffer,
getDirtberrySourceDimensions,
isDirtberryModel,
generateR2Key,
generateThumbnailKey,
uploadMediaWithThumbnail,
Expand Down Expand Up @@ -130,6 +133,8 @@ export const processBatchItem = internalAction({
}

try {
const shouldCropDirtberry = isDirtberryModel(batchJob.generationParams.model)
const dirtberrySourceDimensions = shouldCropDirtberry ? getDirtberrySourceDimensions() : null
// Pollinations API only accepts seeds up to int32 max (2147483647)
const INT32_MAX = 2147483647
const rawSeed = batchJob.generationParams.seed ?? Math.floor(Math.random() * INT32_MAX)
Expand All @@ -140,8 +145,8 @@ export const processBatchItem = internalAction({
prompt: batchJob.generationParams.prompt,
negativePrompt: batchJob.generationParams.negativePrompt,
model: batchJob.generationParams.model,
width: batchJob.generationParams.width,
height: batchJob.generationParams.height,
width: dirtberrySourceDimensions?.width ?? batchJob.generationParams.width,
height: dirtberrySourceDimensions?.height ?? batchJob.generationParams.height,
seed,
enhance: batchJob.generationParams.enhance,
private: batchJob.generationParams.private,
Expand Down Expand Up @@ -191,12 +196,42 @@ export const processBatchItem = internalAction({
const imageBuffer = Buffer.from(await response.arrayBuffer())
const contentType = response.headers.get("content-type") || "image/jpeg"

let uploadBuffer = imageBuffer
let outputWidth = batchJob.generationParams.width ?? 1024
let outputHeight = batchJob.generationParams.height ?? 1024

// Crop Dirtberry outputs before upload/persistence so every surface
// (canvas, gallery, downloads, lightbox) uses the same native asset.
if (shouldCropDirtberry) {
try {
const cropped = await cropDirtberryImageBuffer(imageBuffer)
uploadBuffer = Buffer.from(cropped.buffer)
outputWidth = cropped.width
outputHeight = cropped.height

if (cropped.wasCropped) {
console.log(
`${logger} Applied Dirtberry crop (${cropped.processor}): ${cropped.inputWidth}x${cropped.inputHeight} -> ${cropped.width}x${cropped.height} (trim=${cropped.trimPixels}px)`
)
} else {
console.log(
`${logger} Dirtberry crop skipped: source=${cropped.inputWidth}x${cropped.inputHeight} (image too small to trim safely)`
)
}
} catch (cropError) {
console.error(
`${logger} Dirtberry crop failed, falling back to original image:`,
cropError
)
}
}

// Upload to R2 (and thumbnail for images — videos defer secondary assets)
const r2Key = generateR2Key(batchJob.ownerId, contentType)
console.log(`${logger} Uploading to R2: ${r2Key}`)

const { media: uploadResult, thumbnail: thumbnailResult } = await uploadMediaWithThumbnail(
imageBuffer,
uploadBuffer,
r2Key,
contentType
)
Expand All @@ -217,16 +252,23 @@ export const processBatchItem = internalAction({
previewR2Key: undefined,
previewUrl: undefined,
prompt: batchJob.generationParams.prompt,
width: batchJob.generationParams.width ?? 1024,
height: batchJob.generationParams.height ?? 1024,
width: outputWidth,
height: outputHeight,
model: batchJob.generationParams.model ?? "flux",
seed,
contentType,
sizeBytes: uploadResult.sizeBytes,
generationParams: {
...batchJob.generationParams,
seed,
},
generationParams: shouldCropDirtberry
? {
...batchJob.generationParams,
seed,
width: outputWidth,
height: outputHeight,
}
: {
...batchJob.generationParams,
seed,
},
visibility: batchJob.generationParams.private ? "unlisted" : "public",
})

Expand Down
130 changes: 130 additions & 0 deletions convex/lib/dirtberryCrop.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
"use node"

export const DIRTBERRY_TRIM_FRACTION = 0.03
export const DIRTBERRY_SOURCE_WIDTH = 832
export const DIRTBERRY_SOURCE_HEIGHT = 1216
export const DIRTBERRY_OUTPUT_HEIGHT =
DIRTBERRY_SOURCE_HEIGHT - Math.round(DIRTBERRY_SOURCE_HEIGHT * DIRTBERRY_TRIM_FRACTION) * 2

const JIMP_ENCODABLE_MIME_TYPES = new Set([
"image/bmp",
"image/x-ms-bmp",
"image/gif",
"image/jpeg",
"image/png",
"image/tiff",
] as const)

type JimpEncodableMime = "image/bmp" | "image/x-ms-bmp" | "image/gif" | "image/jpeg" | "image/png" | "image/tiff"

type DirtberryCropRegion = {
top: number
height: number
}

export type DirtberryCropResult = {
buffer: Buffer
width: number
height: number
wasCropped: boolean
inputWidth: number
inputHeight: number
trimPixels: number
processor: "jimp" | "none"
}

export function getDirtberrySourceDimensions(): { width: number; height: number } {
return {
width: DIRTBERRY_SOURCE_WIDTH,
height: DIRTBERRY_SOURCE_HEIGHT,
}
}

export function isDirtberryModel(model?: string): boolean {
const normalized = model?.trim().toLowerCase()
if (!normalized) {
return false
}
return normalized === "dirtberry" || normalized.includes("dirtberry")
}

/**
* Calculate Dirtberry's watermark-removal crop region.
* Trims 3% from top and bottom when the resulting crop is valid.
*/
export function calculateDirtberryCropRegion(height: number): DirtberryCropRegion | null {
if (!Number.isFinite(height) || height <= 0) return null

const trimPixels = Math.round(height * DIRTBERRY_TRIM_FRACTION)
const croppedHeight = height - trimPixels * 2

if (trimPixels <= 0 || croppedHeight <= 0) {
return null
}

return {
top: trimPixels,
height: croppedHeight,
}
}

/**
* Crop Dirtberry images to remove corner watermarks before persistence.
* Returns original data if crop is not applicable.
*/
export async function cropDirtberryImageBuffer(
imageBuffer: Buffer
): Promise<DirtberryCropResult> {
// Use jimp only (pure JS) so Convex linux-arm64 deployments don't fail
// on native sharp module loading.
const { Jimp } = await import("jimp")
const image = await Jimp.read(imageBuffer)
const width = image.bitmap.width
const height = image.bitmap.height

if (!width || !height) {
throw new Error("Unable to read Dirtberry image dimensions")
}

const region = calculateDirtberryCropRegion(height)
if (!region) {
return {
buffer: imageBuffer,
width,
height,
wasCropped: false,
inputWidth: width,
inputHeight: height,
trimPixels: 0,
processor: "none",
}
}

image.crop({
x: 0,
y: region.top,
w: width,
h: region.height,
})

const inputMime = image.mime
const outputMime: JimpEncodableMime =
inputMime && JIMP_ENCODABLE_MIME_TYPES.has(inputMime as JimpEncodableMime)
? (inputMime as JimpEncodableMime)
: "image/jpeg"
Comment on lines +112 to +114
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve output MIME when Dirtberry is re-encoded

This fallback forces outputMime to image/jpeg for any input MIME Jimp cannot encode, but the processors still upload with the original response content-type header and derive the file extension from it. If Dirtberry (or a future variant) returns something like image/webp, the code will upload JPEG bytes labeled as WebP, leading to incorrect metadata/content-type and potential decode issues for clients.

Useful? React with 👍 / 👎.

const croppedBuffer =
outputMime === "image/jpeg"
? Buffer.from(await image.getBuffer(outputMime, { quality: 100 }))
: Buffer.from(await image.getBuffer(outputMime))

return {
buffer: croppedBuffer,
width,
height: region.height,
wasCropped: true,
inputWidth: width,
inputHeight: height,
trimPixels: region.top,
processor: "jimp",
}
}
13 changes: 13 additions & 0 deletions convex/lib/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,19 @@ export {
type ErrorClassification,
} from "./pollinations"

// Dirtberry processing utilities
export {
DIRTBERRY_TRIM_FRACTION,
DIRTBERRY_SOURCE_WIDTH,
DIRTBERRY_SOURCE_HEIGHT,
DIRTBERRY_OUTPUT_HEIGHT,
getDirtberrySourceDimensions,
isDirtberryModel,
calculateDirtberryCropRegion,
cropDirtberryImageBuffer,
type DirtberryCropResult,
} from "./dirtberryCrop"

// R2 storage utilities
export {
generateR2Key,
Expand Down
60 changes: 51 additions & 9 deletions convex/singleGenerationProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ import { internalAction } from "./_generated/server"
import {
buildPollinationsUrl,
classifyApiError,
cropDirtberryImageBuffer,
getDirtberrySourceDimensions,
isDirtberryModel,
generateR2Key,
generateThumbnailKey,
uploadMediaWithThumbnail,
Expand Down Expand Up @@ -119,6 +122,8 @@ export const processGeneration = internalAction({

try {
const params = generation.generationParams
const shouldCropDirtberry = isDirtberryModel(params.model)
const dirtberrySourceDimensions = shouldCropDirtberry ? getDirtberrySourceDimensions() : null
// Pollinations API only accepts seeds up to int32 max (2147483647)
const INT32_MAX = 2147483647
const rawSeed = params.seed ?? Math.floor(Math.random() * INT32_MAX)
Expand All @@ -134,8 +139,8 @@ export const processGeneration = internalAction({
prompt: params.prompt,
negativePrompt: params.negativePrompt,
model: params.model,
width: params.width,
height: params.height,
width: dirtberrySourceDimensions?.width ?? params.width,
height: dirtberrySourceDimensions?.height ?? params.height,
seed,
enhance: params.enhance,
private: params.private,
Expand Down Expand Up @@ -198,12 +203,42 @@ export const processGeneration = internalAction({
const imageBuffer = Buffer.from(await response.arrayBuffer())
const contentType = response.headers.get("content-type") || "image/jpeg"

let uploadBuffer = imageBuffer
let outputWidth = params.width ?? 1024
let outputHeight = params.height ?? 1024

// Crop Dirtberry outputs before upload/persistence so every surface
// (canvas, gallery, downloads, lightbox) uses the same native asset.
if (shouldCropDirtberry) {
try {
const cropped = await cropDirtberryImageBuffer(imageBuffer)
uploadBuffer = Buffer.from(cropped.buffer)
outputWidth = cropped.width
outputHeight = cropped.height

if (cropped.wasCropped) {
console.log(
`${logger} Applied Dirtberry crop (${cropped.processor}): ${cropped.inputWidth}x${cropped.inputHeight} -> ${cropped.width}x${cropped.height} (trim=${cropped.trimPixels}px)`
)
} else {
console.log(
`${logger} Dirtberry crop skipped: source=${cropped.inputWidth}x${cropped.inputHeight} (image too small to trim safely)`
)
}
} catch (cropError) {
console.error(
`${logger} Dirtberry crop failed, falling back to original image:`,
cropError
)
Comment on lines +228 to +232
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Sync persisted dimensions when Dirtberry crop falls back

When cropDirtberryImageBuffer throws, this branch logs and keeps going with the original buffer, but outputWidth/outputHeight stay at the pre-crop request values (params.width/height, typically 832x1144 for Dirtberry). Because the upstream request is forced to source dimensions (832x1216), a crop failure stores metadata that no longer matches the uploaded asset, which can skew aspect-ratio-dependent rendering and downloaded dimension expectations; the same pattern is present in convex/batchProcessor.ts.

Useful? React with 👍 / 👎.

}
}

// Upload to R2 (and thumbnail for images — videos defer secondary assets)
const r2Key = generateR2Key(generation.ownerId, contentType)
console.log(`${logger} Uploading to R2: ${r2Key}`)

const { media: uploadResult, thumbnail: thumbnailResult } = await uploadMediaWithThumbnail(
imageBuffer,
uploadBuffer,
r2Key,
contentType
)
Expand Down Expand Up @@ -240,16 +275,23 @@ export const processGeneration = internalAction({
previewR2Key: undefined,
previewUrl: undefined,
prompt: params.prompt,
width: params.width ?? 1024,
height: params.height ?? 1024,
width: outputWidth,
height: outputHeight,
model: params.model ?? "flux",
seed,
contentType,
sizeBytes: uploadResult.sizeBytes,
generationParams: {
...params,
seed,
},
generationParams: shouldCropDirtberry
? {
...params,
seed,
width: outputWidth,
height: outputHeight,
}
: {
...params,
seed,
},
visibility: params.private ? "unlisted" : "public",
})

Expand Down
7 changes: 6 additions & 1 deletion hooks/queries/use-generate-image.ts
Original file line number Diff line number Diff line change
Expand Up @@ -361,11 +361,16 @@ export function useGenerateImage(
callbacksRef.current.onSettled?.(undefined, err, entry.params)
entry.reject?.(err)
} else if (generatedImage) {
const entryParams = entry.params as GeneratedImage["params"]
const image: GeneratedImage = {
id: generatedImage._id,
url: generatedImage.url,
prompt: generatedImage.prompt,
params: entry.params as GeneratedImage["params"],
params: {
...entryParams,
width: generatedImage.width ?? entryParams.width,
height: generatedImage.height ?? entryParams.height,
},
timestamp: generatedImage.createdAt,
r2Key: generatedImage.r2Key,
sizeBytes: generatedImage.sizeBytes,
Expand Down
Loading
Loading