Initial release
This commit is contained in:
civ
@@ -0,0 +1,498 @@
|
||||
package main
|
||||
|
||||
// engine.go — in-process FFI to libtrellis2.so via purego (no cgo), following
|
||||
// the depth-anything.cpp server pattern. The C ABI is trellis2_capi.h; the
|
||||
// t2_abi_version binding guards against header/library drift.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
)
|
||||
|
||||
const abiVersion = 11
|
||||
|
||||
// Progress stages (enum t2_stage).
|
||||
const (
|
||||
stagePreprocess = 0
|
||||
stageDino = 1
|
||||
stageSSFlow = 2
|
||||
stageSSDec = 3
|
||||
stageSLATFlow = 4
|
||||
stageShapeDec = 5
|
||||
stageMesh = 6
|
||||
stageUpsample = 7
|
||||
stageSLATFlowHR = 8
|
||||
stageShapeDecHR = 9
|
||||
stageTexture = 10
|
||||
)
|
||||
|
||||
var stageNames = map[int]string{
|
||||
stagePreprocess: "preprocess",
|
||||
stageDino: "encoding image (DINOv3)",
|
||||
stageSSFlow: "sampling sparse structure",
|
||||
stageSSDec: "decoding occupancy",
|
||||
stageSLATFlow: "sampling shape SLAT",
|
||||
stageShapeDec: "decoding shape",
|
||||
stageMesh: "extracting mesh",
|
||||
stageUpsample: "upsampling scaffold",
|
||||
stageSLATFlowHR: "sampling shape SLAT (1024)",
|
||||
stageShapeDecHR: "decoding shape (1024)",
|
||||
stageTexture: "sampling material",
|
||||
}
|
||||
|
||||
// Pipeline types (enum t2_pipeline_type) and capability bits (enum t2_caps).
|
||||
const (
|
||||
pipeAuto = 0
|
||||
pipeCoarse = 1
|
||||
pipe512 = 2
|
||||
pipe1024 = 3
|
||||
|
||||
capCoarse = 1
|
||||
cap512 = 2
|
||||
cap1024 = 4
|
||||
capTexture = 8
|
||||
|
||||
backgroundAuto = 0
|
||||
backgroundKeep = 1
|
||||
backgroundBlack = 2
|
||||
backgroundWhite = 3
|
||||
)
|
||||
|
||||
type engine struct {
|
||||
// inference is not thread-safe: one generation at a time.
|
||||
mu sync.Mutex
|
||||
stateMu sync.RWMutex
|
||||
|
||||
pipeline uintptr
|
||||
backend string
|
||||
caps int // bitmask of t2_caps
|
||||
textured bool // PBR texturing enabled
|
||||
models engineModels
|
||||
|
||||
abiVersion func() int32
|
||||
pipelineLoad func(dino, flow, dec, slat, slatHR, shapeDec, shapeEnc, texDec, texFlow, texFlowHR string, flags int32, err unsafe.Pointer, errLen int32) uintptr
|
||||
pipelineFree func(p uintptr)
|
||||
pipelineBackend func(p uintptr) string
|
||||
pipelineCaps func(p uintptr) int32
|
||||
generate func(p uintptr, img unsafe.Pointer, imgLen int32, pipelineType, backgroundMode int32,
|
||||
seed uint64, steps int32, guidance float32, textureSteps int32, cb uintptr, user unsafe.Pointer,
|
||||
preview uintptr, previewUser unsafe.Pointer,
|
||||
err unsafe.Pointer, errLen int32) uintptr
|
||||
meshNVerts func(r uintptr) int32
|
||||
meshNTris func(r uintptr) int32
|
||||
meshVerts func(r uintptr) uintptr
|
||||
meshNormals func(r uintptr) uintptr
|
||||
meshTris func(r uintptr) uintptr
|
||||
meshHasPBR func(r uintptr) int32
|
||||
meshPBR func(r uintptr) uintptr
|
||||
meshFree func(r uintptr)
|
||||
prepareMeshC func(verts unsafe.Pointer, nv int32, tris unsafe.Pointer, nt int32, pbr unsafe.Pointer,
|
||||
componentFilter int32, err unsafe.Pointer, errLen int32) uintptr
|
||||
printRemeshAvailable func() int32
|
||||
preparePrintMeshC func(verts unsafe.Pointer, nv int32, tris unsafe.Pointer, nt int32, pbr unsafe.Pointer,
|
||||
componentFilter int32, alphaRatio, offsetRatio float32,
|
||||
err unsafe.Pointer, errLen int32) uintptr
|
||||
bakeGLB func(verts unsafe.Pointer, nv int32, tris unsafe.Pointer, nt int32, pbr unsafe.Pointer,
|
||||
texSize, componentFilter int32, outLen unsafe.Pointer, err unsafe.Pointer, errLen int32) uintptr
|
||||
bakeProjectedGLB func(targetVerts unsafe.Pointer, targetNV int32, targetTris unsafe.Pointer, targetNT int32,
|
||||
sourceVerts unsafe.Pointer, sourceNV int32, sourceTris unsafe.Pointer, sourceNT int32,
|
||||
sourcePBR unsafe.Pointer, texSize, sourceComponentFilter int32,
|
||||
outLen unsafe.Pointer, err unsafe.Pointer, errLen int32) uintptr
|
||||
freeBuffer func(buf uintptr)
|
||||
}
|
||||
|
||||
// engineModels retains only the paths needed to recreate a freed pipeline. The
|
||||
// GGUF contents themselves remain owned by the C pipeline while it is loaded.
|
||||
type engineModels struct {
|
||||
dino, flow, dec string
|
||||
slat, slatHR, shapeDec string
|
||||
shapeEnc, texDec, texFlow, texFlowHR string
|
||||
}
|
||||
|
||||
// progressSink receives per-stage/step updates for the currently running
|
||||
// generation. Exactly one generation runs at a time (engine.mu), so a single
|
||||
// global callback + current sink is safe.
|
||||
var (
|
||||
progressMu sync.Mutex
|
||||
progressSink func(stage, step, total int)
|
||||
previewSink func(stage, step, total int, blob []byte)
|
||||
)
|
||||
|
||||
var (
|
||||
progressCallback uintptr // created once; purego callbacks are permanent
|
||||
previewCallback uintptr
|
||||
)
|
||||
|
||||
// slatGGUF/shapeDecGGUF may be "" for the coarse path; slatHRGGUF may be "" to
|
||||
// disable the 1024 cascade (512 fine only).
|
||||
func newEngine(libPath, dinoGGUF, flowGGUF, decGGUF, slatGGUF, slatHRGGUF, shapeDecGGUF,
|
||||
shapeEncGGUF, texDecGGUF, texFlowGGUF, texFlowHRGGUF string, startUnloaded bool) (*engine, error) {
|
||||
lib, err := purego.Dlopen(libPath, purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dlopen %s: %w", libPath, err)
|
||||
}
|
||||
|
||||
e := &engine{models: engineModels{
|
||||
dino: dinoGGUF, flow: flowGGUF, dec: decGGUF,
|
||||
slat: slatGGUF, slatHR: slatHRGGUF, shapeDec: shapeDecGGUF,
|
||||
shapeEnc: shapeEncGGUF, texDec: texDecGGUF,
|
||||
texFlow: texFlowGGUF, texFlowHR: texFlowHRGGUF,
|
||||
}}
|
||||
purego.RegisterLibFunc(&e.abiVersion, lib, "t2_abi_version")
|
||||
if got := e.abiVersion(); got != abiVersion {
|
||||
return nil, fmt.Errorf("ABI mismatch: library reports %d, server built for %d", got, abiVersion)
|
||||
}
|
||||
purego.RegisterLibFunc(&e.pipelineLoad, lib, "t2_pipeline_load")
|
||||
purego.RegisterLibFunc(&e.pipelineFree, lib, "t2_pipeline_free")
|
||||
purego.RegisterLibFunc(&e.pipelineBackend, lib, "t2_pipeline_backend")
|
||||
purego.RegisterLibFunc(&e.pipelineCaps, lib, "t2_pipeline_caps")
|
||||
purego.RegisterLibFunc(&e.generate, lib, "t2_generate")
|
||||
purego.RegisterLibFunc(&e.meshNVerts, lib, "t2_mesh_n_verts")
|
||||
purego.RegisterLibFunc(&e.meshNTris, lib, "t2_mesh_n_tris")
|
||||
purego.RegisterLibFunc(&e.meshVerts, lib, "t2_mesh_verts")
|
||||
purego.RegisterLibFunc(&e.meshNormals, lib, "t2_mesh_normals")
|
||||
purego.RegisterLibFunc(&e.meshTris, lib, "t2_mesh_tris")
|
||||
purego.RegisterLibFunc(&e.meshHasPBR, lib, "t2_mesh_has_pbr")
|
||||
purego.RegisterLibFunc(&e.meshPBR, lib, "t2_mesh_pbr")
|
||||
purego.RegisterLibFunc(&e.meshFree, lib, "t2_mesh_free")
|
||||
purego.RegisterLibFunc(&e.prepareMeshC, lib, "t2_prepare_mesh")
|
||||
purego.RegisterLibFunc(&e.printRemeshAvailable, lib, "t2_print_remesh_available")
|
||||
purego.RegisterLibFunc(&e.preparePrintMeshC, lib, "t2_prepare_print_mesh")
|
||||
purego.RegisterLibFunc(&e.bakeGLB, lib, "t2_bake_glb")
|
||||
purego.RegisterLibFunc(&e.bakeProjectedGLB, lib, "t2_bake_projected_glb")
|
||||
purego.RegisterLibFunc(&e.freeBuffer, lib, "t2_free_buffer")
|
||||
|
||||
if progressCallback == 0 {
|
||||
progressCallback = purego.NewCallback(func(user unsafe.Pointer, stage, step, total int32) uintptr {
|
||||
progressMu.Lock()
|
||||
sink := progressSink
|
||||
progressMu.Unlock()
|
||||
if sink != nil {
|
||||
sink(int(stage), int(step), int(total))
|
||||
}
|
||||
return 0
|
||||
})
|
||||
}
|
||||
if previewCallback == 0 {
|
||||
// Live intermediate-preview blobs (T2VOX01 voxel sets). `data` is only
|
||||
// valid during the call, so copy before handing it to the sink.
|
||||
previewCallback = purego.NewCallback(func(user unsafe.Pointer, stage, step, total int32,
|
||||
data unsafe.Pointer, length int32) uintptr {
|
||||
progressMu.Lock()
|
||||
sink := previewSink
|
||||
progressMu.Unlock()
|
||||
if sink != nil && data != nil && length > 0 {
|
||||
blob := make([]byte, int(length))
|
||||
copy(blob, unsafe.Slice((*byte)(data), int(length)))
|
||||
sink(int(stage), int(step), int(total), blob)
|
||||
}
|
||||
return 0
|
||||
})
|
||||
}
|
||||
|
||||
if startUnloaded {
|
||||
e.backend = "GPU (models unloaded)"
|
||||
if os.Getenv("TRELLIS2_DEVICE") == "cpu" {
|
||||
e.backend = "CPU (models unloaded)"
|
||||
}
|
||||
e.caps = configuredCaps(e.models)
|
||||
e.textured = e.caps&capTexture != 0
|
||||
} else {
|
||||
if err := e.loadLocked(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// configuredCaps mirrors t2_pipeline_caps from the model paths that main has
|
||||
// already existence-checked. It keeps /api/info and the quality picker useful
|
||||
// before a lazy first pipeline load.
|
||||
func configuredCaps(m engineModels) int {
|
||||
caps := capCoarse
|
||||
if m.slat != "" && m.shapeDec != "" {
|
||||
caps |= cap512
|
||||
if m.slatHR != "" {
|
||||
caps |= cap1024
|
||||
}
|
||||
if m.shapeEnc != "" && m.texDec != "" && m.texFlow != "" {
|
||||
caps |= capTexture
|
||||
}
|
||||
}
|
||||
return caps
|
||||
}
|
||||
|
||||
// loadLocked recreates the pipeline after an idle unload. e.mu must be held by
|
||||
// callers after construction.
|
||||
func (e *engine) loadLocked() error {
|
||||
if e.pipeline != 0 {
|
||||
return nil
|
||||
}
|
||||
m := e.models
|
||||
errBuf := make([]byte, 512)
|
||||
p := e.pipelineLoad(m.dino, m.flow, m.dec, m.slat, m.slatHR, m.shapeDec,
|
||||
m.shapeEnc, m.texDec, m.texFlow, m.texFlowHR,
|
||||
0 /*flags*/, unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if p == 0 {
|
||||
return fmt.Errorf("pipeline load: %s", cstr(errBuf))
|
||||
}
|
||||
e.stateMu.Lock()
|
||||
e.pipeline = p
|
||||
e.backend = e.pipelineBackend(p)
|
||||
e.caps = int(e.pipelineCaps(p))
|
||||
e.textured = e.caps&capTexture != 0
|
||||
e.stateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Unload releases all resident model buffers. Backend/capability metadata is
|
||||
// retained so /api/info and the quality picker remain useful while idle.
|
||||
func (e *engine) Unload() bool {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.pipeline == 0 {
|
||||
return false
|
||||
}
|
||||
e.pipelineFree(e.pipeline)
|
||||
e.stateMu.Lock()
|
||||
e.pipeline = 0
|
||||
e.stateMu.Unlock()
|
||||
return true
|
||||
}
|
||||
|
||||
func (e *engine) Info() (backend string, caps int, textured, loaded bool) {
|
||||
e.stateMu.RLock()
|
||||
defer e.stateMu.RUnlock()
|
||||
return e.backend, e.caps, e.textured, e.pipeline != 0
|
||||
}
|
||||
|
||||
type meshData struct {
|
||||
NVerts int
|
||||
NTris int
|
||||
Verts []float32 // 3 * NVerts
|
||||
Normals []float32 // 3 * NVerts
|
||||
Tris []int32 // 3 * NTris
|
||||
PBR []float32 // 6 * NVerts (base_color rgb, metallic, roughness, alpha); nil if untextured
|
||||
}
|
||||
|
||||
// Generate runs the full image->mesh pipeline. onProgress and onPreview may be
|
||||
// nil. onPreview receives live intermediate 3D preview blobs (T2VOX01 voxel
|
||||
// sets) as the sparse structure emerges.
|
||||
func (e *engine) Generate(image []byte, pipelineType, backgroundMode int, seed uint64, steps int, guidance float32, textureSteps int,
|
||||
onLoading func(),
|
||||
onProgress func(stage, step, total int),
|
||||
onPreview func(stage, step, total int, blob []byte)) (*meshData, error) {
|
||||
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.pipeline == 0 {
|
||||
if onLoading != nil {
|
||||
onLoading()
|
||||
}
|
||||
if err := e.loadLocked(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
progressMu.Lock()
|
||||
progressSink = onProgress
|
||||
previewSink = onPreview
|
||||
progressMu.Unlock()
|
||||
defer func() {
|
||||
progressMu.Lock()
|
||||
progressSink = nil
|
||||
previewSink = nil
|
||||
progressMu.Unlock()
|
||||
}()
|
||||
|
||||
cb := uintptr(0)
|
||||
if onProgress != nil {
|
||||
cb = progressCallback
|
||||
}
|
||||
pv := uintptr(0)
|
||||
if onPreview != nil {
|
||||
pv = previewCallback
|
||||
}
|
||||
|
||||
errBuf := make([]byte, 512)
|
||||
r := e.generate(e.pipeline, unsafe.Pointer(&image[0]), int32(len(image)), int32(pipelineType),
|
||||
int32(backgroundMode), seed, int32(steps), guidance, int32(textureSteps), cb, nil, pv, nil,
|
||||
unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if r == 0 {
|
||||
return nil, fmt.Errorf("%s", cstr(errBuf))
|
||||
}
|
||||
defer e.meshFree(r)
|
||||
|
||||
nv := int(e.meshNVerts(r))
|
||||
nt := int(e.meshNTris(r))
|
||||
if nv == 0 || nt == 0 {
|
||||
return nil, fmt.Errorf("empty mesh")
|
||||
}
|
||||
|
||||
m := &meshData{NVerts: nv, NTris: nt}
|
||||
m.Verts = copyFloats(e.meshVerts(r), 3*nv)
|
||||
m.Normals = copyFloats(e.meshNormals(r), 3*nv)
|
||||
m.Tris = copyInts(e.meshTris(r), 3*nt)
|
||||
if e.meshHasPBR(r) != 0 {
|
||||
m.PBR = copyFloats(e.meshPBR(r), 6*nv)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// PrepareMesh returns the exact component-filtered, full-density geometry used by
|
||||
// GLB export. It is CPU-only and does not require the model pipeline to be loaded.
|
||||
func (e *engine) PrepareMesh(m *meshData, componentFilter int) (*meshData, error) {
|
||||
if m == nil || m.NVerts == 0 || m.NTris == 0 {
|
||||
return nil, fmt.Errorf("empty mesh")
|
||||
}
|
||||
var pbr unsafe.Pointer
|
||||
if len(m.PBR) == 6*m.NVerts {
|
||||
pbr = unsafe.Pointer(&m.PBR[0])
|
||||
}
|
||||
errBuf := make([]byte, 512)
|
||||
r := e.prepareMeshC(unsafe.Pointer(&m.Verts[0]), int32(m.NVerts),
|
||||
unsafe.Pointer(&m.Tris[0]), int32(m.NTris), pbr,
|
||||
int32(componentFilter),
|
||||
unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if r == 0 {
|
||||
return nil, fmt.Errorf("%s", cstr(errBuf))
|
||||
}
|
||||
defer e.meshFree(r)
|
||||
nv, nt := int(e.meshNVerts(r)), int(e.meshNTris(r))
|
||||
if nv == 0 || nt == 0 {
|
||||
return nil, fmt.Errorf("empty prepared mesh")
|
||||
}
|
||||
out := &meshData{NVerts: nv, NTris: nt}
|
||||
out.Verts = copyFloats(e.meshVerts(r), 3*nv)
|
||||
out.Normals = copyFloats(e.meshNormals(r), 3*nv)
|
||||
out.Tris = copyInts(e.meshTris(r), 3*nt)
|
||||
if e.meshHasPBR(r) != 0 {
|
||||
out.PBR = copyFloats(e.meshPBR(r), 6*nv)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// HasPrintRemesh reports whether this library was built with CGAL Alpha Wrap.
|
||||
func (e *engine) HasPrintRemesh() bool {
|
||||
return e != nil && e.printRemeshAvailable != nil && e.printRemeshAvailable() != 0
|
||||
}
|
||||
|
||||
// PreparePrintMesh component-filters and wraps arbitrary source topology in a
|
||||
// watertight, oriented, intersection-free 2-manifold. Ratios are fractions of
|
||||
// the source bounding-box diagonal. Alpha Wrap creates new geometry, so when the
|
||||
// source is textured the material is projected onto the wrap vertices for an
|
||||
// approximate per-vertex preview; the GLB download rebakes it sharper per texel.
|
||||
func (e *engine) PreparePrintMesh(m *meshData, componentFilter int, alphaRatio, offsetRatio float32) (*meshData, error) {
|
||||
if m == nil || m.NVerts == 0 || m.NTris == 0 {
|
||||
return nil, fmt.Errorf("empty mesh")
|
||||
}
|
||||
if !e.HasPrintRemesh() || e.preparePrintMeshC == nil {
|
||||
return nil, fmt.Errorf("print remeshing is unavailable (library was built without CGAL)")
|
||||
}
|
||||
var pbr unsafe.Pointer
|
||||
if len(m.PBR) == 6*m.NVerts {
|
||||
pbr = unsafe.Pointer(&m.PBR[0])
|
||||
}
|
||||
errBuf := make([]byte, 512)
|
||||
r := e.preparePrintMeshC(unsafe.Pointer(&m.Verts[0]), int32(m.NVerts),
|
||||
unsafe.Pointer(&m.Tris[0]), int32(m.NTris), pbr,
|
||||
int32(componentFilter), alphaRatio, offsetRatio,
|
||||
unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if r == 0 {
|
||||
return nil, fmt.Errorf("%s", cstr(errBuf))
|
||||
}
|
||||
defer e.meshFree(r)
|
||||
nv, nt := int(e.meshNVerts(r)), int(e.meshNTris(r))
|
||||
if nv == 0 || nt == 0 {
|
||||
return nil, fmt.Errorf("empty print mesh")
|
||||
}
|
||||
out := &meshData{NVerts: nv, NTris: nt}
|
||||
out.Verts = copyFloats(e.meshVerts(r), 3*nv)
|
||||
out.Normals = copyFloats(e.meshNormals(r), 3*nv)
|
||||
out.Tris = copyInts(e.meshTris(r), 3*nt)
|
||||
if e.meshHasPBR(r) != 0 {
|
||||
out.PBR = copyFloats(e.meshPBR(r), 6*nv)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// BakeGLB turns a generated mesh into a portable vertex-coloured GLB. texSize
|
||||
// remains an atlas hint for the explicit T2GLB_XATLAS mode. Geometry retains its
|
||||
// original polygon density. CPU-only; it can run while the GPU is idle.
|
||||
func (e *engine) BakeGLB(m *meshData, texSize, componentFilter int) ([]byte, error) {
|
||||
if m == nil || m.NVerts == 0 || m.NTris == 0 {
|
||||
return nil, fmt.Errorf("empty mesh")
|
||||
}
|
||||
var pbr unsafe.Pointer
|
||||
if len(m.PBR) == 6*m.NVerts {
|
||||
pbr = unsafe.Pointer(&m.PBR[0])
|
||||
}
|
||||
var outLen int32
|
||||
errBuf := make([]byte, 512)
|
||||
p := e.bakeGLB(unsafe.Pointer(&m.Verts[0]), int32(m.NVerts),
|
||||
unsafe.Pointer(&m.Tris[0]), int32(m.NTris), pbr,
|
||||
int32(texSize), int32(componentFilter),
|
||||
unsafe.Pointer(&outLen), unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if p == 0 {
|
||||
return nil, fmt.Errorf("%s", cstr(errBuf))
|
||||
}
|
||||
defer e.freeBuffer(p)
|
||||
out := make([]byte, outLen)
|
||||
copy(out, unsafe.Slice((*byte)(unsafe.Pointer(p)), int(outLen)))
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// BakeProjectedGLB UV-unwraps replacement geometry and transfers the dense
|
||||
// source material per texel through the portable CGAL closest-surface backend.
|
||||
func (e *engine) BakeProjectedGLB(target, source *meshData, texSize, sourceComponentFilter int) ([]byte, error) {
|
||||
if target == nil || target.NVerts == 0 || target.NTris == 0 ||
|
||||
source == nil || source.NVerts == 0 || source.NTris == 0 || len(source.PBR) != 6*source.NVerts {
|
||||
return nil, fmt.Errorf("empty projected GLB mesh or missing source PBR")
|
||||
}
|
||||
if !e.HasPrintRemesh() || e.bakeProjectedGLB == nil {
|
||||
return nil, fmt.Errorf("PBR projection is unavailable (library was built without CGAL)")
|
||||
}
|
||||
var outLen int32
|
||||
errBuf := make([]byte, 512)
|
||||
p := e.bakeProjectedGLB(
|
||||
unsafe.Pointer(&target.Verts[0]), int32(target.NVerts),
|
||||
unsafe.Pointer(&target.Tris[0]), int32(target.NTris),
|
||||
unsafe.Pointer(&source.Verts[0]), int32(source.NVerts),
|
||||
unsafe.Pointer(&source.Tris[0]), int32(source.NTris),
|
||||
unsafe.Pointer(&source.PBR[0]), int32(texSize), int32(sourceComponentFilter),
|
||||
unsafe.Pointer(&outLen), unsafe.Pointer(&errBuf[0]), int32(len(errBuf)))
|
||||
if p == 0 {
|
||||
return nil, fmt.Errorf("%s", cstr(errBuf))
|
||||
}
|
||||
defer e.freeBuffer(p)
|
||||
out := make([]byte, outLen)
|
||||
copy(out, unsafe.Slice((*byte)(unsafe.Pointer(p)), int(outLen)))
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func copyFloats(p uintptr, n int) []float32 {
|
||||
src := unsafe.Slice((*float32)(unsafe.Pointer(p)), n)
|
||||
dst := make([]float32, n)
|
||||
copy(dst, src)
|
||||
return dst
|
||||
}
|
||||
|
||||
func copyInts(p uintptr, n int) []int32 {
|
||||
src := unsafe.Slice((*int32)(unsafe.Pointer(p)), n)
|
||||
dst := make([]int32, n)
|
||||
copy(dst, src)
|
||||
return dst
|
||||
}
|
||||
|
||||
func cstr(b []byte) string {
|
||||
for i, c := range b {
|
||||
if c == 0 {
|
||||
return string(b[:i])
|
||||
}
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
module trellis2-server
|
||||
|
||||
go 1.24
|
||||
|
||||
require github.com/ebitengine/purego v0.8.2
|
||||
@@ -0,0 +1,2 @@
|
||||
github.com/ebitengine/purego v0.8.2 h1:jPPGWs2sZ1UgOSgD2bClL0MJIqu58nOmIcBuXr62z1I=
|
||||
github.com/ebitengine/purego v0.8.2/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
|
||||
@@ -0,0 +1,134 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
func testEngine(pipeline uintptr, freed *int) *engine {
|
||||
return &engine{
|
||||
pipeline: pipeline,
|
||||
backend: "test GPU",
|
||||
caps: capCoarse | cap512 | capTexture,
|
||||
textured: true,
|
||||
pipelineFree: func(uintptr) {
|
||||
*freed++
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineUnloadRetainsRuntimeInfoAndReloads(t *testing.T) {
|
||||
freed, loads := 0, 0
|
||||
e := testEngine(41, &freed)
|
||||
e.models.dino = "dino.gguf"
|
||||
e.pipelineLoad = func(dino, flow, dec, slat, slatHR, shapeDec, shapeEnc, texDec, texFlow, texFlowHR string,
|
||||
flags int32, err unsafe.Pointer, errLen int32) uintptr {
|
||||
loads++
|
||||
if dino != "dino.gguf" {
|
||||
t.Fatalf("reload used dino path %q", dino)
|
||||
}
|
||||
return 42
|
||||
}
|
||||
e.pipelineBackend = func(uintptr) string { return "test GPU" }
|
||||
e.pipelineCaps = func(uintptr) int32 { return capCoarse | cap512 | capTexture }
|
||||
|
||||
if !e.Unload() || freed != 1 {
|
||||
t.Fatalf("Unload() = true with one free wanted; freed=%d", freed)
|
||||
}
|
||||
backend, caps, textured, loaded := e.Info()
|
||||
if loaded || backend != "test GPU" || caps != capCoarse|cap512|capTexture || !textured {
|
||||
t.Fatalf("Info() after unload = %q, %d, %v, %v", backend, caps, textured, loaded)
|
||||
}
|
||||
|
||||
e.mu.Lock()
|
||||
err := e.loadLocked()
|
||||
e.mu.Unlock()
|
||||
if err != nil || loads != 1 {
|
||||
t.Fatalf("reload: err=%v loads=%d", err, loads)
|
||||
}
|
||||
_, _, _, loaded = e.Info()
|
||||
if !loaded {
|
||||
t.Fatal("pipeline is not marked loaded after reload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfiguredCapsForLazyStartup(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
m engineModels
|
||||
want int
|
||||
}{
|
||||
{"coarse", engineModels{}, capCoarse},
|
||||
{"512", engineModels{slat: "slat", shapeDec: "shape"}, capCoarse | cap512},
|
||||
{"1024 textured", engineModels{
|
||||
slat: "slat", slatHR: "hr", shapeDec: "shape", shapeEnc: "shapeenc",
|
||||
texDec: "texdec", texFlow: "texflow",
|
||||
}, capCoarse | cap512 | cap1024 | capTexture},
|
||||
{"texture missing encoder", engineModels{
|
||||
slat: "slat", shapeDec: "shape", texDec: "texdec", texFlow: "texflow",
|
||||
}, capCoarse | cap512},
|
||||
{"incomplete fine", engineModels{slat: "slat", texDec: "texdec", texFlow: "texflow"}, capCoarse},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := configuredCaps(tt.m); got != tt.want {
|
||||
t.Fatalf("configuredCaps() = %d, want %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdlePolicyWaitsForQueuedWork(t *testing.T) {
|
||||
freed := 0
|
||||
e := testEngine(41, &freed)
|
||||
s := &server{eng: e, jobs: map[string]*job{}, q: make(chan *job, 1), queued: 1}
|
||||
|
||||
if s.setUnloadIdle(true) {
|
||||
t.Fatal("enabled idle policy unloaded with queued work")
|
||||
}
|
||||
if freed != 0 {
|
||||
t.Fatalf("pipeline freed with queued work: %d", freed)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
s.queued = 0
|
||||
s.mu.Unlock()
|
||||
if !s.unloadModelsIfIdle(nil) || freed != 1 {
|
||||
t.Fatalf("idle unload failed; freed=%d", freed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseExportOptions(t *testing.T) {
|
||||
def := parseExportOptions(httptest.NewRequest("GET", "/api/glb/job", nil))
|
||||
if def.componentFilter != 2 || def.printWrap || def.prepareKey() != "2" {
|
||||
t.Fatalf("default export should preserve all components: %+v", def)
|
||||
}
|
||||
r := httptest.NewRequest("GET", "/api/glb/job?tex=1024&components=largest", nil)
|
||||
o := parseExportOptions(r)
|
||||
if o.textureSize != 1024 || o.componentFilter != 1 {
|
||||
t.Fatalf("parseExportOptions() = %+v", o)
|
||||
}
|
||||
if o.prepareKey() != "1" || o.glbKey() != "1024-1" {
|
||||
t.Fatalf("unexpected export cache keys: %q %q", o.prepareKey(), o.glbKey())
|
||||
}
|
||||
wrapped := parseExportOptions(httptest.NewRequest("GET",
|
||||
"/api/glb/job?print=1&alpha=2.5&offset=0.1", nil))
|
||||
if !wrapped.printWrap || wrapped.alphaRatio != 0.025 || wrapped.offsetRatio != 0.001 ||
|
||||
wrapped.prepareKey() == def.prepareKey() {
|
||||
t.Fatalf("print-wrap options were not parsed/cached independently: %+v", wrapped)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeepAllExportPreviewUsesOriginalMesh(t *testing.T) {
|
||||
original := &meshData{NVerts: 3, NTris: 1}
|
||||
j := &job{mesh: original}
|
||||
s := &server{}
|
||||
got, err := s.preparedExportMesh(j, exportOptions{textureSize: 2048, componentFilter: 2})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != original {
|
||||
t.Fatal("keep-all export preview did not return the exact source mesh")
|
||||
}
|
||||
}
|
||||
+1148
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,349 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const persistedJobVersion = 1
|
||||
|
||||
// persistedJob is deliberately independent of the public job JSON. Runtime
|
||||
// state, caches, and locks never reach disk; only the data needed to restore a
|
||||
// completed generation is retained.
|
||||
type persistedJob struct {
|
||||
Version int `json:"version"`
|
||||
ID string `json:"id"`
|
||||
CreatedAt int64 `json:"createdAt"`
|
||||
StartedAt int64 `json:"startedAt,omitempty"`
|
||||
FinishedAt int64 `json:"finishedAt,omitempty"`
|
||||
DurationMS int64 `json:"durationMs,omitempty"`
|
||||
Quality string `json:"quality,omitempty"`
|
||||
Thumbnail string `json:"thumbnail,omitempty"`
|
||||
Pipeline int `json:"pipeline"`
|
||||
Seed uint64 `json:"seed"`
|
||||
Steps int `json:"steps"`
|
||||
TextureSteps int `json:"textureSteps"`
|
||||
Guidance float32 `json:"guidance"`
|
||||
Frames []frameMeta `json:"frames,omitempty"`
|
||||
LivePreview bool `json:"livePreview,omitempty"`
|
||||
StageTimings []stageTiming `json:"stageTimings,omitempty"`
|
||||
}
|
||||
|
||||
func persistedFrameName(i int) string {
|
||||
return filepath.Join("frames", fmt.Sprintf("%06d.bin", i))
|
||||
}
|
||||
|
||||
// persistJob writes a complete job into a temporary sibling directory and then
|
||||
// renames that directory into place. Startup therefore sees either the previous
|
||||
// complete job or no job, never a half-written mesh/frame set.
|
||||
func (s *server) persistJob(j *job) error {
|
||||
if s.storeDir == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
j.mu.Lock()
|
||||
if j.mesh == nil {
|
||||
j.mu.Unlock()
|
||||
return fmt.Errorf("job %s has no mesh", j.ID)
|
||||
}
|
||||
manifest := persistedJob{
|
||||
Version: persistedJobVersion, ID: j.ID, CreatedAt: j.CreatedAt,
|
||||
StartedAt: j.StartedAt, FinishedAt: j.FinishedAt, DurationMS: j.DurationMS,
|
||||
Quality: j.Quality, Thumbnail: j.Thumbnail,
|
||||
Pipeline: j.pipeline, Seed: j.seed, Steps: j.steps,
|
||||
TextureSteps: j.textureSteps, Guidance: j.guidance,
|
||||
Frames: append([]frameMeta(nil), j.Frames...), LivePreview: j.LivePreview,
|
||||
StageTimings: append([]stageTiming(nil), j.StageTimings...),
|
||||
}
|
||||
mesh := j.mesh
|
||||
input := j.image
|
||||
source := j.source
|
||||
if len(source) == 0 {
|
||||
source = input
|
||||
}
|
||||
previews := append([][]byte(nil), j.previews...)
|
||||
j.mu.Unlock()
|
||||
|
||||
if len(previews) != len(manifest.Frames) {
|
||||
return fmt.Errorf("job %s has %d previews but %d frame records",
|
||||
j.ID, len(previews), len(manifest.Frames))
|
||||
}
|
||||
if err := os.MkdirAll(s.storeDir, 0o755); err != nil {
|
||||
return fmt.Errorf("create job store: %w", err)
|
||||
}
|
||||
tmp, err := os.MkdirTemp(s.storeDir, "."+j.ID+".tmp-")
|
||||
if err != nil {
|
||||
return fmt.Errorf("create temporary job directory: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(tmp)
|
||||
|
||||
if err := writeMeshFile(filepath.Join(tmp, "mesh.t2mesh"), mesh); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(source) > 0 {
|
||||
if err := os.WriteFile(filepath.Join(tmp, "source.img"), source, 0o644); err != nil {
|
||||
return fmt.Errorf("write source image: %w", err)
|
||||
}
|
||||
}
|
||||
if len(input) > 0 {
|
||||
if err := os.WriteFile(filepath.Join(tmp, "input.img"), input, 0o644); err != nil {
|
||||
return fmt.Errorf("write generation input: %w", err)
|
||||
}
|
||||
}
|
||||
if len(previews) > 0 {
|
||||
if err := os.Mkdir(filepath.Join(tmp, "frames"), 0o755); err != nil {
|
||||
return fmt.Errorf("create frame directory: %w", err)
|
||||
}
|
||||
for i, blob := range previews {
|
||||
if err := os.WriteFile(filepath.Join(tmp, persistedFrameName(i)), blob, 0o644); err != nil {
|
||||
return fmt.Errorf("write preview %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
data, err := json.MarshalIndent(manifest, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode manifest: %w", err)
|
||||
}
|
||||
data = append(data, '\n')
|
||||
if err := os.WriteFile(filepath.Join(tmp, "manifest.json"), data, 0o644); err != nil {
|
||||
return fmt.Errorf("write manifest: %w", err)
|
||||
}
|
||||
|
||||
finalDir := filepath.Join(s.storeDir, j.ID)
|
||||
if err := os.Rename(tmp, finalDir); err != nil {
|
||||
return fmt.Errorf("commit job %s: %w", j.ID, err)
|
||||
}
|
||||
|
||||
// Frames are immutable and now durable, so release their duplicate in-memory
|
||||
// copies. The preview endpoint transparently reads them from finalDir.
|
||||
j.mu.Lock()
|
||||
j.persistDir = finalDir
|
||||
j.meshPath = filepath.Join(finalDir, "mesh.t2mesh")
|
||||
if len(input) > 0 {
|
||||
j.inputPath = filepath.Join(finalDir, "input.img")
|
||||
}
|
||||
if len(source) > 0 {
|
||||
j.sourcePath = filepath.Join(finalDir, "source.img")
|
||||
}
|
||||
j.previews = nil
|
||||
j.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// restoreJobs scans only complete, versioned job directories. Temporary
|
||||
// directories left by a killed write and malformed/corrupt entries are ignored
|
||||
// with a log message; one bad asset must not prevent the server from starting.
|
||||
func (s *server) restoreJobs() (int, error) {
|
||||
if s.storeDir == "" {
|
||||
return 0, nil
|
||||
}
|
||||
if err := os.MkdirAll(s.storeDir, 0o755); err != nil {
|
||||
return 0, fmt.Errorf("create job store: %w", err)
|
||||
}
|
||||
entries, err := os.ReadDir(s.storeDir)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("read job store: %w", err)
|
||||
}
|
||||
restored := 0
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") {
|
||||
continue
|
||||
}
|
||||
dir := filepath.Join(s.storeDir, entry.Name())
|
||||
j, err := loadPersistedJob(dir)
|
||||
if err != nil {
|
||||
log.Printf("ignoring persisted job %s: %v", entry.Name(), err)
|
||||
continue
|
||||
}
|
||||
s.mu.Lock()
|
||||
if _, exists := s.jobs[j.ID]; !exists {
|
||||
s.jobs[j.ID] = j
|
||||
restored++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
return restored, nil
|
||||
}
|
||||
|
||||
func loadPersistedJob(dir string) (*job, error) {
|
||||
data, err := os.ReadFile(filepath.Join(dir, "manifest.json"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read manifest: %w", err)
|
||||
}
|
||||
var m persistedJob
|
||||
if err := json.Unmarshal(data, &m); err != nil {
|
||||
return nil, fmt.Errorf("decode manifest: %w", err)
|
||||
}
|
||||
if m.Version != persistedJobVersion {
|
||||
return nil, fmt.Errorf("unsupported manifest version %d", m.Version)
|
||||
}
|
||||
if m.ID == "" || filepath.Base(m.ID) != m.ID || m.ID != filepath.Base(dir) {
|
||||
return nil, fmt.Errorf("invalid job id %q", m.ID)
|
||||
}
|
||||
meshPath := filepath.Join(dir, "mesh.t2mesh")
|
||||
if st, err := os.Stat(meshPath); err != nil || !st.Mode().IsRegular() {
|
||||
if err == nil {
|
||||
err = fmt.Errorf("not a regular file")
|
||||
}
|
||||
return nil, fmt.Errorf("mesh: %w", err)
|
||||
}
|
||||
sourcePath := filepath.Join(dir, "source.img")
|
||||
if st, err := os.Stat(sourcePath); err != nil || !st.Mode().IsRegular() {
|
||||
sourcePath = "" // optional for generations saved before source retention
|
||||
}
|
||||
inputPath := filepath.Join(dir, "input.img")
|
||||
if st, err := os.Stat(inputPath); err != nil || !st.Mode().IsRegular() {
|
||||
inputPath = "" // legacy jobs stored their processed input as source.img
|
||||
}
|
||||
for i := range m.Frames {
|
||||
if st, err := os.Stat(filepath.Join(dir, persistedFrameName(i))); err != nil || !st.Mode().IsRegular() {
|
||||
if err == nil {
|
||||
err = fmt.Errorf("not a regular file")
|
||||
}
|
||||
return nil, fmt.Errorf("preview %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
return &job{
|
||||
ID: m.ID, State: "done", CreatedAt: m.CreatedAt, StartedAt: m.StartedAt,
|
||||
FinishedAt: m.FinishedAt, DurationMS: m.DurationMS,
|
||||
Quality: m.Quality, Thumbnail: m.Thumbnail,
|
||||
PreviewSeq: len(m.Frames), Frames: m.Frames,
|
||||
LivePreview: m.LivePreview || len(m.Frames) > 0, StageTimings: m.StageTimings,
|
||||
pipeline: m.Pipeline, seed: m.Seed, steps: m.Steps,
|
||||
textureSteps: m.TextureSteps, guidance: m.Guidance,
|
||||
persistDir: dir, meshPath: meshPath, inputPath: inputPath, sourcePath: sourcePath,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *server) loadJobMesh(j *job) (*meshData, error) {
|
||||
j.mu.Lock()
|
||||
defer j.mu.Unlock()
|
||||
if j.mesh != nil {
|
||||
return j.mesh, nil
|
||||
}
|
||||
if j.meshPath == "" {
|
||||
return nil, fmt.Errorf("mesh not ready")
|
||||
}
|
||||
mesh, err := readMeshFile(j.meshPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
j.mesh = mesh
|
||||
return mesh, nil
|
||||
}
|
||||
|
||||
func loadJobPreview(j *job, seq int) ([]byte, error) {
|
||||
j.mu.Lock()
|
||||
if seq < 0 || seq >= j.PreviewSeq {
|
||||
j.mu.Unlock()
|
||||
return nil, fmt.Errorf("no such preview frame")
|
||||
}
|
||||
if seq < len(j.previews) && j.previews[seq] != nil {
|
||||
blob := j.previews[seq]
|
||||
j.mu.Unlock()
|
||||
return blob, nil
|
||||
}
|
||||
dir := j.persistDir
|
||||
j.mu.Unlock()
|
||||
if dir == "" {
|
||||
return nil, fmt.Errorf("no such preview frame")
|
||||
}
|
||||
return os.ReadFile(filepath.Join(dir, persistedFrameName(seq)))
|
||||
}
|
||||
|
||||
func writeMeshFile(path string, mesh *meshData) error {
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o644)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create mesh: %w", err)
|
||||
}
|
||||
if err := writeMeshBinary(f, mesh); err != nil {
|
||||
f.Close()
|
||||
return fmt.Errorf("write mesh: %w", err)
|
||||
}
|
||||
if err := f.Sync(); err != nil {
|
||||
f.Close()
|
||||
return fmt.Errorf("sync mesh: %w", err)
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
return fmt.Errorf("close mesh: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func readMeshFile(path string) (*meshData, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open mesh: %w", err)
|
||||
}
|
||||
defer f.Close()
|
||||
var magic [8]byte
|
||||
if _, err := io.ReadFull(f, magic[:]); err != nil {
|
||||
return nil, fmt.Errorf("read mesh magic: %w", err)
|
||||
}
|
||||
pbrWidth := 0
|
||||
switch string(magic[:]) {
|
||||
case "T2MESH01":
|
||||
case "T2MESH02":
|
||||
pbrWidth = 5
|
||||
case "T2MESH03":
|
||||
pbrWidth = 6
|
||||
default:
|
||||
return nil, fmt.Errorf("bad mesh magic %q", magic)
|
||||
}
|
||||
var nv32, nt32 uint32
|
||||
if err := binary.Read(f, binary.LittleEndian, &nv32); err != nil {
|
||||
return nil, fmt.Errorf("read vertex count: %w", err)
|
||||
}
|
||||
if err := binary.Read(f, binary.LittleEndian, &nt32); err != nil {
|
||||
return nil, fmt.Errorf("read triangle count: %w", err)
|
||||
}
|
||||
if nv32 == 0 || nt32 == 0 || nv32 > 100_000_000 || nt32 > 100_000_000 {
|
||||
return nil, fmt.Errorf("invalid mesh size %d vertices, %d triangles", nv32, nt32)
|
||||
}
|
||||
nv, nt := int(nv32), int(nt32)
|
||||
st, err := f.Stat()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stat mesh: %w", err)
|
||||
}
|
||||
expected := int64(16) + int64(nv)*24 + int64(nv*pbrWidth)*4 + int64(nt)*12
|
||||
if st.Size() != expected {
|
||||
return nil, fmt.Errorf("mesh size is %d bytes, expected %d", st.Size(), expected)
|
||||
}
|
||||
m := &meshData{NVerts: nv, NTris: nt, Verts: make([]float32, 3*nv), Normals: make([]float32, 3*nv)}
|
||||
if err := binary.Read(f, binary.LittleEndian, m.Verts); err != nil {
|
||||
return nil, fmt.Errorf("read vertices: %w", err)
|
||||
}
|
||||
if err := binary.Read(f, binary.LittleEndian, m.Normals); err != nil {
|
||||
return nil, fmt.Errorf("read normals: %w", err)
|
||||
}
|
||||
if pbrWidth != 0 {
|
||||
stored := make([]float32, pbrWidth*nv)
|
||||
if err := binary.Read(f, binary.LittleEndian, stored); err != nil {
|
||||
return nil, fmt.Errorf("read PBR attributes: %w", err)
|
||||
}
|
||||
if pbrWidth == 6 {
|
||||
m.PBR = stored
|
||||
} else {
|
||||
m.PBR = make([]float32, 6*nv)
|
||||
for i := 0; i < nv; i++ {
|
||||
copy(m.PBR[6*i:6*i+5], stored[5*i:5*i+5])
|
||||
m.PBR[6*i+5] = 1
|
||||
}
|
||||
}
|
||||
}
|
||||
m.Tris = make([]int32, 3*nt)
|
||||
if err := binary.Read(f, binary.LittleEndian, m.Tris); err != nil {
|
||||
return nil, fmt.Errorf("read triangles: %w", err)
|
||||
}
|
||||
var extra [1]byte
|
||||
if n, err := f.Read(extra[:]); err != io.EOF || n != 0 {
|
||||
return nil, fmt.Errorf("mesh contains trailing data")
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func testPersistedMesh(textured bool) *meshData {
|
||||
m := &meshData{
|
||||
NVerts: 3, NTris: 1,
|
||||
Verts: []float32{0, 0, 0, 1, 0, 0, 0, 1, 0},
|
||||
Normals: []float32{0, 0, 1, 0, 0, 1, 0, 0, 1},
|
||||
Tris: []int32{0, 1, 2},
|
||||
}
|
||||
if textured {
|
||||
m.PBR = []float32{
|
||||
1, 0, 0, 0.1, 0.2, 1,
|
||||
0, 1, 0, 0.3, 0.4, 0.9,
|
||||
0, 0, 1, 0.5, 0.6, 0.8,
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func TestMeshFileRoundTrip(t *testing.T) {
|
||||
for _, textured := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "geometry", true: "pbr"}[textured], func(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "mesh.t2mesh")
|
||||
want := testPersistedMesh(textured)
|
||||
if err := writeMeshFile(path, want); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := readMeshFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("mesh round trip mismatch\n got: %#v\nwant: %#v", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPersistAndRestoreCompletedJob(t *testing.T) {
|
||||
store := t.TempDir()
|
||||
s := &server{jobs: map[string]*job{}, storeDir: store}
|
||||
wantFrames := []frameMeta{
|
||||
{Stage: "sampling sparse structure", Step: 1, Total: 2, Kind: "voxel"},
|
||||
{Stage: "sampling shape SLAT", Step: 2, Total: 2, Kind: "mesh"},
|
||||
}
|
||||
wantPreviews := [][]byte{[]byte("T2VOX01-frame"), []byte("T2MESH01-frame")}
|
||||
j := &job{
|
||||
ID: "0123456789abcdef", State: "running", CreatedAt: 123456789,
|
||||
StartedAt: 123456800, FinishedAt: 123499000, DurationMS: 42200,
|
||||
Quality: "1024", Thumbnail: "data:image/jpeg;base64,dGVzdA==",
|
||||
PreviewSeq: len(wantFrames), Frames: wantFrames, LivePreview: true,
|
||||
StageTimings: []stageTiming{{Stage: "sampling sparse structure", Milliseconds: 21000}},
|
||||
pipeline: pipe1024, seed: 42, steps: 12, textureSteps: 10, guidance: 7.5,
|
||||
previews: wantPreviews, mesh: testPersistedMesh(true),
|
||||
image: []byte("processed-input"), source: []byte("exact-original"),
|
||||
}
|
||||
if err := s.persistJob(j); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(j.previews) != 0 || j.persistDir == "" || j.meshPath == "" ||
|
||||
j.inputPath == "" || j.sourcePath == "" {
|
||||
t.Fatalf("persist did not switch assets to disk: previews=%d dir=%q mesh=%q input=%q source=%q",
|
||||
len(j.previews), j.persistDir, j.meshPath, j.inputPath, j.sourcePath)
|
||||
}
|
||||
|
||||
restarted := &server{jobs: map[string]*job{}, storeDir: store}
|
||||
n, err := restarted.restoreJobs()
|
||||
if err != nil || n != 1 {
|
||||
t.Fatalf("restoreJobs() = %d, %v; want 1, nil", n, err)
|
||||
}
|
||||
got := restarted.jobs[j.ID]
|
||||
if got == nil || got.State != "done" || got.Quality != "1024" ||
|
||||
got.Thumbnail != j.Thumbnail || got.CreatedAt != j.CreatedAt ||
|
||||
got.StartedAt != j.StartedAt || got.FinishedAt != j.FinishedAt ||
|
||||
got.DurationMS != j.DurationMS || got.LivePreview != j.LivePreview ||
|
||||
!reflect.DeepEqual(got.StageTimings, j.StageTimings) {
|
||||
t.Fatalf("restored metadata = %#v", got)
|
||||
}
|
||||
if got.mesh != nil || len(got.previews) != 0 {
|
||||
t.Fatal("restored binary assets should remain lazy until requested")
|
||||
}
|
||||
if source, err := os.ReadFile(got.sourcePath); err != nil || string(source) != "exact-original" {
|
||||
t.Fatalf("restored source image = %q, %v", source, err)
|
||||
}
|
||||
if input, err := os.ReadFile(got.inputPath); err != nil || string(input) != "processed-input" {
|
||||
t.Fatalf("restored generation input = %q, %v", input, err)
|
||||
}
|
||||
rr := httptest.NewRecorder()
|
||||
restarted.handleSource(rr, httptest.NewRequest(http.MethodGet, "/api/source/"+j.ID, nil))
|
||||
if rr.Code != http.StatusOK || rr.Body.String() != "exact-original" {
|
||||
t.Fatalf("GET source = %d %q", rr.Code, rr.Body.String())
|
||||
}
|
||||
mesh, err := restarted.loadJobMesh(got)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(mesh, j.mesh) {
|
||||
t.Fatal("restored mesh mismatch")
|
||||
}
|
||||
for i, want := range wantPreviews {
|
||||
blob, err := loadJobPreview(got, i)
|
||||
if err != nil {
|
||||
t.Fatalf("preview %d: %v", i, err)
|
||||
}
|
||||
if !bytes.Equal(blob, want) {
|
||||
t.Fatalf("preview %d = %q, want %q", i, blob, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegenerateUsesPersistedInputWithoutUpload(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
inputPath := filepath.Join(dir, "input.img")
|
||||
sourcePath := filepath.Join(dir, "source.img")
|
||||
if err := os.WriteFile(inputPath, []byte("processed-input"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(sourcePath, []byte("exact-original"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
old := &job{
|
||||
ID: "old", State: "done", CreatedAt: 1, Quality: "512", Thumbnail: "thumb",
|
||||
inputPath: inputPath, sourcePath: sourcePath, pipeline: pipe512,
|
||||
seed: 7, steps: 8, textureSteps: 9, guidance: 6.5,
|
||||
}
|
||||
s := &server{jobs: map[string]*job{"old": old}, q: make(chan *job, 1)}
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/regenerate/old",
|
||||
strings.NewReader("quality=1024&seed=42&steps=14&texture_steps=15&guidance=8&preview=0"))
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
rr := httptest.NewRecorder()
|
||||
s.handleRegenerate(rr, req)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("POST regenerate = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var response map[string]string
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
regenerated := s.jobs[response["job"]]
|
||||
if regenerated == nil {
|
||||
t.Fatal("regenerated job was not queued")
|
||||
}
|
||||
if string(regenerated.image) != "processed-input" || string(regenerated.source) != "exact-original" {
|
||||
t.Fatalf("regenerated bytes = input %q, source %q", regenerated.image, regenerated.source)
|
||||
}
|
||||
if regenerated.Quality != "1024" || regenerated.pipeline != pipe1024 ||
|
||||
regenerated.background != backgroundKeep || regenerated.seed != 42 ||
|
||||
regenerated.steps != 14 || regenerated.textureSteps != 15 ||
|
||||
regenerated.guidance != 8 || regenerated.LivePreview {
|
||||
t.Fatalf("regenerated settings = %#v", regenerated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateKeepsOriginalSeparateFromProcessedInput(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
mw := multipart.NewWriter(&body)
|
||||
inputPart, err := mw.CreateFormFile("image", "input.png")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inputPart.Write([]byte("processed-png"))
|
||||
sourcePart, err := mw.CreateFormFile("source", "camera-original.webp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sourcePart.Write([]byte("exact-original-webp"))
|
||||
if err := mw.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := &server{jobs: map[string]*job{}, q: make(chan *job, 1)}
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/generate", &body)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
rr := httptest.NewRecorder()
|
||||
s.handleGenerate(rr, req)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("POST generate = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var response map[string]string
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
j := s.jobs[response["job"]]
|
||||
if j == nil || string(j.image) != "processed-png" || string(j.source) != "exact-original-webp" {
|
||||
t.Fatalf("queued job bytes = %#v", j)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerHistoryListAndDelete(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
oldDir := filepath.Join(root, "old")
|
||||
if err := os.Mkdir(oldDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
old := &job{ID: "old", State: "done", CreatedAt: 10, Quality: "512", persistDir: oldDir}
|
||||
newer := &job{ID: "new", State: "done", CreatedAt: 20, Quality: "1024", Thumbnail: "thumb"}
|
||||
active := &job{ID: "active", State: "running", CreatedAt: 30}
|
||||
s := &server{jobs: map[string]*job{"old": old, "new": newer, "active": active}}
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
s.handleJobs(rr, httptest.NewRequest(http.MethodGet, "/api/jobs", nil))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/jobs = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var got []jobSummary
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 2 || got[0].ID != "new" || got[1].ID != "old" {
|
||||
t.Fatalf("history order/content = %#v", got)
|
||||
}
|
||||
|
||||
rr = httptest.NewRecorder()
|
||||
s.handleJob(rr, httptest.NewRequest(http.MethodDelete, "/api/job/old", nil))
|
||||
if rr.Code != http.StatusNoContent {
|
||||
t.Fatalf("DELETE /api/job/old = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if s.jobs["old"] != nil {
|
||||
t.Fatal("deleted job remains in server index")
|
||||
}
|
||||
if _, err := os.Stat(oldDir); !os.IsNotExist(err) {
|
||||
t.Fatalf("deleted job directory still exists: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreIgnoresIncompleteTemporaryDirectory(t *testing.T) {
|
||||
store := t.TempDir()
|
||||
if err := os.Mkdir(filepath.Join(store, ".unfinished.tmp-123"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := &server{jobs: map[string]*job{}, storeDir: store}
|
||||
n, err := s.restoreJobs()
|
||||
if err != nil || n != 0 {
|
||||
t.Fatalf("restoreJobs() = %d, %v; want 0, nil", n, err)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user