mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-03 04:46:54 -04:00
d0e6bf3aa7
Image generation (and the tts/transcript/embeddings/vad/rerank/llm helpers) pass the request context to loader.Load so distributed routing decisions reach the request's X-LocalAI-Node holder. That context also governs cancellation of the load, so when a client disconnects mid-load the LoadModel RPC is aborted, stopLoadProcess tears down the backend process, and every retry restarts from scratch. Heavy diffusers/LLM models on a slow host (e.g. a shared-memory iGPU) take long enough to load that the request routinely ends first, so the model never finishes loading and the UI shows "NetworkError when attempting to fetch resource". Wrap the load context with context.WithoutCancel: the routing holder value still propagates, but the request's cancellation no longer aborts the load, so it runs to completion and caches for the next request. Inference keeps the cancellable request context, so a disconnect still stops generation. Adds a regression spec asserting a canceled request context does not cancel the model load while the routing holder still reaches the router. Fixes #10636 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Assisted-by: Claude:claude-opus-4-8 [Claude Code]
97 lines
2.9 KiB
Go
97 lines
2.9 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 carries the request context into the load so distributed
|
|
// routing decisions reach the request's X-LocalAI-Node holder via
|
|
// distributedhdr.Stamp. context.WithoutCancel keeps those values but drops
|
|
// the request's cancellation, so a slow first load still completes and
|
|
// caches if the client disconnects instead of aborting the LoadModel RPC and
|
|
// tearing down the backend process (issue #10636). Inference below keeps the
|
|
// cancellable ctx, so a disconnect still stops generation.
|
|
opts := ModelOptions(modelConfig, appConfig, model.WithContext(context.WithoutCancel(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{
|
|
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,
|
|
}
|
|
|
|
startTime := time.Now()
|
|
originalFn := fn
|
|
fn = func() error {
|
|
err := originalFn()
|
|
duration := time.Since(startTime)
|
|
|
|
errStr := ""
|
|
if err != nil {
|
|
errStr = err.Error()
|
|
}
|
|
|
|
trace.RecordBackendTrace(trace.BackendTrace{
|
|
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
|
|
}
|
|
}
|
|
|
|
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
|