Fix authentication for file uploads

internal/client/client.go

@@ -3,30 +3,32 @@ package client
 import (
 	"context"
 	"fmt"
+	"os"
 
 	"github.com/replicate/cli/internal"
 	"github.com/replicate/cli/internal/config"
 	"github.com/replicate/replicate-go"
 )
 
 func NewClient(opts ...replicate.ClientOption) (*replicate.Client, error) {
-	token, err := config.GetAPIToken()
+	token, err := getToken()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get API token: %w", err)
 	}
 
+	baseURL := getBaseURL()
+
 	// Validate token when connecting to api.replicate.com.
 	// Alternate API hosts proxying Replicate may not require a token.
-	if token == "" && config.GetAPIBaseURL() == config.DefaultBaseURL {
+	if token == "" && baseURL == config.DefaultBaseURL {
 		return nil, fmt.Errorf("please authenticate with `replicate auth login`")
 	}
 
 	return NewClientWithAPIToken(token, opts...)
 }
 
 func NewClientWithAPIToken(token string, opts ...replicate.ClientOption) (*replicate.Client, error) {
-	baseURL := config.GetAPIBaseURL()
-
+	baseURL := getBaseURL()
 	userAgent := fmt.Sprintf("replicate-cli/%s", internal.Version())
 
 	opts = append([]replicate.ClientOption{
@@ -57,3 +59,19 @@ func VerifyToken(ctx context.Context, token string) (bool, error) {
 
 	return true, nil
 }
+
+func getToken() (string, error) {
+	token, exists := os.LookupEnv("REPLICATE_API_TOKEN")
+	if !exists {
+		return config.GetAPIToken()
+	}
+	return token, nil
+}
+
+func getBaseURL() string {
+	baseURL, exists := os.LookupEnv("REPLICATE_BASE_URL")
+	if !exists {
+		baseURL = config.GetAPIBaseURL()
+	}
+	return baseURL
+}

internal/util/upload.go

@@ -9,6 +9,8 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+
+	"github.com/replicate/cli/internal/config"
 )
 
 // uploadFile uploads a file to Replicate's experimental DreamBooth API and returns the URL
@@ -27,7 +29,16 @@ func UploadFile(ctx context.Context, path string) (string, error) {
 		return "", fmt.Errorf("failed to create upload request: %w", err)
 	}
 
-	request.Header.Set("Authorization", "Token "+os.Getenv("REPLICATE_API_TOKEN"))
+	token, exists := os.LookupEnv("REPLICATE_API_TOKEN")
+	if !exists {
+		var err error
+		token, err = config.GetAPIToken()
+		if err != nil {
+			return "", fmt.Errorf("failed to get API token: %w", err)
+		}
+	}
+
+	request.Header.Set("Authorization", fmt.Sprintf("Token %s", token))
 	resp, err := http.DefaultClient.Do(request)
 	if err != nil {
 		return "", fmt.Errorf("failed to get upload URL: %w", err)