Manireddy1508 commited on
Commit
f29389b
·
verified ·
1 Parent(s): 6e52c0e

Update utils/planner.py

Browse files
Files changed (1) hide show
  1. utils/planner.py +14 -57
utils/planner.py CHANGED
@@ -1,64 +1,20 @@
1
  # utils/planner.py
2
 
3
- import openai
4
  import os
5
- from openai import OpenAI
6
- from dotenv import load_dotenv
7
  import json
8
- load_dotenv()
9
-
10
- client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
11
-
12
- SYSTEM_INSTRUCTIONS = """
13
- You are a scene planning assistant for an AI image generation system.
14
- Your job is to take the user's prompt and return a structured JSON with:
15
- - scene (environment, setting)
16
- - subject (main actor)
17
- - objects (main product or items)
18
- - layout (foreground/background elements and their placement)
19
- - rules (validation rules to ensure visual correctness)
20
- Respond ONLY in raw JSON format. Do NOT include explanations.
21
- """
22
-
23
- def extract_scene_plan(prompt: str) -> dict:
24
- try:
25
- response = client.chat.completions.create(
26
- model="gpt-4o-mini-2024-07-18",
27
- messages=[
28
- {"role": "system", "content": SYSTEM_INSTRUCTIONS},
29
- {"role": "user", "content": prompt}
30
- ],
31
- temperature=0.3,
32
- max_tokens=500
33
- )
34
- json_output = response.choices[0].message.content
35
- print("🧠 Scene Plan (Raw):", json_output)
36
- return json.loads(json_output) # Be cautious: Use `json.loads()` if possible
37
-
38
- except Exception as e:
39
- print("❌ extract_scene_plan() Error:", e)
40
- return {
41
- "scene": None,
42
- "subject": None,
43
- "objects": [],
44
- "layout": {},
45
- "rules": {}
46
- }
47
-
48
-
49
- # utils/planner.py
50
-
51
- import openai
52
- import os
53
- from openai import OpenAI
54
  from dotenv import load_dotenv
55
- import json
56
- load_dotenv()
57
 
 
 
 
 
58
  client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
59
 
 
60
  # 🧠 Scene Plan Extractor
61
- SYSTEM_INSTRUCTIONS = """
 
62
  You are a scene planning assistant for an AI image generation system.
63
  Your job is to take the user's prompt and return a structured JSON with:
64
  - scene (environment, setting)
@@ -74,7 +30,7 @@ def extract_scene_plan(prompt: str) -> dict:
74
  response = client.chat.completions.create(
75
  model="gpt-4o-mini-2024-07-18",
76
  messages=[
77
- {"role": "system", "content": SYSTEM_INSTRUCTIONS},
78
  {"role": "user", "content": prompt}
79
  ],
80
  temperature=0.3,
@@ -94,19 +50,19 @@ def extract_scene_plan(prompt: str) -> dict:
94
  "rules": {}
95
  }
96
 
97
- # 🧠 GPT-Based Prompt Variation Generator
 
 
98
  def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
99
  try:
100
  system_msg = f"""
101
  You are a creative prompt variation generator for an AI image generation system.
102
-
103
  Given a base user prompt and its structured scene plan, generate {n} diverse image generation prompts.
104
  Each prompt should:
105
  - Be visually rich and descriptive
106
  - Include stylistic or contextual variation
107
  - Reference the same product and environment
108
  - Stay faithful to the base prompt and extracted plan
109
-
110
  Respond ONLY with a JSON array of strings. No explanations.
111
  """
112
 
@@ -129,4 +85,5 @@ Respond ONLY with a JSON array of strings. No explanations.
129
 
130
  except Exception as e:
131
  print("❌ generate_prompt_variations_from_scene() Error:", e)
132
- return [base_prompt] # fallback to original if anything fails
 
 
1
  # utils/planner.py
2
 
 
3
  import os
 
 
4
  import json
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  from dotenv import load_dotenv
6
+ from openai import OpenAI
 
7
 
8
+ # ----------------------------
9
+ # 🔐 Load Environment & Client
10
+ # ----------------------------
11
+ load_dotenv()
12
  client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
13
 
14
+ # ----------------------------
15
  # 🧠 Scene Plan Extractor
16
+ # ----------------------------
17
+ SCENE_SYSTEM_INSTRUCTIONS = """
18
  You are a scene planning assistant for an AI image generation system.
19
  Your job is to take the user's prompt and return a structured JSON with:
20
  - scene (environment, setting)
 
30
  response = client.chat.completions.create(
31
  model="gpt-4o-mini-2024-07-18",
32
  messages=[
33
+ {"role": "system", "content": SCENE_SYSTEM_INSTRUCTIONS},
34
  {"role": "user", "content": prompt}
35
  ],
36
  temperature=0.3,
 
50
  "rules": {}
51
  }
52
 
53
+ # ----------------------------
54
+ # 🧠 Prompt Variation Generator
55
+ # ----------------------------
56
  def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
57
  try:
58
  system_msg = f"""
59
  You are a creative prompt variation generator for an AI image generation system.
 
60
  Given a base user prompt and its structured scene plan, generate {n} diverse image generation prompts.
61
  Each prompt should:
62
  - Be visually rich and descriptive
63
  - Include stylistic or contextual variation
64
  - Reference the same product and environment
65
  - Stay faithful to the base prompt and extracted plan
 
66
  Respond ONLY with a JSON array of strings. No explanations.
67
  """
68
 
 
85
 
86
  except Exception as e:
87
  print("❌ generate_prompt_variations_from_scene() Error:", e)
88
+ return [base_prompt] # fallback to original
89
+