mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
Richard Palethorpe
799cc9f211
feat: bound backend admission and expose running traces Add process-wide backend execution admission without blocking UI or administrative HTTP work. Represent backend operations while they are in flight, surface running traces with immediate log links, and tie streaming admission leases to the gRPC receive lifecycle. Assisted-by: OpenAI Codex: GPT-5 Signed-off-by: Richard Palethorpe <io@richiejp.com>
110 lines
3.3 KiB
Go
110 lines
3.3 KiB
Go
package backend
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/mudler/LocalAI/core/config"
|
|
"github.com/mudler/LocalAI/core/trace"
|
|
|
|
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
|
model "github.com/mudler/LocalAI/pkg/model"
|
|
)
|
|
|
|
func ImageGeneration(ctx context.Context, height, width, step, seed int, positive_prompt, negative_prompt, src, dst string, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig, refImages []string) (func() error, error) {
|
|
|
|
// model.WithContext(ctx) overrides the app-context default set in
|
|
// ModelOptions so distributed routing decisions reach the request's
|
|
// X-LocalAI-Node holder via distributedhdr.Stamp.
|
|
opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
|
|
inferenceModel, err := loader.Load(
|
|
opts...,
|
|
)
|
|
if err != nil {
|
|
recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)
|
|
return nil, err
|
|
}
|
|
|
|
fn := func() error {
|
|
_, err := inferenceModel.GenerateImage(
|
|
ctx,
|
|
&proto.GenerateImageRequest{
|
|
// ModelIdentity is ModelConfig.Model — the SAME expression
|
|
// ModelOptions feeds to model.WithModel above, which the
|
|
// backend receives as ModelOptions.Model at LoadModel. Equal
|
|
// by construction, so the check cannot false-reject (#10952).
|
|
ModelIdentity: modelConfig.Model,
|
|
Height: int32(height),
|
|
Width: int32(width),
|
|
Step: int32(step),
|
|
Seed: int32(seed),
|
|
CLIPSkip: int32(modelConfig.Diffusers.ClipSkip),
|
|
PositivePrompt: positive_prompt,
|
|
NegativePrompt: negative_prompt,
|
|
Dst: dst,
|
|
Src: src,
|
|
EnableParameters: modelConfig.Diffusers.EnableParameters,
|
|
RefImages: refImages,
|
|
})
|
|
return err
|
|
}
|
|
|
|
if appConfig.EnableTracing {
|
|
trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes)
|
|
|
|
traceData := map[string]any{
|
|
"positive_prompt": positive_prompt,
|
|
"negative_prompt": negative_prompt,
|
|
"height": height,
|
|
"width": width,
|
|
"step": step,
|
|
"seed": seed,
|
|
"source_image": src,
|
|
"destination": dst,
|
|
}
|
|
|
|
originalFn := fn
|
|
fn = func() error {
|
|
startTime := time.Now()
|
|
traceID := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceImageGeneration, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(positive_prompt, 200)})
|
|
defer trace.CancelBackendTrace(traceID)
|
|
err := originalFn()
|
|
duration := time.Since(startTime)
|
|
|
|
errStr := ""
|
|
if err != nil {
|
|
errStr = err.Error()
|
|
}
|
|
|
|
trace.RecordBackendTrace(trace.BackendTrace{
|
|
ID: traceID,
|
|
Timestamp: startTime,
|
|
Duration: duration,
|
|
Type: trace.BackendTraceImageGeneration,
|
|
ModelName: modelConfig.Name,
|
|
Backend: modelConfig.Backend,
|
|
Summary: trace.TruncateString(positive_prompt, 200),
|
|
Error: errStr,
|
|
Data: traceData,
|
|
})
|
|
|
|
return err
|
|
}
|
|
}
|
|
originalFn := fn
|
|
fn = func() error {
|
|
release, err := AcquireGlobalBackendSlot()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer release()
|
|
return originalFn()
|
|
}
|
|
|
|
return fn, nil
|
|
}
|
|
|
|
// ImageGenerationFunc is a test-friendly indirection to call image generation logic.
|
|
// Tests can override this variable to provide a stub implementation.
|
|
var ImageGenerationFunc = ImageGeneration
|