Added delete method

igolaizola · Aug 2, 2024 · ca4348f · ca4348f
1 parent dcab99e
commit ca4348f
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 19 deletions.
diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go
@@ -73,6 +73,7 @@ func newGenerateCommand() *ffcli.Command {
 	fs.StringVar(&cfg.Token, "token", "", "runway token")
 
 	fs.StringVar(&cfg.Model, "model", "gen2", "model to use (gen2 or gen3)")
+	fs.StringVar(&cfg.Folder, "folder", "", "runway folder to store assets (optional)")
 	fs.StringVar(&cfg.Image, "image", "", "source image")
 	fs.StringVar(&cfg.Text, "text", "", "source text")
 	fs.StringVar(&cfg.Output, "output", "", "output file (optional, if omitted it won't be saved)")
@@ -113,6 +114,7 @@ func newExtendCommand() *ffcli.Command {
 	fs.StringVar(&cfg.Output, "output", "", "output file (optional, if omitted it won't be saved)")
 	fs.IntVar(&cfg.N, "n", 1, "extend the video by this many times")
 	fs.StringVar(&cfg.Model, "model", "gen2", "model to use (gen2 or gen3)")
+	fs.StringVar(&cfg.Folder, "folder", "", "runway folder to store assets (optional)")
 	fs.BoolVar(&cfg.Interpolate, "interpolate", true, "interpolate frames (optional)")
 	fs.BoolVar(&cfg.Upscale, "upscale", false, "upscale frames (optional)")
 	fs.BoolVar(&cfg.Watermark, "watermark", false, "add watermark (optional)")

diff --git a/pkg/cmd/extend/extend.go b/pkg/cmd/extend/extend.go
@@ -24,6 +24,7 @@ type Config struct {
 	Output      string
 	N           int
 	Model       string
+	Folder      string
 	Interpolate bool
 	Upscale     bool
 	Watermark   bool
@@ -42,10 +43,11 @@ func Run(ctx context.Context, cfg *Config) error {
 		return fmt.Errorf("token is required")
 	}
 	client, err := runway.New(&runway.Config{
-		Token: cfg.Token,
-		Wait:  cfg.Wait,
-		Debug: cfg.Debug,
-		Proxy: cfg.Proxy,
+		Token:  cfg.Token,
+		Wait:   cfg.Wait,
+		Debug:  cfg.Debug,
+		Proxy:  cfg.Proxy,
+		Folder: cfg.Folder,
 	})
 	if err != nil {
 		return fmt.Errorf("vidai: couldn't create client: %w", err)
@@ -83,10 +85,18 @@ func Run(ctx context.Context, cfg *Config) error {
 		name := filepath.Base(img)
 
 		// Generate video
-		imageURL, err := client.Upload(ctx, name, b)
+		imageURL, assetID, err := client.Upload(ctx, name, b)
 		if err != nil {
 			return fmt.Errorf("vidai: couldn't upload image: %w", err)
 		}
+		defer func() {
+			// Delete asset
+			deleteCTX, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+			if err := client.Delete(deleteCTX, assetID); err != nil {
+				log.Println(fmt.Errorf("vidai: couldn't delete asset: %w", err))
+			}
+		}()
 		gen, err := client.Generate(ctx, &runway.GenerateRequest{
 			Model:       cfg.Model,
 			AssetURL:    imageURL,

diff --git a/pkg/cmd/generate/generate.go b/pkg/cmd/generate/generate.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
 	"os"
 	"path/filepath"
 	"strings"
@@ -20,6 +21,7 @@ type Config struct {
 
 	Output      string
 	Model       string
+	Folder      string
 	Image       string
 	Text        string
 	Extend      int
@@ -40,10 +42,11 @@ func Run(ctx context.Context, cfg *Config) error {
 		return fmt.Errorf("token is required")
 	}
 	client, err := runway.New(&runway.Config{
-		Token: cfg.Token,
-		Wait:  cfg.Wait,
-		Debug: cfg.Debug,
-		Proxy: cfg.Proxy,
+		Token:  cfg.Token,
+		Wait:   cfg.Wait,
+		Debug:  cfg.Debug,
+		Proxy:  cfg.Proxy,
+		Folder: cfg.Folder,
 	})
 	if err != nil {
 		return fmt.Errorf("vidai: couldn't create client: %w", err)
@@ -58,10 +61,19 @@ func Run(ctx context.Context, cfg *Config) error {
 		}
 		fileName = filepath.Base(cfg.Image)
 
-		imageURL, err = client.Upload(ctx, fileName, b)
+		var assetID string
+		imageURL, assetID, err = client.Upload(ctx, fileName, b)
 		if err != nil {
 			return fmt.Errorf("vidai: couldn't upload image: %w", err)
 		}
+		defer func() {
+			// Delete asset
+			deleteCTX, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+			if err := client.Delete(deleteCTX, assetID); err != nil {
+				log.Println(fmt.Errorf("vidai: couldn't delete asset: %w", err))
+			}
+		}()
 	}
 	gen, err := client.Generate(ctx, &runway.GenerateRequest{
 		Model:       cfg.Model,

diff --git a/pkg/runway/runway.go b/pkg/runway/runway.go
@@ -102,7 +102,7 @@ type datasetResponse struct {
 	} `json:"dataset"`
 }
 
-func (c *Client) Upload(ctx context.Context, name string, data []byte) (string, error) {
+func (c *Client) Upload(ctx context.Context, name string, data []byte) (string, string, error) {
 	ext := strings.TrimPrefix(".", filepath.Ext(name))
 	file := &uploadFile{
 		data:      data,
@@ -127,16 +127,16 @@ func (c *Client) Upload(ctx context.Context, name string, data []byte) (string,
 		}
 		var uploadResp uploadResponse
 		if _, err := c.do(ctx, "POST", "uploads", uploadReq, &uploadResp); err != nil {
-			return "", fmt.Errorf("runway: couldn't obtain upload url: %w", err)
+			return "", "", fmt.Errorf("runway: couldn't obtain upload url: %w", err)
 		}
 		if len(uploadResp.UploadURLs) == 0 {
-			return "", fmt.Errorf("runway: no upload urls returned")
+			return "", "", fmt.Errorf("runway: no upload urls returned")
 		}
 
 		// Upload file
 		uploadURL := uploadResp.UploadURLs[0]
 		if _, err := c.do(ctx, "PUT", uploadURL, file, nil); err != nil {
-			return "", fmt.Errorf("runway: couldn't upload file: %w", err)
+			return "", "", fmt.Errorf("runway: couldn't upload file: %w", err)
 		}
 
 		// Complete upload
@@ -154,12 +154,12 @@ func (c *Client) Upload(ctx context.Context, name string, data []byte) (string,
 		}
 		var completeResp uploadCompleteResponse
 		if _, err := c.do(ctx, "POST", completeURL, completeReq, &completeResp); err != nil {
-			return "", fmt.Errorf("runway: couldn't complete upload: %w", err)
+			return "", "", fmt.Errorf("runway: couldn't complete upload: %w", err)
 		}
 
 		c.log("runway: upload complete %s", completeResp.URL)
 		if completeResp.URL == "" {
-			return "", fmt.Errorf("runway: empty image url for type %s", t)
+			return "", "", fmt.Errorf("runway: empty image url for type %s", t)
 		}
 		imageURL = completeResp.URL
 
@@ -190,13 +190,37 @@ func (c *Client) Upload(ctx context.Context, name string, data []byte) (string,
 	}
 	var datasetResp datasetResponse
 	if _, err := c.do(ctx, "POST", "datasets", datasetReq, &datasetResp); err != nil {
-		return "", fmt.Errorf("runway: couldn't create dataset: %w", err)
+		return "", "", fmt.Errorf("runway: couldn't create dataset: %w", err)
 	}
 	if datasetResp.Dataset.URL == "" || datasetResp.Dataset.ID == "" {
-		return "", fmt.Errorf("runway: empty dataset url or id")
+		return "", "", fmt.Errorf("runway: empty dataset url or id")
 	}
 
-	return imageURL, nil
+	return imageURL, datasetResp.Dataset.ID, nil
+}
+
+type deleteRequest struct {
+	AsTeamID int `json:"asTeamId"`
+}
+
+type deleteResponse struct {
+	Success bool `json:"success"`
+}
+
+func (c *Client) Delete(ctx context.Context, assetID string) error {
+	path := fmt.Sprintf("assets/%s", assetID)
+	req := &deleteRequest{
+		AsTeamID: c.teamID,
+	}
+	var resp deleteResponse
+	b, err := c.do(ctx, "DELETE", path, req, &resp)
+	if err != nil {
+		return fmt.Errorf("runway: couldn't delete asset %s: %w", assetID, err)
+	}
+	if !resp.Success {
+		return fmt.Errorf("runway: couldn't delete asset %s: %s", assetID, string(b))
+	}
+	return nil
 }
 
 type createGen2TaskRequest struct {