Spaces:
Sleeping
Sleeping
| """ | |
| Image Engine — Generates scene images via OpenRouter using black-forest-labs/flux.2-pro. | |
| Saves images as compressed JPEG files to avoid huge base64 data URIs. | |
| """ | |
| import os | |
| import re | |
| import io | |
| import base64 | |
| import requests | |
| from PIL import Image | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "") | |
| MODEL_ID = "black-forest-labs/flux.2-pro" | |
| OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" | |
| # Directory to store generated scene images | |
| IMG_DIR = os.path.join(os.path.dirname(__file__), "generated_scenes") | |
| os.makedirs(IMG_DIR, exist_ok=True) | |
| _img_counter = 0 | |
| def _save_as_jpeg(data_uri: str) -> str: | |
| """Convert a base64 data URI (PNG) to a compressed JPEG file. Returns file path.""" | |
| global _img_counter | |
| _img_counter += 1 | |
| # Strip data URI prefix | |
| b64_data = data_uri.split(",", 1)[1] if "," in data_uri else data_uri | |
| raw = base64.b64decode(b64_data) | |
| img = Image.open(io.BytesIO(raw)).convert("RGB") | |
| path = os.path.join(IMG_DIR, f"scene_{_img_counter}.jpg") | |
| img.save(path, "JPEG", quality=72) | |
| size_kb = os.path.getsize(path) / 1024 | |
| print(f"[ImageEngine] Saved JPEG ({size_kb:.0f} KB): {path}") | |
| return path | |
| def generate_scene_image(prompt: str) -> str | None: | |
| """ | |
| Generate a scene image from a text prompt via OpenRouter chat/completions. | |
| Returns a local file path to a JPEG or None on failure. | |
| """ | |
| headers = { | |
| "Authorization": f"Bearer {OPENROUTER_API_KEY}", | |
| "Content-Type": "application/json", | |
| } | |
| full_prompt = ( | |
| f"A cinematic wide-angle theatrical stage scene, dramatic lighting, " | |
| f"rich colors, film still quality: {prompt}" | |
| ) | |
| payload = { | |
| "model": MODEL_ID, | |
| "messages": [ | |
| {"role": "user", "content": full_prompt} | |
| ], | |
| "modalities": ["image"], | |
| } | |
| try: | |
| print(f"[ImageEngine] Generating image: {prompt[:60]}...") | |
| resp = requests.post( | |
| OPENROUTER_URL, | |
| headers=headers, | |
| json=payload, | |
| timeout=60, | |
| ) | |
| if resp.status_code != 200: | |
| print(f"[ImageEngine] API Error ({resp.status_code}): {resp.text[:200]}") | |
| return None | |
| data = resp.json() | |
| message = data.get("choices", [{}])[0].get("message", {}) | |
| # Flux returns images in message.images[] array | |
| images = message.get("images", []) | |
| if images: | |
| img = images[0] | |
| if isinstance(img, dict): | |
| url = img.get("image_url", {}).get("url") or img.get("url") | |
| if url and url.startswith("data:image"): | |
| return _save_as_jpeg(url) | |
| elif url: | |
| return url # External URL, use directly | |
| # Fallback: check content field | |
| content = message.get("content") | |
| if content: | |
| if content.startswith("data:image"): | |
| return _save_as_jpeg(content) | |
| md_match = re.search(r'!\[.*?\]\((.*?)\)', content) | |
| if md_match: | |
| return md_match.group(1) | |
| print(f"[ImageEngine] Could not extract image from response.") | |
| return None | |
| except Exception as e: | |
| print(f"[ImageEngine] Error: {e}") | |
| return None | |