feat(imagegen): add support for quite flag for image generation

- feat: add support for quite flag during image generation
- refactor: move imaggen logic to imagegen package
- exit with exit code 1 when image generation fail

Author: Aman Shaw

helper.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	url_package "net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -17,7 +16,6 @@ import (
 	"github.com/aandrew-me/tgpt/v2/providers"
 	"github.com/aandrew-me/tgpt/v2/providers/gemini"
 	"github.com/aandrew-me/tgpt/v2/structs"
-	"github.com/aandrew-me/tgpt/v2/utils"
 	http "github.com/bogdanfinn/fhttp"
 
 	"github.com/olekukonko/ts"
@@ -547,90 +545,6 @@ func handleStatus400(resp *http.Response) {
 // 	}
 // }
 
-
-func generateImagePollinations(prompt string) {
-	bold.Println("Generating image with pollinations.ai...")
-
-	client, err := client.NewClient()
-	if err != nil {
-		fmt.Fprintln(os.Stderr, err)
-		os.Exit(1)
-	}
-
-	full_prompt := url_package.QueryEscape(prompt);
-
-	randId := utils.RandomString(20)
-	filename := randId + ".jpg"
-
-	model := "flux"
-
-	if *apiModel != "" {
-		model = *apiModel
-	}
-
-	fmt.Println()
-
-	link := fmt.Sprintf("https://image.pollinations.ai/prompt/%v", full_prompt)
-
-	params := url_package.Values{}
-
-	seed := utils.GenerateRandomNumber(5)
-
-	params.Add("model", model)
-	params.Add("width", "1024")
-	params.Add("height", "1024")
-	params.Add("nologo", "true")
-	params.Add("safe", "false")
-	params.Add("nsfw", "true")
-	params.Add("isChild", "false")
-	params.Add("seed", seed)
-
-	urlObj, err := url_package.Parse(link)
-	if err != nil {
-		fmt.Println("Error parsing URL:", err)
-		return
-	}
-
-	urlObj.RawQuery = params.Encode()
-
-
-	req, _ := http.NewRequest("GET", urlObj.String(), nil)
-
-	res, err := client.Do(req)
-
-	if err != nil {
-		fmt.Fprint(os.Stderr, err);
-	}
-
-	defer res.Body.Close()
-
-
-	if res.StatusCode == http.StatusOK {
-		file, err := os.Create(filename)
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "Error: %v", err)
-
-			return
-		}
-		defer file.Close()
-	
-		// Copy the response body (image data) to the file
-		_, err = io.Copy(file, res.Body)
-		if err != nil {
-			fmt.Fprintf(os.Stderr,"Error: %v", err)
-			
-			return
-		}
-
-		fmt.Printf("Saved image as %v\n", filename)
-	} else {
-		body, _ := io.ReadAll(res.Body)
-		responseText := string(body)
-
-		fmt.Fprintf(os.Stderr,"Some error has occurred. Try again (perhaps with a different model).\nError: %v", responseText)
-	}
-}
-
 func downloadImage(url string, destDir string) error {
 	client, err := client.NewClient()
 	if err != nil {
@@ -799,13 +713,4 @@ func makeRequestAndGetData(input string, params structs.Params, extraOptions str
 	}
 
 	return ""
-}
-
-func generateImg(prompt string, provider string) {	
-	if provider == "pollinations" || provider == "" {
-		generateImagePollinations(prompt)
-
-	} else {
-		fmt.Fprintln(os.Stderr, "Such a provider doesn't exist")
-	}
-}
+}

imagegen/imagegen.go

@@ -0,0 +1,111 @@
+package imagegen
+
+import (
+	"fmt"
+	"io"
+	"os"
+
+	url_package "net/url"
+
+	http "github.com/bogdanfinn/fhttp"
+
+	"github.com/aandrew-me/tgpt/v2/client"
+	"github.com/aandrew-me/tgpt/v2/structs"
+	"github.com/aandrew-me/tgpt/v2/utils"
+	"github.com/fatih/color"
+)
+
+var bold = color.New(color.Bold)
+
+func GenerateImg(prompt string, params structs.Params, isQuite bool) {
+	if params.Provider == "pollinations" || params.Provider == "" {
+		if !isQuite {
+			bold.Println("Generating image with pollinations.ai...")
+		}
+		filename := generateImagePollinations(prompt, params)
+		if !isQuite {
+			fmt.Printf("Saved image as %v\n", filename)
+		}
+
+	} else {
+		fmt.Fprintln(os.Stderr, "Such a provider doesn't exist")
+	}
+}
+
+func generateImagePollinations(prompt string, params structs.Params) string {
+
+	client, err := client.NewClient()
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+
+	full_prompt := url_package.QueryEscape(prompt)
+
+	randId := utils.RandomString(20)
+	filename := randId + ".jpg"
+
+	model := "flux"
+
+	if params.ApiModel != "" {
+		model = params.ApiModel
+	}
+
+	link := fmt.Sprintf("https://image.pollinations.ai/prompt/%v", full_prompt)
+
+	queryParams := url_package.Values{}
+
+	seed := utils.GenerateRandomNumber(5)
+
+	queryParams.Add("model", model)
+	queryParams.Add("width", "1024")
+	queryParams.Add("height", "1024")
+	queryParams.Add("nologo", "true")
+	queryParams.Add("safe", "false")
+	queryParams.Add("nsfw", "true")
+	queryParams.Add("isChild", "false")
+	queryParams.Add("seed", seed)
+
+	urlObj, err := url_package.Parse(link)
+	if err != nil {
+		fmt.Println("Error parsing URL:", err)
+		os.Exit(1)
+	}
+
+	urlObj.RawQuery = queryParams.Encode()
+
+	req, _ := http.NewRequest("GET", urlObj.String(), nil)
+
+	res, err := client.Do(req)
+
+	if err != nil {
+		fmt.Fprint(os.Stderr, err)
+	}
+
+	defer res.Body.Close()
+
+	if res.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(res.Body)
+		responseText := string(body)
+
+		fmt.Fprintf(os.Stderr, "Some error has occurred. Try again (perhaps with a different model).\nError: %v", responseText)
+
+	}
+
+	file, err := os.Create(filename)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v", err)
+		os.Exit(1)
+	}
+	defer file.Close()
+
+	// Copy the response body (image data) to the file
+	_, err = io.Copy(file, res.Body)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v", err)
+		os.Exit(1)
+	}
+
+	return filename
+
+}

main.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/aandrew-me/tgpt/v2/imagegen"
 	"github.com/aandrew-me/tgpt/v2/structs"
 	"github.com/aandrew-me/tgpt/v2/utils"
 	"github.com/atotto/clipboard"
@@ -168,6 +169,38 @@ func main() {
 			fmt.Println("tgpt", localVersion)
 		case *isChangelog:
 			getVersionHistory()
+		case *isImage:
+			params := structs.Params{
+				ApiKey:       *apiKey,
+				ApiModel:     *apiModel,
+				Provider:     *provider,
+				Max_length:   *max_length,
+				Temperature:  *temperature,
+				Top_p:        *top_p,
+				Preprompt:    *preprompt,
+				Url:          *url,
+				PrevMessages: "",
+				ThreadID:     "",
+			}
+
+			if len(prompt) > 1 {
+				trimmedPrompt := strings.TrimSpace(prompt)
+				if len(trimmedPrompt) < 1 {
+					fmt.Fprintln(os.Stderr, "You need to provide some text")
+					fmt.Fprintln(os.Stderr, `Example: tgpt -img "cat"`)
+					os.Exit(1)
+				}
+
+				imagegen.GenerateImg(trimmedPrompt, params, *isQuiet)
+
+			} else {
+				formattedInput := getFormattedInputStdin()
+				if !*isQuiet {
+					fmt.Println()
+				}
+
+				imagegen.GenerateImg(formattedInput, params, *isQuiet)
+			}
 		case *isWhole:
 			if len(prompt) > 1 {
 				trimmedPrompt := strings.TrimSpace(prompt)
@@ -337,22 +370,6 @@ func main() {
 
 			}
 
-		case *isImage:
-			if len(prompt) > 1 {
-				trimmedPrompt := strings.TrimSpace(prompt)
-				if len(trimmedPrompt) < 1 {
-					fmt.Fprintln(os.Stderr, "You need to provide some text")
-					fmt.Fprintln(os.Stderr, `Example: tgpt -img "cat"`)
-					os.Exit(1)
-				}
-
-				generateImg(trimmedPrompt, *provider)
-			} else {
-				formattedInput := getFormattedInputStdin()
-				fmt.Println()
-
-				generateImg(formattedInput, *provider)
-			}
 		case *isHelp:
 			showHelpMessage()
 		default:
@@ -549,7 +566,7 @@ func showHelpMessage() {
 
 	bold.Println("\nProvider: isou")
 	fmt.Println("Free provider with web search")
-	
+
 	bold.Println("\nProvider: koboldai")
 	fmt.Println("Uses koboldcpp/HF_SPACE_Tiefighter-13B only, answers from novels")