Spaces:
Sleeping
Sleeping
File size: 4,731 Bytes
9828167 f275e0a 9828167 7b3babf 9828167 7b3babf f275e0a cf6d25f 7b3babf 9828167 7b3babf 9828167 7b3babf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
from openai import OpenAI
from pydantic import BaseModel
import replicate
from dotenv import load_dotenv
import boto3
import os
from uuid import uuid4
import requests
from meta_data import meta_data_helper_function
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Optional
load_dotenv()
gpt_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
def generate_prompt(category, keyword, subkeyword, aspect_ratio):
system_prompt = """You are professional prompt generator for the images which is to be used in the ads for the performance marketing.
User will give you category, keyword and subkeyword as input. You will generate the prompt for the given category, keyword and subkeyword.
Make sure the prompt which you would generate for generating the images should relevant to information provided. The images should be stocky.
Strictly, their should be no text overlays mentioned in the prompt. I don't want any text on images so don't mention any text overlays in the prompt.
For generating the images, google/imagen-4 model will be used so, craft the prompt for that model."""
user_prompt = f"Generate me the stocky images for the given category: {category} and keyword: {keyword} and subkeyword: {subkeyword}. The aspect ratio of the images should be: {aspect_ratio}."
class GeneratedPrompt(BaseModel):
prompt: str
messages = [
{
"role": "system",
"content": [
{
"type": "text",
"text": system_prompt
},
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": user_prompt
},
]
}
]
completion = gpt_client.beta.chat.completions.parse(
model="gpt-4o",
messages=messages,
response_format=GeneratedPrompt,
)
prompt_response = completion.choices[0].message
if prompt_response.parsed:
return prompt_response.parsed
else:
return prompt_response.refusal
def generate_images(prompt, aspect_ratio):
replicate_client = replicate.Client(api_token=os.getenv("REPLICATE_API_TOKEN"))
output = replicate_client.run(
"google/imagen-4-ultra",
input={
"prompt": prompt,
"aspect_ratio": aspect_ratio
}
)
print(output)
if isinstance(output, list) and output:
first = output[0]
url = getattr(first, "url", str(first))
urls = [url]
elif isinstance(output, str):
urls = [output]
elif hasattr(output, "url"):
urls = [getattr(output, "url")]
return urls[0]
def fetch_image_bytes(url):
r = requests.get(url, timeout=60)
r.raise_for_status()
return r.content
def init_s3():
return boto3.client(
"s3",
endpoint_url=os.getenv("R2_ENDPOINT"),
aws_access_key_id=os.getenv("R2_ACCESS_KEY"),
aws_secret_access_key=os.getenv("R2_SECRET_KEY"),
region_name="auto",
)
def upload_to_r2(image_bytes):
s3_client = init_s3()
filename = f"{uuid4().hex}.png"
file_key = f"infinityverse/{filename}"
s3_client.put_object(
Bucket=os.getenv("R2_BUCKET_NAME"),
Key=file_key,
Body=image_bytes,
ContentType="image/png",
)
r2_url = f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
return r2_url
def _process_one(aspect_ratio: str, category: str, keyword: str, subkeyword: str) -> str:
prompt_obj = generate_prompt(category, keyword, subkeyword, aspect_ratio)
prompt_text = getattr(prompt_obj, "prompt", prompt_obj)
image_urls = generate_images(prompt_text, aspect_ratio)
image_bytes = fetch_image_bytes(image_urls)
image_with_metadata = meta_data_helper_function(image_bytes)
r2_url = upload_to_r2(image_with_metadata)
return r2_url
def get_images(category: str, keyword: str, subkeyword: str) -> List[str]:
max_workers = 5
aspect_ratios = ["1:1", "1:1", "16:9", "16:9", "16:9"]
urls: List[Optional[str]] = [None] * len(aspect_ratios)
workers = max(1, min(max_workers, len(aspect_ratios)))
with ThreadPoolExecutor(max_workers=workers) as pool:
futures = {
pool.submit(_process_one, ar, category, keyword, subkeyword): idx
for idx, ar in enumerate(aspect_ratios)
}
for fut in as_completed(futures):
idx = futures[fut]
try:
urls[idx] = fut.result()
except Exception as e:
print(e)
urls[idx] = None
return [u for u in urls if u is not None] |