Spaces:
Paused
Paused
File size: 2,975 Bytes
9e3bd6c d3a3bf1 69a2277 f29389b 69a2277 f29389b 69a2277 f29389b 69a2277 f29389b 69a2277 f29389b 69a2277 f29389b 69a2277 f29389b |
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 |
# utils/planner.py
import os
import json
from dotenv import load_dotenv
from openai import OpenAI
# ----------------------------
# 🔐 Load Environment & Client
# ----------------------------
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
# ----------------------------
# 🧠 Scene Plan Extractor
# ----------------------------
SCENE_SYSTEM_INSTRUCTIONS = """
You are a scene planning assistant for an AI image generation system.
Your job is to take the user's prompt and return a structured JSON with:
- scene (environment, setting)
- subject (main actor)
- objects (main product or items)
- layout (foreground/background elements and their placement)
- rules (validation rules to ensure visual correctness)
Respond ONLY in raw JSON format. Do NOT include explanations.
"""
def extract_scene_plan(prompt: str) -> dict:
try:
response = client.chat.completions.create(
model="gpt-4o-mini-2024-07-18",
messages=[
{"role": "system", "content": SCENE_SYSTEM_INSTRUCTIONS},
{"role": "user", "content": prompt}
],
temperature=0.3,
max_tokens=500
)
json_output = response.choices[0].message.content
print("🧠 Scene Plan (Raw):", json_output)
return json.loads(json_output)
except Exception as e:
print("❌ extract_scene_plan() Error:", e)
return {
"scene": None,
"subject": None,
"objects": [],
"layout": {},
"rules": {}
}
# ----------------------------
# 🧠 Prompt Variation Generator
# ----------------------------
def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
try:
system_msg = f"""
You are a creative prompt variation generator for an AI image generation system.
Given a base user prompt and its structured scene plan, generate {n} diverse image generation prompts.
Each prompt should:
- Be visually rich and descriptive
- Include stylistic or contextual variation
- Reference the same product and environment
- Stay faithful to the base prompt and extracted plan
Respond ONLY with a JSON array of strings. No explanations.
"""
response = client.chat.completions.create(
model="gpt-4o-mini-2024-07-18",
messages=[
{"role": "system", "content": system_msg},
{"role": "user", "content": json.dumps({
"base_prompt": base_prompt,
"scene_plan": scene_plan
})}
],
temperature=0.7,
max_tokens=600
)
content = response.choices[0].message.content
print("🧠 Prompt Variations (Raw):", content)
return json.loads(content)
except Exception as e:
print("❌ generate_prompt_variations_from_scene() Error:", e)
return [base_prompt] # fallback to original
|