update to download metadata and preview image too

Author: yansigit

README.md

@@ -11,4 +11,9 @@ A Simple and easy tool to download models from [Civitai](https://civitai.com/)
 3. Run the binary and download a model you like. That's it!
 ```bash
 ./civitai-downloader <sub_dir> <model_url|AIR>
-```
+```
+
+## TODO
+
+- [v] Add support for downloading models from [Civitai](https://civitai.com/)
+- [v] Download related image and metadata along with the model.

downloader/downloader.go

@@ -1,6 +1,7 @@
 package downloader
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
@@ -11,20 +12,55 @@ import (
 	"github.com/schollz/progressbar/v3"
 )
 
-func DownloadFile(outputPath string, url string) error {
-	dir := filepath.Dir(outputPath)
-	if err := os.MkdirAll(dir, 0755); err != nil {
-		return fmt.Errorf("failed to create directories: %v", err)
+// ModelVersion represents the model version information from the Civitai API
+type ModelVersion struct {
+	ID          int64   `json:"id"`
+	ModelID     int64   `json:"modelId"`
+	Name        string  `json:"name"`
+	Files       []File  `json:"files"`
+	Images      []Image `json:"images"`
+	Description string  `json:"description"`
+	DownloadURL string  `json:"downloadUrl"`
+}
+
+// File represents a file associated with the model version
+type File struct {
+	ID          int64   `json:"id"`
+	SizeKB      float64 `json:"sizeKB"`
+	Name        string  `json:"name"`
+	Type        string  `json:"type"`
+	DownloadURL string  `json:"downloadUrl"`
+}
+
+// Image represents an image associated with the model version
+type Image struct {
+	URL    string `json:"url"`
+	Type   string `json:"type"`
+	Width  int    `json:"width"`
+	Height int    `json:"height"`
+}
+
+const (
+	APIModelVersions = "https://civitai.com/api/v1/model-versions/"
+)
+
+// DownloadFile downloads a single file from the given URL to the specified path and returns the saved filename
+func DownloadFile(outputPath, url, modelVersionId string) (string, error) {
+	outputDir := filepath.Dir(outputPath)
+	if _, err := os.Stat(outputDir); os.IsNotExist(err) {
+		if err := os.MkdirAll(outputDir, 0755); err != nil {
+			return "", fmt.Errorf("failed to create directory: %w", err)
+		}
 	}
 
 	resp, err := http.Get(url)
 	if err != nil {
-		return fmt.Errorf("failed to fetch %s: %v", url, err)
+		return "", fmt.Errorf("failed to download file: %w", err)
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("HTTP error for %s: %v", url, resp.StatusCode)
+		return "", fmt.Errorf("unexpected status code: %d", resp.StatusCode)
 	}
 
 	header := resp.Header.Get("content-disposition")
@@ -37,7 +73,7 @@ func DownloadFile(outputPath string, url string) error {
 
 	file, err := os.Create(outputPath)
 	if err != nil {
-		return fmt.Errorf("failed to create file %s: %v", outputPath, err)
+		return "", fmt.Errorf("failed to create file: %w", err)
 	}
 	defer file.Close()
 
@@ -57,5 +93,95 @@ func DownloadFile(outputPath string, url string) error {
 	)
 
 	_, err = io.Copy(io.MultiWriter(file, bar), resp.Body)
-	return err
+	if err != nil {
+		return "", fmt.Errorf("failed to write to file: %w", err)
+	}
+
+	return outputPath, nil
+}
+
+// DownloadAll downloads all available files for a given model ID
+func DownloadAll(modelType, baseModelPath, modelID string) error {
+	modelURL := fmt.Sprintf("%s%s", APIModelVersions, modelID)
+	resp, err := http.Get(modelURL)
+	if err != nil {
+		return fmt.Errorf("failed to fetch model version: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("failed to fetch model version. Status code: %d", resp.StatusCode)
+	}
+
+	var modelVersion ModelVersion
+	err = json.NewDecoder(resp.Body).Decode(&modelVersion)
+	if err != nil {
+		return fmt.Errorf("failed to decode model version response: %w", err)
+	}
+
+	// Create a subdirectory based on modelType
+	dir := filepath.Join(baseModelPath, modelType)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create model type directory: %w", err)
+	}
+
+	outputPath, err := DownloadFile(filepath.Join(dir, filepath.Base(baseModelPath)), modelVersion.DownloadURL, modelID)
+	if err != nil {
+		return fmt.Errorf("failed to download model file: %w", err)
+	}
+
+	baseName := strings.TrimSuffix(filepath.Base(outputPath), filepath.Ext(outputPath))
+
+	for _, image := range modelVersion.Images {
+		if image.Type == "image" {
+			imgPath := fmt.Sprintf("%s.preview.png", filepath.Join(filepath.Dir(outputPath), baseName))
+			if _, err := DownloadFile(imgPath, image.URL, modelID); err != nil {
+				return fmt.Errorf("failed to download image: %w", err)
+			}
+		} else if image.Type == "video" {
+			imgPath := fmt.Sprintf("%s.preview.mp4", filepath.Join(filepath.Dir(outputPath), baseName))
+			if _, err := DownloadFile(imgPath, image.URL, modelID); err != nil {
+				return fmt.Errorf("failed to download image: %w", err)
+			}
+		}
+	}
+
+	metadataPath := fmt.Sprintf("%s.civitai.info", filepath.Join(filepath.Dir(outputPath), baseName))
+	metadata, err := json.MarshalIndent(modelVersion, "", "  ")
+	if err != nil {
+		return fmt.Errorf("failed to marshal metadata: %w", err)
+	}
+	if err := os.WriteFile(metadataPath, metadata, 0644); err != nil {
+		return fmt.Errorf("failed to save metadata: %w", err)
+	}
+
+	if modelVersion.Description != "" {
+		descPath := fmt.Sprintf("%s.description.txt", filepath.Join(filepath.Dir(outputPath), baseName))
+		if err := os.WriteFile(descPath, []byte(modelVersion.Description), 0644); err != nil {
+			return fmt.Errorf("failed to save description: %w", err)
+		}
+	}
+
+	fmt.Printf("Successfully downloaded model files to %s\n", filepath.Dir(outputPath))
+	return nil
+}
+
+// GetModelID retrieves the model ID from the URL
+func GetModelID(url string) (string, error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return "", fmt.Errorf("failed to fetch model page: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+	}
+
+	parts := url[len("https://civitai.com/models/"):]
+	if len(parts) == 0 {
+		return "", fmt.Errorf("invalid URL format")
+	}
+	modelID := parts
+	return modelID, nil
 }

main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -44,22 +45,25 @@ func main() {
 		fmt.Println("Base model path not specified, using current directory")
 	}
 
-	var downloadURL string
+	var modelVersionId string
 	if strings.HasPrefix(modelIdentifier, "urn:air:") {
 		parts := strings.Split(modelIdentifier, ":")
 		if modelType == "" {
 			modelType = parts[3]
 		}
 		modelInfo := strings.Split(parts[len(parts)-1], "@")
-		version := modelInfo[1]
-		downloadURL = fmt.Sprintf("https://civitai.com/api/download/models/%s?token=%s", version, token)
+		modelVersionId = modelInfo[1]
 	} else {
-		downloadURL = modelIdentifier + "?token=" + token
+		re := regexp.MustCompile(`https://civitai.com/models/(\d+)`)
+		matches := re.FindStringSubmatch(modelIdentifier)
+		if len(matches) == 2 {
+			modelVersionId = matches[1]
+		}
 	}
-
+	fmt.Println("Model version ID:", modelVersionId)
 	outputPath := filepath.Join(baseModelPath, modelType, fmt.Sprintf("temp-%d.safetensors", time.Now().UnixNano()))
 
-	err = downloader.DownloadFile(outputPath, downloadURL)
+	err = downloader.DownloadAll(modelType, baseModelPath, modelVersionId)
 	if err != nil {
 		fmt.Printf("Error downloading %s: %v\n", modelType, err)
 		return