-
Notifications
You must be signed in to change notification settings - Fork 0
feat: add dirtberry #43
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
0a2c225
27fbd78
db11a87
12057dd
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| 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" | ||
| 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", | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,6 +19,9 @@ import { internalAction } from "./_generated/server" | |
| import { | ||
| buildPollinationsUrl, | ||
| classifyApiError, | ||
| cropDirtberryImageBuffer, | ||
| getDirtberrySourceDimensions, | ||
| isDirtberryModel, | ||
| generateR2Key, | ||
| generateThumbnailKey, | ||
| uploadMediaWithThumbnail, | ||
|
|
@@ -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) | ||
|
|
@@ -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, | ||
|
|
@@ -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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When 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 | ||
| ) | ||
|
|
@@ -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", | ||
| }) | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This fallback forces
outputMimetoimage/jpegfor any input MIME Jimp cannot encode, but the processors still upload with the original responsecontent-typeheader and derive the file extension from it. If Dirtberry (or a future variant) returns something likeimage/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 👍 / 👎.