raannakasturi commited on
Commit
10266ab
·
1 Parent(s): 227f9f9

Refactor image generation to use a client with retry logic and multiple providers

Browse files
Files changed (1) hide show
  1. image.py +22 -13
image.py CHANGED
@@ -5,6 +5,9 @@ import re
5
  import requests
6
  from urllib.parse import quote
7
  from PIL import Image
 
 
 
8
 
9
  def extract_summary(text):
10
  text = text.replace("#", "").strip().lower()
@@ -21,21 +24,27 @@ def fix_base64_padding(data):
21
 
22
  def generate_image(title, summary):
23
  try:
24
- extracted_summary = extract_summary(summary)
25
  negative="low quality, blurry, pixelated, bad anatomy, bad hands, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, poorly rendered hands, bad face, fused face, cloned face, worst face, three crus, extra crus, fused crus, worst feet, three feet, fused feet, fused thigh, three thigh, extra thigh, worst thigh, missing fingers, extra fingers, ugly fingers, long fingers, bad composition, horn, extra eyes, huge eyes, 2girl, amputation, disconnected limbs, cartoon, cg, 3d, unreal, animate, cgi, render, artwork, illustration, 3d render, cinema 4d, artstation, octane render, mutated body parts, painting, oil painting, 2d, sketch, bad photography, bad photo, deviant art, aberrations, abstract, anime, black and white, collapsed, conjoined, creative, drawing, extra windows, harsh lighting, jpeg artifacts, low saturation, monochrome, multiple levels, overexposed, oversaturated, photoshop, rotten, surreal, twisted, UI, underexposed, unnatural, unreal engine, unrealistic, video game, deformed body features",
26
- image_api_endpoint = "image.pollinations.ai"
27
- params = f"negative={quote(str(negative))}&seed=2342342340&width=1024&height=576&nologo=True&private=False&model=turbo&enhance=True"
28
- data_header = {"Content-Type": "application/json"}
29
- prompt = quote(f"[[(({title.strip()}))]]: {extracted_summary.strip()}")
30
- image_url = f"https://{image_api_endpoint}/prompt/{prompt}?{params}"
31
- response = requests.get(image_url, headers=data_header)
32
- if response.status_code != 200:
33
- print(f"Error generating image: {response.text}")
34
- return None
35
- img_data = response.content
 
 
 
 
 
 
 
 
36
  if img_data:
37
- data = base64.b64encode(img_data).decode("utf-8")
38
- return f"data:image/png;base64,{data}"
39
  return None
40
  except Exception as e:
41
  print(f"An error occurred during image generation: {e}")
 
5
  import requests
6
  from urllib.parse import quote
7
  from PIL import Image
8
+ from g4f.client import Client
9
+ from g4f.Provider import RetryProvider, Airforce, Blackbox, BlackboxCreateAgent, PollinationsAI
10
+
11
 
12
  def extract_summary(text):
13
  text = text.replace("#", "").strip().lower()
 
24
 
25
  def generate_image(title, summary):
26
  try:
 
27
  negative="low quality, blurry, pixelated, bad anatomy, bad hands, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, poorly rendered hands, bad face, fused face, cloned face, worst face, three crus, extra crus, fused crus, worst feet, three feet, fused feet, fused thigh, three thigh, extra thigh, worst thigh, missing fingers, extra fingers, ugly fingers, long fingers, bad composition, horn, extra eyes, huge eyes, 2girl, amputation, disconnected limbs, cartoon, cg, 3d, unreal, animate, cgi, render, artwork, illustration, 3d render, cinema 4d, artstation, octane render, mutated body parts, painting, oil painting, 2d, sketch, bad photography, bad photo, deviant art, aberrations, abstract, anime, black and white, collapsed, conjoined, creative, drawing, extra windows, harsh lighting, jpeg artifacts, low saturation, monochrome, multiple levels, overexposed, oversaturated, photoshop, rotten, surreal, twisted, UI, underexposed, unnatural, unreal engine, unrealistic, video game, deformed body features",
28
+ extracted_summary = extract_summary(summary)
29
+ prompt = quote(f"[[IMAGE GENERATED SHOULD BE SAFE FOR WORK (SFW). NO NUDES OR REVEALING IMAGES]] [[(({title.strip()}))]]: {extracted_summary.strip()}")
30
+ client = Client(
31
+ image_provider=RetryProvider(
32
+ providers=[Airforce, Blackbox, BlackboxCreateAgent, PollinationsAI],
33
+ shuffle=True,
34
+ single_provider_retry=True,
35
+ max_retries=3,
36
+ )
37
+ )
38
+ img_data = client.images.generate(
39
+ model="flux",
40
+ prompt=prompt,
41
+ negative_prompt=negative,
42
+ response_format="b64_json",
43
+ width=1024,
44
+ height=576,
45
+ ).data[0].b64_json
46
  if img_data:
47
+ return f"data:image/png;base64,{img_data}"
 
48
  return None
49
  except Exception as e:
50
  print(f"An error occurred during image generation: {e}")