Spaces:
Running
Running
| import os, io, zipfile, replicate, time, logging, requests, streamlit as st, boto3, threading | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from typing import Dict, Any, List, Tuple, Optional, Union | |
| from uuid import uuid4 | |
| from urllib.parse import urlparse | |
| from functools import lru_cache | |
| import os, base64, logging | |
| from openai import OpenAI | |
| from helpers_function.helper_meta_data import meta_data_helper_function | |
| from database.operations import start_job, finish_job | |
| from database.connections import get_results_collection | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| def _encode_image_to_base64(image_path): | |
| try: | |
| with open(image_path, "rb") as f: | |
| return base64.b64encode(f.read()).decode("utf-8") | |
| except Exception: | |
| logger.exception(f"Failed to base64 encode image: {image_path}") | |
| return "" | |
| logger = logging.getLogger("imagegen_service") | |
| logging.basicConfig(level=logging.INFO) | |
| CANONICAL_STATIC_BASE = "https://adgen.static.lookfinity.com" | |
| REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN") | |
| MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4) | |
| REQUEST_TIMEOUT = 30 | |
| RETRY_ATTEMPTS = 3 | |
| def _public_base_url() -> str: | |
| base = (os.getenv("NEW_BASE") or "").strip().rstrip("/") | |
| return base or CANONICAL_STATIC_BASE | |
| def get_cost_per_image( | |
| model_key: str, | |
| width: int | None = None, | |
| height: int | None = None, | |
| ) -> float: | |
| cfg = MODEL_REGISTRY.get(model_key, {}) | |
| price = cfg.get("price") | |
| unit = cfg.get("price_unit") | |
| if not price: | |
| return 0.0 | |
| if unit == "per_image": | |
| return round(price, 4) | |
| return 0.0 | |
| MODEL_REGISTRY: Dict[str, Dict[str, Any]] = { | |
| "imagegen-4-ultra": { | |
| "id": "google/imagen-4-ultra", | |
| "price": 0.06, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "imagen-4": { | |
| "id": "google/imagen-4", | |
| "price": 0.04, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "nano-banana": { | |
| "id": "google/nano-banana", | |
| "price": 0.039, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "qwen": { | |
| "id": "qwen/qwen-image", | |
| "price": 0.025, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "seedream-3": { | |
| "id": "bytedance/seedream-3", | |
| "price": 0.03, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3","21:9"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "recraft-v3": { | |
| "id": "recraft-ai/recraft-v3", | |
| "price": 0.04, | |
| "price_unit": "per_image", | |
| "aspect_ratios": [ | |
| "1:1","4:3","3:4","3:2","2:3","16:9","9:16", | |
| "1:2","2:1","7:5","5:7","4:5","5:4","3:5","5:3" | |
| ], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "photon": { | |
| "id": "luma/photon", | |
| "price": 0.03, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","3:4","4:3","9:16","16:9","9:21","21:9"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "ideogram-v3-quality": { | |
| "id": "ideogram-ai/ideogram-v3-quality", | |
| "price": 0.09, | |
| "price_unit": "per_image", | |
| "aspect_ratios": [ | |
| "1:3","3:1","1:2","2:1","9:16","16:9","10:16", | |
| "16:10","2:3","3:2","3:4","4:3","4:5","5:4","1:1" | |
| ], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "ideogram-v3-turbo": { | |
| "id": "ideogram-ai/ideogram-v3-turbo", | |
| "price": 0.03, | |
| "price_unit": "per_image", | |
| "aspect_ratios": [ | |
| "1:3","3:1","1:2","2:1","9:16","16:9","10:16", | |
| "16:10","2:3","3:2","3:4","4:3","4:5","5:4","1:1" | |
| ], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "flux-2-max": { | |
| "id": "black-forest-labs/flux-2-max", | |
| "price": 0.04, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| "p-image": { | |
| "id": "prunaai/p-image", | |
| "price": 0.005, | |
| "price_unit": "per_image", | |
| "type": "custom", | |
| }, | |
| "z-image-turbo": { | |
| "id": "prunaai/z-image-turbo", | |
| "price": 0.01, | |
| "price_unit": "per_image", | |
| "type": "custom", | |
| }, | |
| "bytedance/seedream-4.5": { | |
| "id": "bytedance/seedream-4.5", | |
| "size": "2k", | |
| "price": 0.04, | |
| "price_unit": "per_image", | |
| "aspect_ratios": ["1:1","2:3","3:2","3:4","4:3","9:16","16:9","21:9"], | |
| "param_name": "aspect_ratio", | |
| }, | |
| } | |
| _thread_local = threading.local() | |
| def get_model_config(model_key: str) -> Optional[Dict[str, Any]]: | |
| return MODEL_REGISTRY.get(model_key) | |
| def _get_model_config_cached(model_key: str) -> Optional[Dict[str, Any]]: | |
| return MODEL_REGISTRY.get(model_key) | |
| def _s3(): | |
| if not hasattr(_thread_local, "s3"): | |
| needed = ["R2_ENDPOINT", "R2_ACCESS_KEY", "R2_SECRET_KEY", "R2_BUCKET_NAME"] | |
| if any(not os.getenv(k) for k in needed): | |
| _thread_local.s3 = None | |
| return None | |
| try: | |
| _thread_local.s3 = 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", | |
| ) | |
| except Exception as e: | |
| logger.error(f"S3 init failed: {e}") | |
| _thread_local.s3 = None | |
| return _thread_local.s3 | |
| def _upload_to_r2(image_bytes: bytes) -> Optional[str]: | |
| s3 = _s3() | |
| if not s3: | |
| return None | |
| for attempt in range(RETRY_ATTEMPTS): | |
| try: | |
| filename = f"{uuid4().hex}.png" | |
| key = f"adgenesis_image_text/creative_adgenesis/images/{filename}" | |
| s3.put_object( | |
| Bucket=os.getenv("R2_BUCKET_NAME"), | |
| Key=key, | |
| Body=image_bytes, | |
| ContentType="image/png", | |
| ) | |
| return f"{_public_base_url()}/{key}" | |
| except Exception as e: | |
| if attempt == RETRY_ATTEMPTS - 1: | |
| logger.error(f"R2 upload failed: {e}") | |
| return None | |
| time.sleep(2 ** attempt) | |
| return None | |
| def _generate_one(model_key: str,prompt: str, aspect_ratio: str,params: Dict[str, Any],) -> List[str]: | |
| if not REPLICATE_API_TOKEN: | |
| return [] | |
| cfg = _get_model_config_cached(model_key) | |
| if not cfg: | |
| return [] | |
| for attempt in range(RETRY_ATTEMPTS): | |
| try: | |
| if cfg.get("type") == "custom": | |
| input_payload = { | |
| "prompt": prompt, | |
| "width": 1024, | |
| "height": 1024, | |
| } | |
| else: | |
| param_name = cfg.get("param_name", "aspect_ratio") | |
| input_payload = {"prompt": prompt,param_name: aspect_ratio} | |
| output = replicate.run(cfg["id"], input=input_payload) | |
| urls: List[str] = [] | |
| if isinstance(output, list) and output: | |
| urls = [getattr(output[0], "url", str(output[0]))] | |
| elif isinstance(output, str): | |
| urls = [output] | |
| elif hasattr(output, "url"): | |
| urls = [output.url] | |
| if urls: | |
| return urls | |
| except Exception as e: | |
| if attempt == RETRY_ATTEMPTS - 1: | |
| logger.error(f"replicate run failed: {e}") | |
| return [] | |
| time.sleep(1) | |
| return [] | |
| def _fetch(url: Union[str, Any]) -> Optional[bytes]: | |
| url_str = getattr(url, "url", str(url)) | |
| for attempt in range(RETRY_ATTEMPTS): | |
| try: | |
| r = requests.get( | |
| url_str, timeout=REQUEST_TIMEOUT, stream=True, | |
| headers={"Cache-Control":"no-cache","Pragma":"no-cache","User-Agent":"ImageBot/1.0"} | |
| ) | |
| r.raise_for_status() | |
| buf = b"" | |
| for chunk in r.iter_content(8192): | |
| buf += chunk | |
| return buf | |
| except Exception: | |
| if attempt == RETRY_ATTEMPTS - 1: | |
| return None | |
| time.sleep(1) | |
| return None | |
| def _process_one(args: Tuple[str, str, str, int, bool, Dict[str, Any]]) -> Dict[str, Any]: | |
| model_key, prompt, aspect_ratio, idx, private_mode, params = args | |
| out = {"index": idx, "success": False, "source_url": None, "r2_url": None, "error": None} | |
| try: | |
| urls = _generate_one(model_key, prompt, aspect_ratio, params) | |
| if not urls: | |
| out["error"] = "No URLs returned" | |
| return out | |
| src = urls[0] | |
| out["source_url"] = src | |
| b = _fetch(src) | |
| if not b: | |
| out["error"] = "Fetch failed" | |
| return out | |
| image_with_metadata = meta_data_helper_function(b) | |
| if private_mode: | |
| out["r2_url"] = "data:image/png;base64," + base64.b64encode(image_with_metadata).decode() | |
| else: | |
| out["r2_url"] = _upload_to_r2(image_with_metadata) | |
| out["success"] = bool(out["r2_url"]) | |
| except Exception as e: | |
| out["error"] = str(e) | |
| return out | |
| def _generate_images_parallel( | |
| model_key: str, | |
| aspect_ratio: str, | |
| prompt: str, | |
| num_images: int, | |
| *, | |
| params: Dict[str, Any], | |
| private_mode: bool = False, | |
| ) -> Tuple[List[str], List[str], List[str]]: | |
| args = [ | |
| (model_key, prompt, aspect_ratio, i, private_mode, params) | |
| for i in range(num_images) | |
| ] | |
| r2, src, errs = [], [], [] | |
| with ThreadPoolExecutor(max_workers=min(MAX_WORKERS, num_images)) as ex: | |
| for fut in as_completed({ex.submit(_process_one, a): a for a in args}): | |
| res = fut.result() | |
| if res["success"]: | |
| r2.append(res["r2_url"]) | |
| src.append(res["source_url"]) | |
| else: | |
| errs.append(res["error"] or "Generation failed") | |
| return list(dict.fromkeys(r2)), list(dict.fromkeys(src)), errs | |
| def generate_images_parallel(model_key: str,aspect_ratio: str,prompt: str,num_images: int,*,params: Dict[str, Any],private_mode: bool = False,): | |
| return _generate_images_parallel( | |
| model_key, | |
| aspect_ratio, | |
| prompt.strip(), | |
| num_images, | |
| params=params, | |
| private_mode=private_mode, | |
| ) | |
| def handle_image_generation_optimized( | |
| *, | |
| model_key: str, | |
| aspect_ratio: str, | |
| prompt: str, | |
| num_images: int, | |
| debug_mode: bool = False, | |
| category: Optional[str] = None, | |
| platform: Optional[str] = None, | |
| uid: str, | |
| private_mode: bool = False, | |
| params: Dict[str, Any] | None = None, | |
| ): | |
| """ | |
| Streamlit-friendly wrapper: kicks off parallel gen, persists a job row, | |
| and renders results in-place (no return value). | |
| """ | |
| if not REPLICATE_API_TOKEN: | |
| st.error("Missing REPLICATE_API_TOKEN. Set it as an environment variable.") | |
| return | |
| if not prompt.strip(): | |
| st.warning("Please enter a prompt.") | |
| return | |
| created_by = uid | |
| results_col = None if private_mode else get_results_collection() | |
| db_job_id = None | |
| if results_col is not None: | |
| try: | |
| db_job_id = start_job( | |
| results_col, | |
| type="generation", | |
| created_by=created_by, | |
| category=(category or "general"), | |
| inputs={"model_key": model_key, "aspect_ratio": aspect_ratio, "num_images": num_images}, | |
| settings={"platform": platform}, | |
| user_prompt=prompt.strip(), | |
| ) | |
| except Exception as e: | |
| logger.error(f"start_job failed: {e}") | |
| progress = st.progress(0, text="Starting generation...") | |
| status = st.empty() | |
| start = time.time() | |
| try: | |
| with status.container(): | |
| st.info(f"Generating {num_images} image(s)") | |
| progress.progress(10, text="Running...") | |
| r2_urls, source_urls, errors = _generate_images_parallel( | |
| model_key, | |
| aspect_ratio, | |
| prompt.strip(), | |
| num_images, | |
| params=params or {}, | |
| private_mode=private_mode, | |
| ) | |
| urls = r2_urls if private_mode else (r2_urls or source_urls) | |
| if results_col is not None and db_job_id: | |
| try: | |
| finish_job( | |
| results_col, | |
| db_job_id, | |
| status="completed" if urls else "failed", | |
| outputs_urls=urls or [], | |
| provider_update={"errors": errors} if errors else None, | |
| ) | |
| except Exception as e: | |
| logger.error(f"finish_job failed: {e}") | |
| progress.progress(100, text="Complete!") | |
| took = time.time() - start | |
| if urls: | |
| with status.container(): | |
| message = f"Generated {len(urls)} image(s) in {took:.1f}s." | |
| if not private_mode: | |
| message += f" Job ID: {db_job_id or 'N/A'}" | |
| else: | |
| message += " Private mode: results stay local to this session." | |
| st.success(message) | |
| cols = st.columns(min(4, len(urls)) or 1) | |
| image_bytes_list = [] | |
| for i, u in enumerate(urls): | |
| with cols[i % len(cols)]: | |
| try: | |
| if isinstance(u, str) and u.startswith("data:image"): | |
| try: | |
| _, encoded = u.split(",", 1) | |
| b = base64.b64decode(encoded) | |
| except Exception: | |
| b = None | |
| else: | |
| b = _fetch(u) | |
| if b is None: | |
| st.error("Failed to load image") | |
| continue | |
| image_bytes_list.append((f"image_{i + 1}.png", b)) | |
| st.image(b, width='stretch') | |
| st.download_button( | |
| f"Download image ", | |
| b, | |
| file_name=f"image_{i + 1}.png", | |
| mime="image/png", | |
| width='stretch', | |
| ) | |
| except Exception as e: | |
| st.error(f"Display failed: {e}") | |
| if len(image_bytes_list) > 1: | |
| zip_buffer = io.BytesIO() | |
| with zipfile.ZipFile(zip_buffer, "w") as zf: | |
| for fname, b in image_bytes_list: | |
| zf.writestr(fname, b) | |
| zip_buffer.seek(0) | |
| st.download_button( | |
| " Download All Images", | |
| data=zip_buffer, | |
| file_name="all_images.zip", | |
| mime="application/zip", | |
| width='stretch', | |
| ) | |
| else: | |
| with status.container(): | |
| st.error("No images were generated.") | |
| if errors and debug_mode: | |
| with st.expander("Generation Errors", expanded=True): | |
| for e in errors: | |
| st.error(e) | |
| except Exception as e: | |
| if results_col is not None and db_job_id: | |
| try: | |
| finish_job(results_col, db_job_id, status="failed") | |
| except Exception: | |
| pass | |
| with status.container(): | |
| st.error(f"Generation failed: {e}") | |
| def generate_image(file_path, size, quality, category, sentiment, user_prompt, platform, blur, i=None): | |
| try: | |
| api_key = os.getenv("OPENAI_API_KEY") | |
| if not api_key: | |
| logger.critical("OPENAI_API_KEY is not set.") | |
| raise RuntimeError("OPENAI_API_KEY is missing") | |
| client = OpenAI(api_key=api_key) | |
| with open(file_path, "rb") as img_file: | |
| background = "blurred background." if blur else " not blurred background." | |
| result = client.images.edit( | |
| model="gpt-image-1", | |
| prompt=( | |
| f"You are a top-tier performance digital marketer and creative strategist with 15+ years of expertise in affiliate marketing.\n" | |
| f"Your objective is to analyze the provided winning ad image, deconstruct its concept, visual composition, and color scheme, and generate a fresh, conversion-focused ad visual tailored for the {category} niche.\n" | |
| f"The new design should convey a {sentiment} sentiment and incorporate the user instruction: \n {user_prompt}.\n If user has given multple choices or options to be include in the image so choose randomly relevant to the reference image." | |
| f"Create a visually compelling ad optimized for {platform} Ads that is scroll-stopping, pattern-interrupting, and designed to drive high CTR and Conversion Rate. Utilize striking color combinations, dynamic contrast levels, and strategic layout compositions to command attention while aligning with the target audience avatar.\n" | |
| f"Make sure the images should be realistic, not be stocky at all and raw which should look like they are shot from an iPhone with {background}." | |
| ), | |
| image=img_file, | |
| size=size, | |
| quality=quality, | |
| ) | |
| image_base64 = result.data[0].b64_json | |
| image_bytes = base64.b64decode(image_base64) | |
| logger.info(f"Successfully generated image for {file_path}") | |
| return image_bytes | |
| except Exception as e: | |
| logger.exception(f"Failed to generate image for {file_path}: {e}") | |
| raise |