diff --git a/backend/go/stablediffusion-ggml/gosd.cpp b/backend/go/stablediffusion-ggml/gosd.cpp index 12e355a42..589abba0e 100644 --- a/backend/go/stablediffusion-ggml/gosd.cpp +++ b/backend/go/stablediffusion-ggml/gosd.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "gosd.h" // #include "preprocessing.hpp" @@ -85,7 +86,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) { fflush(stderr); } -int load_model(char *model, char* options[], int threads, int diff) { +int load_model(char *model, char *model_path, char* options[], int threads, int diff) { fprintf (stderr, "Loading model!\n"); sd_set_log_callback(sd_log_cb, NULL); @@ -103,6 +104,8 @@ int load_model(char *model, char* options[], int threads, int diff) { char *vae_path = ""; char *scheduler = ""; char *sampler = ""; + char *lora_dir = model_path; + bool lora_dir_allocated = false; fprintf(stderr, "parsing options\n"); @@ -132,6 +135,20 @@ int load_model(char *model, char* options[], int threads, int diff) { if (!strcmp(optname, "sampler")) { sampler = optval; } + if (!strcmp(optname, "lora_dir")) { + // Path join with model dir + if (model_path && strlen(model_path) > 0) { + std::filesystem::path model_path_str(model_path); + std::filesystem::path lora_path(optval); + std::filesystem::path full_lora_path = model_path_str / lora_path; + lora_dir = strdup(full_lora_path.string().c_str()); + lora_dir_allocated = true; + fprintf(stderr, "Lora dir resolved to: %s\n", lora_dir); + } else { + lora_dir = optval; + fprintf(stderr, "No model path provided, using lora dir as-is: %s\n", lora_dir); + } + } } fprintf(stderr, "parsed options\n"); @@ -176,7 +193,7 @@ int load_model(char *model, char* options[], int threads, int diff) { ctx_params.vae_path = vae_path; ctx_params.taesd_path = ""; ctx_params.control_net_path = ""; - ctx_params.lora_model_dir = ""; + ctx_params.lora_model_dir = lora_dir; ctx_params.embedding_dir = ""; ctx_params.stacked_id_embed_dir = ""; ctx_params.vae_decode_only = false; @@ -189,12 +206,21 @@ int load_model(char *model, char* options[], int threads, int diff) { if (sd_ctx == NULL) { fprintf (stderr, "failed loading model (generic error)\n"); + // Clean up allocated memory + if (lora_dir_allocated && lora_dir) { + free(lora_dir); + } return 1; } fprintf (stderr, "Created context: OK\n"); sd_c = sd_ctx; + // Clean up allocated memory + if (lora_dir_allocated && lora_dir) { + free(lora_dir); + } + return 0; } diff --git a/backend/go/stablediffusion-ggml/gosd.go b/backend/go/stablediffusion-ggml/gosd.go index 43f01fee2..221ba4294 100644 --- a/backend/go/stablediffusion-ggml/gosd.go +++ b/backend/go/stablediffusion-ggml/gosd.go @@ -29,9 +29,14 @@ func (sd *SDGGML) Load(opts *pb.ModelOptions) error { sd.threads = int(opts.Threads) + modelPath := opts.ModelPath + modelFile := C.CString(opts.ModelFile) defer C.free(unsafe.Pointer(modelFile)) + modelPathC := C.CString(modelPath) + defer C.free(unsafe.Pointer(modelPathC)) + var options **C.char // prepare the options array to pass to C @@ -70,7 +75,7 @@ func (sd *SDGGML) Load(opts *pb.ModelOptions) error { sd.cfgScale = opts.CFGScale - ret := C.load_model(modelFile, options, C.int(opts.Threads), C.int(diffusionModel)) + ret := C.load_model(modelFile, modelPathC, options, C.int(opts.Threads), C.int(diffusionModel)) if ret != 0 { return fmt.Errorf("could not load model") } diff --git a/backend/go/stablediffusion-ggml/gosd.h b/backend/go/stablediffusion-ggml/gosd.h index 8208bd27a..45db2e450 100644 --- a/backend/go/stablediffusion-ggml/gosd.h +++ b/backend/go/stablediffusion-ggml/gosd.h @@ -1,7 +1,7 @@ #ifdef __cplusplus extern "C" { #endif -int load_model(char *model, char* options[], int threads, int diffusionModel); +int load_model(char *model, char *model_path, char* options[], int threads, int diffusionModel); int gen_image(char *text, char *negativeText, int width, int height, int steps, int seed, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char **ref_images, int ref_images_count); #ifdef __cplusplus }