npcpy 1.0.26__py3-none-any.whl → 1.2.32__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (148) hide show
  1. npcpy/__init__.py +0 -7
  2. npcpy/data/audio.py +16 -99
  3. npcpy/data/image.py +43 -42
  4. npcpy/data/load.py +83 -124
  5. npcpy/data/text.py +28 -28
  6. npcpy/data/video.py +8 -32
  7. npcpy/data/web.py +51 -23
  8. npcpy/ft/diff.py +110 -0
  9. npcpy/ft/ge.py +115 -0
  10. npcpy/ft/memory_trainer.py +171 -0
  11. npcpy/ft/model_ensembler.py +357 -0
  12. npcpy/ft/rl.py +360 -0
  13. npcpy/ft/sft.py +248 -0
  14. npcpy/ft/usft.py +128 -0
  15. npcpy/gen/audio_gen.py +24 -0
  16. npcpy/gen/embeddings.py +13 -13
  17. npcpy/gen/image_gen.py +262 -117
  18. npcpy/gen/response.py +615 -415
  19. npcpy/gen/video_gen.py +53 -7
  20. npcpy/llm_funcs.py +1869 -437
  21. npcpy/main.py +1 -1
  22. npcpy/memory/command_history.py +844 -510
  23. npcpy/memory/kg_vis.py +833 -0
  24. npcpy/memory/knowledge_graph.py +892 -1845
  25. npcpy/memory/memory_processor.py +81 -0
  26. npcpy/memory/search.py +188 -90
  27. npcpy/mix/debate.py +192 -3
  28. npcpy/npc_compiler.py +1672 -801
  29. npcpy/npc_sysenv.py +593 -1266
  30. npcpy/serve.py +3120 -0
  31. npcpy/sql/ai_function_tools.py +257 -0
  32. npcpy/sql/database_ai_adapters.py +186 -0
  33. npcpy/sql/database_ai_functions.py +163 -0
  34. npcpy/sql/model_runner.py +19 -19
  35. npcpy/sql/npcsql.py +706 -507
  36. npcpy/sql/sql_model_compiler.py +156 -0
  37. npcpy/tools.py +183 -0
  38. npcpy/work/plan.py +13 -279
  39. npcpy/work/trigger.py +3 -3
  40. npcpy-1.2.32.dist-info/METADATA +803 -0
  41. npcpy-1.2.32.dist-info/RECORD +54 -0
  42. npcpy/data/dataframes.py +0 -171
  43. npcpy/memory/deep_research.py +0 -125
  44. npcpy/memory/sleep.py +0 -557
  45. npcpy/modes/_state.py +0 -78
  46. npcpy/modes/alicanto.py +0 -1075
  47. npcpy/modes/guac.py +0 -785
  48. npcpy/modes/mcp_npcsh.py +0 -822
  49. npcpy/modes/npc.py +0 -213
  50. npcpy/modes/npcsh.py +0 -1158
  51. npcpy/modes/plonk.py +0 -409
  52. npcpy/modes/pti.py +0 -234
  53. npcpy/modes/serve.py +0 -1637
  54. npcpy/modes/spool.py +0 -312
  55. npcpy/modes/wander.py +0 -549
  56. npcpy/modes/yap.py +0 -572
  57. npcpy/npc_team/alicanto.npc +0 -2
  58. npcpy/npc_team/alicanto.png +0 -0
  59. npcpy/npc_team/assembly_lines/test_pipeline.py +0 -181
  60. npcpy/npc_team/corca.npc +0 -13
  61. npcpy/npc_team/foreman.npc +0 -7
  62. npcpy/npc_team/frederic.npc +0 -6
  63. npcpy/npc_team/frederic4.png +0 -0
  64. npcpy/npc_team/guac.png +0 -0
  65. npcpy/npc_team/jinxs/automator.jinx +0 -18
  66. npcpy/npc_team/jinxs/bash_executer.jinx +0 -31
  67. npcpy/npc_team/jinxs/calculator.jinx +0 -11
  68. npcpy/npc_team/jinxs/edit_file.jinx +0 -96
  69. npcpy/npc_team/jinxs/file_chat.jinx +0 -14
  70. npcpy/npc_team/jinxs/gui_controller.jinx +0 -28
  71. npcpy/npc_team/jinxs/image_generation.jinx +0 -29
  72. npcpy/npc_team/jinxs/internet_search.jinx +0 -30
  73. npcpy/npc_team/jinxs/local_search.jinx +0 -152
  74. npcpy/npc_team/jinxs/npcsh_executor.jinx +0 -31
  75. npcpy/npc_team/jinxs/python_executor.jinx +0 -8
  76. npcpy/npc_team/jinxs/screen_cap.jinx +0 -25
  77. npcpy/npc_team/jinxs/sql_executor.jinx +0 -33
  78. npcpy/npc_team/kadiefa.npc +0 -3
  79. npcpy/npc_team/kadiefa.png +0 -0
  80. npcpy/npc_team/npcsh.ctx +0 -9
  81. npcpy/npc_team/npcsh_sibiji.png +0 -0
  82. npcpy/npc_team/plonk.npc +0 -2
  83. npcpy/npc_team/plonk.png +0 -0
  84. npcpy/npc_team/plonkjr.npc +0 -2
  85. npcpy/npc_team/plonkjr.png +0 -0
  86. npcpy/npc_team/sibiji.npc +0 -5
  87. npcpy/npc_team/sibiji.png +0 -0
  88. npcpy/npc_team/spool.png +0 -0
  89. npcpy/npc_team/templates/analytics/celona.npc +0 -0
  90. npcpy/npc_team/templates/hr_support/raone.npc +0 -0
  91. npcpy/npc_team/templates/humanities/eriane.npc +0 -4
  92. npcpy/npc_team/templates/it_support/lineru.npc +0 -0
  93. npcpy/npc_team/templates/marketing/slean.npc +0 -4
  94. npcpy/npc_team/templates/philosophy/maurawa.npc +0 -0
  95. npcpy/npc_team/templates/sales/turnic.npc +0 -4
  96. npcpy/npc_team/templates/software/welxor.npc +0 -0
  97. npcpy/npc_team/yap.png +0 -0
  98. npcpy/routes.py +0 -958
  99. npcpy/work/mcp_helpers.py +0 -357
  100. npcpy/work/mcp_server.py +0 -194
  101. npcpy-1.0.26.data/data/npcpy/npc_team/alicanto.npc +0 -2
  102. npcpy-1.0.26.data/data/npcpy/npc_team/alicanto.png +0 -0
  103. npcpy-1.0.26.data/data/npcpy/npc_team/automator.jinx +0 -18
  104. npcpy-1.0.26.data/data/npcpy/npc_team/bash_executer.jinx +0 -31
  105. npcpy-1.0.26.data/data/npcpy/npc_team/calculator.jinx +0 -11
  106. npcpy-1.0.26.data/data/npcpy/npc_team/celona.npc +0 -0
  107. npcpy-1.0.26.data/data/npcpy/npc_team/corca.npc +0 -13
  108. npcpy-1.0.26.data/data/npcpy/npc_team/edit_file.jinx +0 -96
  109. npcpy-1.0.26.data/data/npcpy/npc_team/eriane.npc +0 -4
  110. npcpy-1.0.26.data/data/npcpy/npc_team/file_chat.jinx +0 -14
  111. npcpy-1.0.26.data/data/npcpy/npc_team/foreman.npc +0 -7
  112. npcpy-1.0.26.data/data/npcpy/npc_team/frederic.npc +0 -6
  113. npcpy-1.0.26.data/data/npcpy/npc_team/frederic4.png +0 -0
  114. npcpy-1.0.26.data/data/npcpy/npc_team/guac.png +0 -0
  115. npcpy-1.0.26.data/data/npcpy/npc_team/gui_controller.jinx +0 -28
  116. npcpy-1.0.26.data/data/npcpy/npc_team/image_generation.jinx +0 -29
  117. npcpy-1.0.26.data/data/npcpy/npc_team/internet_search.jinx +0 -30
  118. npcpy-1.0.26.data/data/npcpy/npc_team/kadiefa.npc +0 -3
  119. npcpy-1.0.26.data/data/npcpy/npc_team/kadiefa.png +0 -0
  120. npcpy-1.0.26.data/data/npcpy/npc_team/lineru.npc +0 -0
  121. npcpy-1.0.26.data/data/npcpy/npc_team/local_search.jinx +0 -152
  122. npcpy-1.0.26.data/data/npcpy/npc_team/maurawa.npc +0 -0
  123. npcpy-1.0.26.data/data/npcpy/npc_team/npcsh.ctx +0 -9
  124. npcpy-1.0.26.data/data/npcpy/npc_team/npcsh_executor.jinx +0 -31
  125. npcpy-1.0.26.data/data/npcpy/npc_team/npcsh_sibiji.png +0 -0
  126. npcpy-1.0.26.data/data/npcpy/npc_team/plonk.npc +0 -2
  127. npcpy-1.0.26.data/data/npcpy/npc_team/plonk.png +0 -0
  128. npcpy-1.0.26.data/data/npcpy/npc_team/plonkjr.npc +0 -2
  129. npcpy-1.0.26.data/data/npcpy/npc_team/plonkjr.png +0 -0
  130. npcpy-1.0.26.data/data/npcpy/npc_team/python_executor.jinx +0 -8
  131. npcpy-1.0.26.data/data/npcpy/npc_team/raone.npc +0 -0
  132. npcpy-1.0.26.data/data/npcpy/npc_team/screen_cap.jinx +0 -25
  133. npcpy-1.0.26.data/data/npcpy/npc_team/sibiji.npc +0 -5
  134. npcpy-1.0.26.data/data/npcpy/npc_team/sibiji.png +0 -0
  135. npcpy-1.0.26.data/data/npcpy/npc_team/slean.npc +0 -4
  136. npcpy-1.0.26.data/data/npcpy/npc_team/spool.png +0 -0
  137. npcpy-1.0.26.data/data/npcpy/npc_team/sql_executor.jinx +0 -33
  138. npcpy-1.0.26.data/data/npcpy/npc_team/test_pipeline.py +0 -181
  139. npcpy-1.0.26.data/data/npcpy/npc_team/turnic.npc +0 -4
  140. npcpy-1.0.26.data/data/npcpy/npc_team/welxor.npc +0 -0
  141. npcpy-1.0.26.data/data/npcpy/npc_team/yap.png +0 -0
  142. npcpy-1.0.26.dist-info/METADATA +0 -827
  143. npcpy-1.0.26.dist-info/RECORD +0 -139
  144. npcpy-1.0.26.dist-info/entry_points.txt +0 -11
  145. /npcpy/{modes → ft}/__init__.py +0 -0
  146. {npcpy-1.0.26.dist-info → npcpy-1.2.32.dist-info}/WHEEL +0 -0
  147. {npcpy-1.0.26.dist-info → npcpy-1.2.32.dist-info}/licenses/LICENSE +0 -0
  148. {npcpy-1.0.26.dist-info → npcpy-1.2.32.dist-info}/top_level.txt +0 -0
npcpy/gen/image_gen.py CHANGED
@@ -7,7 +7,8 @@ import PIL
7
7
  from PIL import Image
8
8
 
9
9
  from litellm import image_generation
10
-
10
+ import requests
11
+ from urllib.request import urlopen
11
12
 
12
13
 
13
14
  def generate_image_diffusers(
@@ -17,18 +18,84 @@ def generate_image_diffusers(
17
18
  height: int = 512,
18
19
  width: int = 512,
19
20
  ):
20
- """Generate an image using the Stable Diffusion API."""
21
- from diffusers import StableDiffusionPipeline
21
+ """Generate an image using the Stable Diffusion API with memory optimization."""
22
+ import torch
23
+ import gc
24
+
25
+ try:
26
+ torch_dtype = torch.float16 if device != "cpu" and torch.cuda.is_available() else torch.float32
27
+
28
+ if 'Qwen' in model:
29
+ from diffusers import DiffusionPipeline
30
+
31
+ pipe = DiffusionPipeline.from_pretrained(
32
+ model,
33
+ torch_dtype=torch_dtype,
34
+ use_safetensors=True,
35
+ variant="fp16" if torch_dtype == torch.float16 else None,
36
+ )
37
+ else:
38
+ from diffusers import StableDiffusionPipeline
39
+
40
+ pipe = StableDiffusionPipeline.from_pretrained(
41
+ model,
42
+ torch_dtype=torch_dtype,
43
+ use_safetensors=True,
44
+ variant="fp16" if torch_dtype == torch.float16 else None,
45
+ )
46
+
47
+ if hasattr(pipe, 'enable_attention_slicing'):
48
+ pipe.enable_attention_slicing()
49
+
50
+ if hasattr(pipe, 'enable_model_cpu_offload') and device != "cpu":
51
+ pipe.enable_model_cpu_offload()
52
+ else:
53
+ pipe = pipe.to(device)
54
+
55
+ if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
56
+ try:
57
+ pipe.enable_xformers_memory_efficient_attention()
58
+ except Exception:
59
+ pass
60
+
61
+ with torch.inference_mode():
62
+ result = pipe(prompt, height=height, width=width, num_inference_steps=20)
63
+ image = result.images[0]
64
+
65
+ del pipe
66
+ if torch.cuda.is_available():
67
+ torch.cuda.empty_cache()
68
+ gc.collect()
69
+
70
+ return image
71
+
72
+ except Exception as e:
73
+ if 'pipe' in locals():
74
+ del pipe
75
+ if torch.cuda.is_available():
76
+ torch.cuda.empty_cache()
77
+ gc.collect()
78
+
79
+ if "out of memory" in str(e).lower() or "killed" in str(e).lower():
80
+ print(f"Memory error during image generation: {e}")
81
+ print("Suggestions:")
82
+ print("1. Try using a smaller model like 'CompVis/stable-diffusion-v1-4'")
83
+ print("2. Reduce image dimensions (e.g., height=256, width=256)")
84
+ print("3. Use a cloud service for image generation")
85
+ raise MemoryError(f"Insufficient memory for image generation with model {model}. Try a smaller model or reduce image size.")
86
+ else:
87
+ raise e
22
88
 
23
- pipe = StableDiffusionPipeline.from_pretrained(model)
24
- pipe = pipe.to(device)
89
+ import os
90
+ import base64
91
+ import io
92
+ from typing import Union, List, Optional
25
93
 
26
- # Generate the image
27
- image = pipe(prompt, height=height, width=width)
28
- image = image.images[0]
29
-
30
- return image
94
+ import PIL
95
+ from PIL import Image
31
96
 
97
+ import requests
98
+ from urllib.request import urlopen
32
99
 
33
100
  def openai_image_gen(
34
101
  prompt: str,
@@ -40,63 +107,77 @@ def openai_image_gen(
40
107
  ):
41
108
  """Generate or edit an image using the OpenAI API."""
42
109
  from openai import OpenAI
43
-
110
+
44
111
  client = OpenAI()
45
-
112
+
46
113
  if height is None:
47
114
  height = 1024
48
115
  if width is None:
49
- width = 1024
50
- # Prepare image for editing if provided
116
+ width = 1024
117
+
118
+ size_str = f"{width}x{height}"
119
+
51
120
  if attachments is not None:
52
- # Process the attachments into the format OpenAI expects
53
121
  processed_images = []
54
-
122
+ files_to_close = []
55
123
  for attachment in attachments:
56
124
  if isinstance(attachment, str):
57
- # Assume it's a file path
58
- processed_images.append(open(attachment, "rb"))
125
+ file_handle = open(attachment, "rb")
126
+ processed_images.append(file_handle)
127
+ files_to_close.append(file_handle)
59
128
  elif isinstance(attachment, bytes):
60
- processed_images.append(io.BytesIO(attachment))
129
+ img_byte_arr = io.BytesIO(attachment)
130
+ img_byte_arr.name = 'image.png' # FIX: Add filename hint
131
+ processed_images.append(img_byte_arr)
61
132
  elif isinstance(attachment, Image.Image):
62
133
  img_byte_arr = io.BytesIO()
63
134
  attachment.save(img_byte_arr, format='PNG')
64
135
  img_byte_arr.seek(0)
136
+ img_byte_arr.name = 'image.png' # FIX: Add filename hint
65
137
  processed_images.append(img_byte_arr)
66
- # Use images.edit for image editing
67
- result = client.images.edit(
68
- model=model,
69
- image=processed_images[0], # Edit only supports a single image
70
- prompt=prompt,
71
- n=n_images,
72
- size=f"{width}x{height}",
73
- )
138
+
139
+ try:
140
+ result = client.images.edit(
141
+ model=model,
142
+ image=processed_images[0],
143
+ prompt=prompt,
144
+ n=n_images,
145
+ size=size_str,
146
+ )
147
+ finally:
148
+ # This ensures any files we opened are properly closed
149
+ for f in files_to_close:
150
+ f.close()
74
151
  else:
75
- # Use images.generate for new image generation
76
152
  result = client.images.generate(
77
153
  model=model,
78
154
  prompt=prompt,
79
155
  n=n_images,
80
- size=f"{width}x{height}",
156
+ size=size_str,
81
157
  )
82
- if model =='gpt-image-1':
83
- image_base64 = result.data[0].b64_json
84
- image_bytes = base64.b64decode(image_base64)
85
- image = Image.open(io.BytesIO(image_bytes))
86
- image.save('generated_image.png')
87
- elif model == 'dall-e-2' or model == 'dall-e-3':
88
- image_base64 = result.data[0].url
89
- import requests
90
- response = requests.get(image_base64)
91
- image = Image.open(io.BytesIO(response.content))
92
- image.save('generated_image.png')
158
+
159
+ collected_images = []
160
+ for item_data in result.data:
161
+ if model == 'gpt-image-1':
162
+ image_base64 = item_data.b64_json
163
+ image_bytes = base64.b64decode(image_base64)
164
+ image = Image.open(io.BytesIO(image_bytes))
165
+ elif model == 'dall-e-2' or model == 'dall-e-3':
166
+ image_url = item_data.url
167
+ response = requests.get(image_url)
168
+ response.raise_for_status()
169
+ image = Image.open(io.BytesIO(response.content))
170
+ else:
171
+ image = item_data
172
+ collected_images.append(image)
93
173
 
94
- return image
174
+ return collected_images
175
+
95
176
 
96
177
 
97
178
  def gemini_image_gen(
98
179
  prompt: str,
99
- model: str = "imagen-3.0-generate-002",
180
+ model: str = "gemini-2.5-flash",
100
181
  attachments: Union[List[Union[str, bytes, Image.Image]], None] = None,
101
182
  n_images: int = 1,
102
183
  api_key: Optional[str] = None,
@@ -105,61 +186,114 @@ def gemini_image_gen(
105
186
  from google import genai
106
187
  from google.genai import types
107
188
  from io import BytesIO
189
+ import requests
108
190
 
109
- # Use environment variable if api_key not provided
110
191
  if api_key is None:
111
192
  api_key = os.environ.get('GEMINI_API_KEY')
112
193
 
113
194
  client = genai.Client(api_key=api_key)
114
-
195
+ collected_images = []
196
+
115
197
  if attachments is not None:
116
- # Process the attachments
117
- processed_images = []
198
+ processed_contents = [prompt]
199
+
118
200
  for attachment in attachments:
119
201
  if isinstance(attachment, str):
120
- # Assume it's a file path
121
- processed_images.append(Image.open(attachment))
202
+ if attachment.startswith(('http://', 'https://')):
203
+ image_bytes = requests.get(attachment).content
204
+ mime_type = 'image/jpeg'
205
+ processed_contents.append(
206
+ types.Part.from_bytes(
207
+ data=image_bytes,
208
+ mime_type=mime_type
209
+ )
210
+ )
211
+ else:
212
+ with open(attachment, 'rb') as f:
213
+ image_bytes = f.read()
214
+
215
+ if attachment.lower().endswith('.png'):
216
+ mime_type = 'image/png'
217
+ elif attachment.lower().endswith('.jpg') or attachment.lower().endswith('.jpeg'):
218
+ mime_type = 'image/jpeg'
219
+ elif attachment.lower().endswith('.webp'):
220
+ mime_type = 'image/webp'
221
+ else:
222
+ mime_type = 'image/jpeg'
223
+
224
+ processed_contents.append(
225
+ types.Part.from_bytes(
226
+ data=image_bytes,
227
+ mime_type=mime_type
228
+ )
229
+ )
122
230
  elif isinstance(attachment, bytes):
123
- processed_images.append(Image.open(BytesIO(attachment)))
231
+ processed_contents.append(
232
+ types.Part.from_bytes(
233
+ data=attachment,
234
+ mime_type='image/jpeg'
235
+ )
236
+ )
124
237
  elif isinstance(attachment, Image.Image):
125
- processed_images.append(attachment)
238
+ img_byte_arr = BytesIO()
239
+ attachment.save(img_byte_arr, format='PNG')
240
+ img_byte_arr.seek(0)
241
+ processed_contents.append(
242
+ types.Part.from_bytes(
243
+ data=img_byte_arr.getvalue(),
244
+ mime_type='image/png'
245
+ )
246
+ )
126
247
 
127
- # Use generate_content for image editing with gemini-2.0-flash-exp-image-generation
128
248
  response = client.models.generate_content(
129
- model="gemini-2.0-flash-exp-image-generation",
130
- contents=[prompt, processed_images],
131
- config=types.GenerateContentConfig(
132
- response_modalities=['TEXT', 'IMAGE']
133
- )
249
+ model=model,
250
+ contents=processed_contents,
134
251
  )
135
252
 
136
- # Process the response
137
- images = []
138
- for part in response.candidates[0].content.parts:
139
- if part.inline_data is not None:
140
- image = Image.open(BytesIO(part.inline_data.data))
141
- images.append(image)
253
+ if hasattr(response, 'candidates') and response.candidates:
254
+ for candidate in response.candidates:
255
+ for part in candidate.content.parts:
256
+ if hasattr(part, 'inline_data') and part.inline_data:
257
+ image_data = part.inline_data.data
258
+ collected_images.append(Image.open(BytesIO(image_data)))
259
+
260
+ if not collected_images and hasattr(response, 'text'):
261
+ print(f"Gemini response text: {response.text}")
142
262
 
143
- return images[0] if images else None
263
+ return collected_images
144
264
  else:
145
- # Use generate_images for new image generation with imagen model
146
- response = client.models.generate_images(
147
- model=model,
148
- prompt=prompt,
149
- config=types.GenerateImagesConfig(
150
- number_of_images=n_images,
265
+ if 'imagen' in model:
266
+ response = client.models.generate_images(
267
+ model=model,
268
+ prompt=prompt,
269
+ config=types.GenerateImagesConfig(
270
+ number_of_images=n_images,
271
+ )
151
272
  )
152
- )
153
-
154
- # Process the response
155
- images = []
156
- for generated_image in response.generated_images:
157
- image = Image.open(BytesIO(generated_image.image.image_bytes))
158
- images.append(image)
273
+ for generated_image in response.generated_images:
274
+ collected_images.append(Image.open(BytesIO(generated_image.image.image_bytes)))
275
+ return collected_images
276
+
277
+ elif 'flash-image' in model or 'image-preview' in model or '2.5-flash' in model:
278
+ response = client.models.generate_content(
279
+ model=model,
280
+ contents=[prompt],
281
+ )
282
+
283
+ if hasattr(response, 'candidates') and response.candidates:
284
+ for candidate in response.candidates:
285
+ for part in candidate.content.parts:
286
+ if hasattr(part, 'inline_data') and part.inline_data:
287
+ image_data = part.inline_data.data
288
+ collected_images.append(Image.open(BytesIO(image_data)))
289
+
290
+ if not collected_images and hasattr(response, 'text'):
291
+ print(f"Gemini response text: {response.text}")
292
+
293
+ return collected_images
159
294
 
160
- return images[0] if images else None
161
-
162
-
295
+ else:
296
+ raise ValueError(f"Unsupported Gemini image model or API usage for new generation: '{model}'")
163
297
  def generate_image(
164
298
  prompt: str,
165
299
  model: str ,
@@ -188,27 +322,37 @@ def generate_image(
188
322
  save_path (str): Path to save the generated image.
189
323
 
190
324
  Returns:
191
- PIL.Image: The generated image.
325
+ List[PIL.Image.Image]: A list of generated PIL Image objects.
192
326
  """
193
- # Set default model if none provided
327
+ from urllib.request import urlopen
328
+
194
329
  if model is None:
195
330
  if provider == "openai":
196
331
  model = "dall-e-2"
197
332
  elif provider == "diffusers":
198
333
  model = "runwayml/stable-diffusion-v1-5"
199
334
  elif provider == "gemini":
200
- model = "imagen-3.0-generate-002"
335
+ model = "gemini-2.5-flash-image-preview"
201
336
 
202
- # Generate or edit the image based on provider
337
+ all_generated_pil_images = []
338
+
203
339
  if provider == "diffusers":
204
- image = generate_image_diffusers(
205
- prompt=prompt,
206
- model=model,
207
- height=height,
208
- width=width
209
- )
340
+ for _ in range(n_images):
341
+ try:
342
+ image = generate_image_diffusers(
343
+ prompt=prompt,
344
+ model=model,
345
+ height=height,
346
+ width=width
347
+ )
348
+ all_generated_pil_images.append(image)
349
+ except MemoryError as e:
350
+ raise e
351
+ except Exception as e:
352
+ raise e
353
+
210
354
  elif provider == "openai":
211
- image = openai_image_gen(
355
+ images = openai_image_gen(
212
356
  prompt=prompt,
213
357
  model=model,
214
358
  attachments=attachments,
@@ -216,29 +360,26 @@ def generate_image(
216
360
  width=width,
217
361
  n_images=n_images
218
362
  )
363
+ all_generated_pil_images.extend(images)
364
+
219
365
  elif provider == "gemini":
220
- image = gemini_image_gen(
366
+ images = gemini_image_gen(
221
367
  prompt=prompt,
222
368
  model=model,
223
369
  attachments=attachments,
224
370
  n_images=n_images,
225
371
  api_key=api_key
226
372
  )
373
+ all_generated_pil_images.extend(images)
374
+
227
375
  else:
228
- # Validate image size for litellm
229
376
  valid_sizes = ["256x256", "512x512", "1024x1024", "1024x1792", "1792x1024"]
230
- size = f"{height}x{width}"
231
- if size not in valid_sizes:
232
- raise ValueError(
233
- f"Invalid image size: {size}. Please use one of the following: {', '.join(valid_sizes)}"
234
- )
377
+ size = f"{width}x{height}"
235
378
 
236
- # litellm doesn't support editing yet, so raise an error if attachments provided
237
379
  if attachments is not None:
238
380
  raise ValueError("Image editing not supported with litellm provider")
239
381
 
240
- # Generate image using litellm
241
- result = image_generation(
382
+ image_response = image_generation(
242
383
  prompt=prompt,
243
384
  model=f"{provider}/{model}",
244
385
  n=n_images,
@@ -246,17 +387,25 @@ def generate_image(
246
387
  api_key=api_key,
247
388
  api_base=api_url,
248
389
  )
249
- # Convert the URL result to a PIL image
250
- # This assumes image_generation returns a URL
251
- from urllib.request import urlopen
252
- with urlopen(result) as response:
253
- image = Image.open(response)
254
-
255
- # Save the image if a save path is provided
256
- if save_path and image:
257
- image.save(save_path)
258
-
259
- return image
390
+ for item_data in image_response.data:
391
+ if hasattr(item_data, 'url') and item_data.url:
392
+ with urlopen(item_data.url) as response:
393
+ all_generated_pil_images.append(Image.open(response))
394
+ elif hasattr(item_data, 'b64_json') and item_data.b64_json:
395
+ image_bytes = base64.b64decode(item_data.b64_json)
396
+ all_generated_pil_images.append(Image.open(io.BytesIO(image_bytes)))
397
+ else:
398
+ print(f"Warning: litellm ImageResponse item has no URL or b64_json: {item_data}")
399
+
400
+ if save_path:
401
+ for i, img_item in enumerate(all_generated_pil_images):
402
+ temp_save_path = f"{os.path.splitext(save_path)[0]}_{i}{os.path.splitext(save_path)[1]}"
403
+ if isinstance(img_item, Image.Image):
404
+ img_item.save(temp_save_path)
405
+ else:
406
+ print(f"Warning: Attempting to save non-PIL image item: {type(img_item)}. Skipping save for this item.")
407
+
408
+ return all_generated_pil_images
260
409
 
261
410
 
262
411
  def edit_image(
@@ -285,14 +434,12 @@ def edit_image(
285
434
  Returns:
286
435
  PIL.Image: The edited image.
287
436
  """
288
- # Set default model based on provider
289
437
  if model is None:
290
438
  if provider == "openai":
291
439
  model = "gpt-image-1"
292
440
  elif provider == "gemini":
293
- model = "gemini-2.0-flash-exp-image-generation"
441
+ model = "gemini-2.5-flash-image-preview"
294
442
 
295
- # Use the generate_image function with attachments
296
443
  image = generate_image(
297
444
  prompt=prompt,
298
445
  provider=provider,
@@ -304,6 +451,4 @@ def edit_image(
304
451
  api_key=api_key
305
452
  )
306
453
 
307
- return image
308
-
309
-
454
+ return image