# app/services/image_generation_service.py import os import logging import mimetypes import traceback import tempfile from typing import List, Tuple, Optional from google import genai from google.genai import types from app.core.config import settings # Basic logger for this module. The app-level startup should configure logging more globally. logger = logging.getLogger(__name__) # Ensure the Gemini API key exists if not settings.GEMINI_API_KEY: raise ValueError("GEMINI_API_KEY not found in environment or .env file") client = genai.Client(api_key=settings.GEMINI_API_KEY) IMAGE_GEN_MODEL = settings.IMAGE_GEN_MODEL def _safe_info(msg: str, *args, **kwargs) -> None: """Log info but ignore low-level IO errors (BrokenPipeError) that can occur in containers.""" try: logger.info(msg, *args, **kwargs) except (BrokenPipeError, OSError): # stdout/stderr might be closed by the launcher; swallow these errors pass except Exception: # don't allow logging problems to interrupt business logic pass def _safe_error(msg: str, *args, exc: Optional[BaseException] = None, **kwargs) -> None: """Log error and exception safely.""" try: if exc is not None: logger.exception(msg, *args, **kwargs) else: logger.error(msg, *args, **kwargs) except (BrokenPipeError, OSError): pass except Exception: pass def _try_delete_uploaded(uploaded_obj) -> None: """ Attempt to delete a previously uploaded file from GenAI service. This is best-effort: we handle multiple possible attribute names and suppress errors. """ if uploaded_obj is None: return # Common fields that might identify the uploaded file candidates = [] # google genai's uploaded file object may have attributes like: name, uri, id for attr in ("name", "uri", "id", "file_id"): val = getattr(uploaded_obj, attr, None) if val: candidates.append((attr, val)) # Try a few delete patterns (best-effort) try: # If the client provides a delete method, call it. This API surface may vary by SDK version. if hasattr(client.files, "delete"): try: # If uploaded_obj has a 'name' attribute this is common for many SDKs name = getattr(uploaded_obj, "name", None) if name: client.files.delete(name) _safe_info(f"[GenAI Info] Deleted uploaded file by name: {name}") return except Exception: # continue to other attempts pass # If we have a URI with a resource name, try to request deletion via client.files.delete with URI uri = getattr(uploaded_obj, "uri", None) if uri: try: client.files.delete(uri) _safe_info(f"[GenAI Info] Deleted uploaded file by uri: {uri}") return except Exception: pass # Last resort: try deleting by id if present file_id = getattr(uploaded_obj, "id", None) or getattr(uploaded_obj, "file_id", None) if file_id: try: client.files.delete(file_id) _safe_info(f"[GenAI Info] Deleted uploaded file by id: {file_id}") return except Exception: pass except Exception as e: _safe_error(f"[GenAI Error] Exception while attempting to delete uploaded file {uploaded_obj}: {e}", exc=e) async def generate_image_from_files_and_prompt( image_files: List[Tuple[bytes, str]], # list of (bytes, original_filename) prompt: str, ) -> Optional[bytes]: """ Uploads provided image bytes to GenAI, requests an image generation with the prompt, and returns the generated image bytes (or None on failure). Args: image_files: list of tuples (file_bytes, original_filename). The function expects at least one image; two images is common for 'replace' flows. prompt: textual prompt to guide generation. Returns: bytes of generated image (if present) or None on failure. """ temp_file_paths: List[str] = [] uploaded_file_infos: List[object] = [] parts: List[types.Part] = [] try: # 1) Write incoming bytes to persistent temp files and upload them for img_bytes, original_filename in image_files: # Determine suffix and mime type suffix = os.path.splitext(original_filename or "")[1] or "" mime_type, _ = mimetypes.guess_type(original_filename or "") if not mime_type: mime_type = "application/octet-stream" # create a temp file and persist its path (we'll cleanup in finally) with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(img_bytes) tmp_path = tmp.name temp_file_paths.append(tmp_path) _safe_info("[GenAI Info] Created temporary file: %s (mime: %s)", tmp_path, mime_type) # Upload using client.files.upload. Keep the uploaded file info for potential cleanup try: # The SDK may accept either a path string or a file-like object. uploaded = client.files.upload(file=tmp_path) uploaded_file_infos.append(uploaded) # The uploaded object may have attributes like 'uri' and 'mime_type' uri = getattr(uploaded, "uri", None) uploaded_mime = getattr(uploaded, "mime_type", mime_type) parts.append( types.Part.from_uri( file_uri=uri if uri else tmp_path, # fallback to path if SDK didn't return uri mime_type=uploaded_mime or mime_type, ) ) _safe_info("[GenAI Info] Uploaded file %s -> uri=%s", tmp_path, uri) except Exception as e: _safe_error(f"[GenAI Error] Upload failed for {tmp_path}: {e}", exc=e) # continue loop so we still try to clean up temp files; then fail raise # 2) Add prompt as part parts.append(types.Part.from_text(text=prompt)) contents = [types.Content(role="user", parts=parts)] # 3) Prepare config (safety settings preserved) generate_content_config = types.GenerateContentConfig( response_modalities=["IMAGE", "TEXT"], safety_settings=[ types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"), types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"), types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"), types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"), ], ) _safe_info("[GenAI Info] Requesting generation with model=%s", IMAGE_GEN_MODEL) response = client.models.generate_content( model=IMAGE_GEN_MODEL, contents=contents, config=generate_content_config, ) # 4) Parse response - look for image bytes in candidate parts if getattr(response, "candidates", None): for candidate in response.candidates: content = getattr(candidate, "content", None) if not content: _safe_info("[GenAI Warning] Candidate has no content; skipping.") continue parts_list = getattr(content, "parts", []) or [] for part in parts_list: # inline_data.data is where binary bytes typically live in SDK responses inline_data = getattr(part, "inline_data", None) if inline_data is not None: data_bytes = getattr(inline_data, "data", None) if data_bytes: _safe_info("[GenAI Info] Found image bytes in response; returning bytes.") return data_bytes # alternate: some SDKs return base64-encoded strings in text fields; handle common fallbacks text_val = getattr(part, "text", None) if text_val and isinstance(text_val, str) and text_val.startswith("data:image/"): # data URL -> decode base64 portion try: header, b64 = text_val.split(",", 1) import base64 data_bytes = base64.b64decode(b64) _safe_info("[GenAI Info] Extracted image bytes from data URL in text part.") return data_bytes except Exception: _safe_error("[GenAI Error] Failed to decode data URL in text part.", exc=None) # If part has 'uri' that points to a generated asset, attempt to fetch it (best-effort) part_uri = getattr(part, "uri", None) if part_uri: try: # Use client.files.get or HTTP fetch as available. We'll try client.files.get if present. if hasattr(client.files, "get"): fetched = client.files.get(part_uri) # fetched could include bytes in different fields; try common ones data_bytes = getattr(fetched, "data", None) or getattr(fetched, "content", None) if data_bytes: _safe_info("[GenAI Info] Fetched generated image from part.uri via client.files.get") return data_bytes else: # Fall back to HTTP GET if the URI is an http(s) link if str(part_uri).startswith("http"): import requests as _req r = _req.get(part_uri, timeout=10) if r.status_code == 200: _safe_info("[GenAI Info] Fetched generated image from HTTP URI in part.") return r.content except Exception as e: _safe_error(f"[GenAI Error] Failed to fetch part URI {part_uri}: {e}", exc=e) # If we reach here no image bytes found _safe_info("[GenAI Warning] No image bytes found in response candidates; returning None.") return None except Exception as e: # Log full traceback safely for debugging _safe_error(f"[GenAI Error] Image generation failed: {e}", exc=e) try: # print stack trace to logger (safe wrapper) tb = traceback.format_exc() _safe_error(f"[GenAI Error] Traceback:\n{tb}") except Exception: pass return None finally: # Cleanup temporary files created locally for path in temp_file_paths: try: if path and os.path.exists(path): os.remove(path) _safe_info("[GenAI Info] Deleted temporary file: %s", path) except Exception as e: _safe_error(f"[GenAI Error] Failed to delete temporary file {path}: {e}", exc=e) # Attempt to delete uploaded files from GenAI (best-effort) for uploaded in uploaded_file_infos: try: _try_delete_uploaded(uploaded) except Exception as e: _safe_error(f"[GenAI Error] Failed to delete uploaded file record {uploaded}: {e}", exc=e)