mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-08-09 06:58:41 -04:00
Merge verified Odysseus fixes
This commit is contained in:
@@ -0,0 +1,30 @@
|
||||
// swift-tools-version: 6.2
|
||||
import PackageDescription
|
||||
|
||||
let package = Package(
|
||||
name: "odysseus-mlx-image-bridge",
|
||||
platforms: [.macOS(.v26)],
|
||||
products: [
|
||||
.executable(name: "odysseus-mlx-inpaint", targets: ["OdysseusMLXInpaint"]),
|
||||
.executable(name: "odysseus-mlx-colorize", targets: ["OdysseusMLXColorize"]),
|
||||
],
|
||||
dependencies: [
|
||||
.package(url: "https://github.com/xocialize/mlx-lama-swift", branch: "main"),
|
||||
.package(url: "https://github.com/xocialize/mlx-ddcolor-swift", branch: "main"),
|
||||
],
|
||||
targets: [
|
||||
.executableTarget(
|
||||
name: "OdysseusMLXInpaint",
|
||||
dependencies: [
|
||||
.product(name: "LaMa", package: "mlx-lama-swift"),
|
||||
.product(name: "MIGAN", package: "mlx-lama-swift"),
|
||||
]
|
||||
),
|
||||
.executableTarget(
|
||||
name: "OdysseusMLXColorize",
|
||||
dependencies: [
|
||||
.product(name: "DDColor", package: "mlx-ddcolor-swift"),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,80 @@
|
||||
import Foundation
|
||||
import CoreGraphics
|
||||
import ImageIO
|
||||
import UniformTypeIdentifiers
|
||||
import MLX
|
||||
import DDColor
|
||||
|
||||
struct Args {
|
||||
var model = ""
|
||||
var image = ""
|
||||
var output = ""
|
||||
var tier = ""
|
||||
}
|
||||
|
||||
func value(after flag: String, in args: [String]) -> String? {
|
||||
guard let i = args.firstIndex(of: flag), i + 1 < args.count else { return nil }
|
||||
return args[i + 1]
|
||||
}
|
||||
|
||||
func parseArgs() throws -> Args {
|
||||
let argv = Array(CommandLine.arguments.dropFirst())
|
||||
var out = Args()
|
||||
out.model = value(after: "--model", in: argv) ?? ""
|
||||
out.image = value(after: "--image", in: argv) ?? ""
|
||||
out.output = value(after: "--output", in: argv) ?? ""
|
||||
out.tier = value(after: "--tier", in: argv) ?? ""
|
||||
guard !out.model.isEmpty, !out.image.isEmpty, !out.output.isEmpty else {
|
||||
throw BridgeError.usage("usage: odysseus-mlx-colorize --model weights.safetensors --image input.png --output output.png [--tier tiny|large]")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeCGImage(_ path: String) throws -> CGImage {
|
||||
let url = URL(fileURLWithPath: path)
|
||||
guard let src = CGImageSourceCreateWithURL(url as CFURL, nil),
|
||||
let cg = CGImageSourceCreateImageAtIndex(src, 0, nil) else {
|
||||
throw BridgeError.decode(path)
|
||||
}
|
||||
return cg
|
||||
}
|
||||
|
||||
func encodePNG(_ image: CGImage, _ path: String) throws {
|
||||
let url = URL(fileURLWithPath: path)
|
||||
guard let dest = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil) else {
|
||||
throw BridgeError.encode(path)
|
||||
}
|
||||
CGImageDestinationAddImage(dest, image, nil)
|
||||
guard CGImageDestinationFinalize(dest) else { throw BridgeError.encode(path) }
|
||||
}
|
||||
|
||||
enum BridgeError: Error, CustomStringConvertible {
|
||||
case usage(String)
|
||||
case decode(String)
|
||||
case encode(String)
|
||||
|
||||
var description: String {
|
||||
switch self {
|
||||
case .usage(let s): return s
|
||||
case .decode(let p): return "failed to decode image: \(p)"
|
||||
case .encode(let p): return "failed to write PNG: \(p)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
do {
|
||||
let args = try parseArgs()
|
||||
let image = try decodeCGImage(args.image)
|
||||
let text = (args.tier + " " + args.model).lowercased()
|
||||
let tier: DDColorTier = text.contains("tiny") ? .tiny : .large
|
||||
let colorizer = try DDColorColorizer.fromPretrained(
|
||||
args.model,
|
||||
config: DDColorConfig(tier: tier),
|
||||
dtype: .float16
|
||||
)
|
||||
let output = colorizer(image)
|
||||
try encodePNG(output, args.output)
|
||||
} catch {
|
||||
fputs("\(error)\n", stderr)
|
||||
exit(1)
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
import Foundation
|
||||
import CoreGraphics
|
||||
import ImageIO
|
||||
import UniformTypeIdentifiers
|
||||
import MLX
|
||||
import LaMa
|
||||
import MIGAN
|
||||
|
||||
struct Args {
|
||||
var model = ""
|
||||
var image = ""
|
||||
var mask = ""
|
||||
var output = ""
|
||||
var mode = ""
|
||||
}
|
||||
|
||||
func value(after flag: String, in args: [String]) -> String? {
|
||||
guard let i = args.firstIndex(of: flag), i + 1 < args.count else { return nil }
|
||||
return args[i + 1]
|
||||
}
|
||||
|
||||
func parseArgs() throws -> Args {
|
||||
let argv = Array(CommandLine.arguments.dropFirst())
|
||||
var out = Args()
|
||||
out.model = value(after: "--model", in: argv) ?? ""
|
||||
out.image = value(after: "--image", in: argv) ?? ""
|
||||
out.mask = value(after: "--mask", in: argv) ?? ""
|
||||
out.output = value(after: "--output", in: argv) ?? ""
|
||||
out.mode = value(after: "--mode", in: argv) ?? ""
|
||||
guard !out.model.isEmpty, !out.image.isEmpty, !out.mask.isEmpty, !out.output.isEmpty else {
|
||||
throw BridgeError.usage("usage: odysseus-mlx-inpaint --model weights.safetensors --image input.png --mask mask.png --output output.png [--mode best|fast]")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeCGImage(_ path: String) throws -> CGImage {
|
||||
let url = URL(fileURLWithPath: path)
|
||||
guard let src = CGImageSourceCreateWithURL(url as CFURL, nil),
|
||||
let cg = CGImageSourceCreateImageAtIndex(src, 0, nil) else {
|
||||
throw BridgeError.decode(path)
|
||||
}
|
||||
return cg
|
||||
}
|
||||
|
||||
func encodePNG(_ image: CGImage, _ path: String) throws {
|
||||
let url = URL(fileURLWithPath: path)
|
||||
guard let dest = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil) else {
|
||||
throw BridgeError.encode(path)
|
||||
}
|
||||
CGImageDestinationAddImage(dest, image, nil)
|
||||
guard CGImageDestinationFinalize(dest) else { throw BridgeError.encode(path) }
|
||||
}
|
||||
|
||||
enum BridgeError: Error, CustomStringConvertible {
|
||||
case usage(String)
|
||||
case decode(String)
|
||||
case encode(String)
|
||||
|
||||
var description: String {
|
||||
switch self {
|
||||
case .usage(let s): return s
|
||||
case .decode(let p): return "failed to decode image: \(p)"
|
||||
case .encode(let p): return "failed to write PNG: \(p)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
do {
|
||||
let args = try parseArgs()
|
||||
let source = try decodeCGImage(args.image)
|
||||
let mask = try decodeCGImage(args.mask)
|
||||
let lower = args.model.lowercased()
|
||||
let mode = args.mode.lowercased()
|
||||
let output: CGImage
|
||||
|
||||
if lower.contains("mi-gan") || lower.contains("migan") || mode == "fast" {
|
||||
let resolution = lower.contains("512") ? 512 : 256
|
||||
let inpainter = try MIGANInpainter.fromPretrained(args.model, resolution: resolution, dtype: .float16)
|
||||
output = inpainter(source, mask: mask)
|
||||
} else {
|
||||
let inpainter = try LaMaInpainter.fromPretrained(args.model, dtype: .bfloat16)
|
||||
output = inpainter(source, mask: mask)
|
||||
}
|
||||
try encodePNG(output, args.output)
|
||||
} catch {
|
||||
fputs("\(error)\n", stderr)
|
||||
exit(1)
|
||||
}
|
||||
Reference in New Issue
Block a user