diff --git a/internal/client/client.go b/internal/client/client.go index 38478e2..3941e26 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -3,6 +3,7 @@ package client import ( "context" "fmt" + "os" "github.com/replicate/cli/internal" "github.com/replicate/cli/internal/config" @@ -10,14 +11,16 @@ import ( ) 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`") } @@ -25,8 +28,7 @@ func NewClient(opts ...replicate.ClientOption) (*replicate.Client, error) { } 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 +} diff --git a/internal/util/upload.go b/internal/util/upload.go index 6e4492e..67d67f1 100644 --- a/internal/util/upload.go +++ b/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)