From 26d33e589ed8e9f6d86b8f18ceec66fa2f52df33 Mon Sep 17 00:00:00 2001
From: Mattt Zmuda <mattt@replicate.com>
Date: Wed, 20 Dec 2023 16:25:11 -0800
Subject: [PATCH 1/2] Override stored token with REPLICATE_API_TOKEN
 environment variable

---
 internal/client/client.go | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

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
+}

From d038cc270bde1a3d61da12efcbad0f89e19ba664 Mon Sep 17 00:00:00 2001
From: Mattt Zmuda <mattt@replicate.com>
Date: Wed, 20 Dec 2023 16:32:43 -0800
Subject: [PATCH 2/2] Load token from config for upload endpoint

---
 internal/util/upload.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

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)