Skip to content
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 72 additions & 21 deletions core/gallery/models.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"os"
"path/filepath"
"slices"
"strings"

"dario.cat/mergo"
Expand Down Expand Up @@ -285,21 +286,32 @@ func GetLocalModelConfiguration(basePath string, name string) (*ModelConfig, err
return ReadConfigFile[ModelConfig](galleryFile)
}

func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
additionalFiles := []string{}
func listModelFiles(systemState *system.SystemState, name string) ([]string, error) {

configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))
if err := utils.VerifyPath(configFile, systemState.Model.ModelsPath); err != nil {
return fmt.Errorf("failed to verify path %s: %w", configFile, err)
return nil, fmt.Errorf("failed to verify path %s: %w", configFile, err)
}

// os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths.
name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name))
if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil {
return nil, fmt.Errorf("failed to verify path %s: %w", galleryFile, err)
}

additionalFiles := []string{}
allFiles := []string{}

// Galleryname is the name of the model in this case
dat, err := os.ReadFile(configFile)
if err == nil {
modelConfig := &config.ModelConfig{}

err = yaml.Unmarshal(dat, &modelConfig)
if err != nil {
return err
return nil, err
}
if modelConfig.Model != "" {
additionalFiles = append(additionalFiles, modelConfig.ModelFileName())
Expand All @@ -310,26 +322,15 @@ func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
}
}

// os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths.
name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name))
if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil {
return fmt.Errorf("failed to verify path %s: %w", galleryFile, err)
}

var filesToRemove []string

// Delete all the files associated to the model
// read the model config
galleryconfig, err := ReadConfigFile[ModelConfig](galleryFile)
if err == nil && galleryconfig != nil {
for _, f := range galleryconfig.Files {
fullPath := filepath.Join(systemState.Model.ModelsPath, f.Filename)
if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
return fmt.Errorf("failed to verify path %s: %w", fullPath, err)
return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
}
filesToRemove = append(filesToRemove, fullPath)
allFiles = append(allFiles, fullPath)
}
} else {
log.Error().Err(err).Msgf("failed to read gallery file %s", configFile)
Expand All @@ -338,18 +339,68 @@ func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
for _, f := range additionalFiles {
fullPath := filepath.Join(filepath.Join(systemState.Model.ModelsPath, f))
if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
return fmt.Errorf("failed to verify path %s: %w", fullPath, err)
return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
}
filesToRemove = append(filesToRemove, fullPath)
allFiles = append(allFiles, fullPath)
}

filesToRemove = append(filesToRemove, galleryFile)
allFiles = append(allFiles, galleryFile)

// skip duplicates
filesToRemove = utils.Unique(filesToRemove)
allFiles = utils.Unique(allFiles)

return allFiles, nil
}

func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))

filesToRemove, err := listModelFiles(systemState, name)
if err != nil {
return err
}

allOtherFiles := []string{}
// Get all files of all other models
fi, err := os.ReadDir(systemState.Model.ModelsPath)
if err != nil {
return err
}
for _, f := range fi {
if f.IsDir() {
continue
}
if strings.HasPrefix(f.Name(), "._gallery_") {
continue
}
if !strings.HasSuffix(f.Name(), ".yaml") || !strings.HasSuffix(f.Name(), ".yml") {
continue
}
if f.Name() == name {
continue
}

name := strings.TrimSuffix(f.Name(), ".yaml")
name = strings.TrimSuffix(name, ".yml")

log.Debug().Msgf("Checking file %s", f.Name())
files, err := listModelFiles(systemState, name)
if err != nil {
log.Debug().Err(err).Msgf("failed to list files for model %s", f.Name())
continue
}
allOtherFiles = append(allOtherFiles, files...)
}

log.Debug().Msgf("Files to remove: %+v", filesToRemove)
log.Debug().Msgf("All other files: %+v", allOtherFiles)

// Removing files
for _, f := range filesToRemove {
if slices.Contains(allOtherFiles, f) {
log.Debug().Msgf("Skipping file %s because it is part of another model", f)
continue
}
if e := os.Remove(f); e != nil {
log.Error().Err(e).Msgf("failed to remove file %s", f)
}
Expand Down
Loading