ChenyuRabbitLove commited on
Commit
c74caed
·
1 Parent(s): fef0c2e

refactor: update image generation response structure to return base64 images, simplify Dockerfile permissions, and clean up unused code in image service

Browse files
Dockerfile CHANGED
@@ -46,11 +46,9 @@ COPY --from=frontend-builder --chown=nextjs:nodejs /app/.next/cache ./frontend/.
46
  COPY --from=backend /backend ./backend
47
  COPY --from=backend /usr/local/lib/python3.9/site-packages /usr/local/lib/python3.9/site-packages
48
 
49
- # Create directories and set permissions
50
- RUN mkdir -p /app/generated_images && \
51
- chown -R nextjs:nodejs /app/backend /app/generated_images && \
52
- chmod -R 755 /app/backend && \
53
- chmod 755 /app/generated_images
54
 
55
  COPY start.sh .
56
  RUN chmod +x start.sh
 
46
  COPY --from=backend /backend ./backend
47
  COPY --from=backend /usr/local/lib/python3.9/site-packages /usr/local/lib/python3.9/site-packages
48
 
49
+ # Set permissions for backend
50
+ RUN chown -R nextjs:nodejs /app/backend && \
51
+ chmod -R 755 /app/backend
 
 
52
 
53
  COPY start.sh .
54
  RUN chmod +x start.sh
app/page.tsx CHANGED
@@ -25,8 +25,7 @@ interface ImageGenerationRequest {
25
  interface ImageGenerationResponse {
26
  success: boolean;
27
  message: string;
28
- filename?: string;
29
- filenames?: string[];
30
  count?: number;
31
  }
32
 
@@ -153,10 +152,10 @@ export default function Home() {
153
 
154
  const data: ImageGenerationResponse = await response.json();
155
 
156
- if (data.success && (data.filenames || data.filename)) {
157
- const filenames = data.filenames || (data.filename ? [data.filename] : []);
158
- const imageUrls = filenames.map(filename =>
159
- `/api/v1/images/download/${filename}`
160
  );
161
  setGeneratedImages(imageUrls);
162
  setLoadingImages([]);
@@ -187,7 +186,7 @@ export default function Home() {
187
  };
188
 
189
  const downloadAllImages = () => {
190
- generatedImages.forEach((imageUrl, index) => {
191
  const link = document.createElement('a');
192
  link.href = imageUrl;
193
  link.download = `fox-cat-${index + 1}.png`;
 
25
  interface ImageGenerationResponse {
26
  success: boolean;
27
  message: string;
28
+ images?: string[];
 
29
  count?: number;
30
  }
31
 
 
152
 
153
  const data: ImageGenerationResponse = await response.json();
154
 
155
+ if (data.success && data.images && data.images.length > 0) {
156
+ // Convert base64 data to data URLs for display
157
+ const imageUrls = data.images.map(base64Data =>
158
+ `data:image/png;base64,${base64Data}`
159
  );
160
  setGeneratedImages(imageUrls);
161
  setLoadingImages([]);
 
186
  };
187
 
188
  const downloadAllImages = () => {
189
+ generatedImages.forEach((imageUrl: string, index: number) => {
190
  const link = document.createElement('a');
191
  link.href = imageUrl;
192
  link.download = `fox-cat-${index + 1}.png`;
backend/app/api/v1/endpoints/images.py CHANGED
@@ -1,6 +1,4 @@
1
  from fastapi import APIRouter, HTTPException, status
2
- from fastapi.responses import FileResponse
3
- import os
4
  import logging
5
 
6
  from ....schemas.image import ImageGenerationRequest, ImageGenerationResponse
@@ -15,6 +13,7 @@ router = APIRouter(prefix="/images", tags=["images"])
15
  async def generate_image(request: ImageGenerationRequest):
16
  """
17
  Generate an image using OpenAI DALL-E, optionally with a reference image
 
18
  """
19
  try:
20
  result = await image_service.generate_image(
@@ -34,8 +33,7 @@ async def generate_image(request: ImageGenerationRequest):
34
  return ImageGenerationResponse(
35
  success=True,
36
  message=result["message"],
37
- filename=result["filename"],
38
- filenames=result.get("filenames", []),
39
  count=result.get("count", 0),
40
  )
41
 
@@ -47,28 +45,3 @@ async def generate_image(request: ImageGenerationRequest):
47
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
48
  detail="An unexpected error occurred while generating the image",
49
  )
50
-
51
-
52
- @router.get("/download/{filename}")
53
- async def download_image(filename: str):
54
- """
55
- Download a generated image by filename
56
- """
57
- try:
58
- filepath = os.path.join(image_service.output_dir, filename)
59
-
60
- if not os.path.exists(filepath):
61
- raise HTTPException(
62
- status_code=status.HTTP_404_NOT_FOUND, detail="Image not found"
63
- )
64
-
65
- return FileResponse(filepath, media_type="image/png", filename=filename)
66
-
67
- except HTTPException:
68
- raise
69
- except Exception as e:
70
- logger.error(f"Error downloading image: {str(e)}")
71
- raise HTTPException(
72
- status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
73
- detail="An error occurred while downloading the image",
74
- )
 
1
  from fastapi import APIRouter, HTTPException, status
 
 
2
  import logging
3
 
4
  from ....schemas.image import ImageGenerationRequest, ImageGenerationResponse
 
13
  async def generate_image(request: ImageGenerationRequest):
14
  """
15
  Generate an image using OpenAI DALL-E, optionally with a reference image
16
+ Returns base64 encoded image data directly
17
  """
18
  try:
19
  result = await image_service.generate_image(
 
33
  return ImageGenerationResponse(
34
  success=True,
35
  message=result["message"],
36
+ images=result.get("images", []),
 
37
  count=result.get("count", 0),
38
  )
39
 
 
45
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
46
  detail="An unexpected error occurred while generating the image",
47
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
backend/app/schemas/image.py CHANGED
@@ -30,7 +30,7 @@ class ImageGenerationResponse(BaseModel):
30
 
31
  success: bool
32
  message: str
33
- image_url: Optional[str] = None
34
- filename: Optional[str] = None
35
- filenames: Optional[List[str]] = None
36
  count: Optional[int] = None
 
30
 
31
  success: bool
32
  message: str
33
+ images: Optional[List[str]] = Field(
34
+ default=None, description="List of base64 encoded generated images"
35
+ )
36
  count: Optional[int] = None
backend/app/services/image_service.py CHANGED
@@ -19,13 +19,6 @@ class ImageGenerationService:
19
  timeout=60.0, # Increase timeout for Hugging Face environment
20
  max_retries=2, # Reduce retries to fail faster
21
  )
22
- self.output_dir = "generated_images"
23
- self._ensure_output_directory()
24
-
25
- def _ensure_output_directory(self):
26
- """Ensure the output directory exists"""
27
- if not os.path.exists(self.output_dir):
28
- os.makedirs(self.output_dir)
29
 
30
  async def _fallback_to_dalle(
31
  self, prompt: str, size: str, n: int, model: str
@@ -45,33 +38,25 @@ class ImageGenerationService:
45
  response_format="b64_json",
46
  )
47
 
48
- generated_filenames = []
49
 
50
  for i, image_data in enumerate(response.data):
51
  try:
52
- image_bytes = base64.b64decode(image_data.b64_json)
53
-
54
- # Generate unique filename and save
55
- filename = f"{uuid.uuid4()}.png"
56
- filepath = os.path.join(self.output_dir, filename)
57
-
58
- with open(filepath, "wb") as f:
59
- f.write(image_bytes)
60
-
61
- generated_filenames.append(filename)
62
- logger.info(f"Fallback image {i+1} saved successfully: {filename}")
63
 
64
  except Exception as e:
65
- logger.warning(f"Failed to save fallback image {i+1}: {str(e)}")
66
  continue
67
 
68
- if generated_filenames:
69
  return {
70
  "success": True,
71
- "message": f"Generated {len(generated_filenames)}/{n} images using DALL-E fallback (reference image ignored due to network restrictions)",
72
- "filename": generated_filenames[0],
73
- "filenames": generated_filenames,
74
- "count": len(generated_filenames),
75
  }
76
  else:
77
  raise Exception("Fallback also failed to generate any images")
@@ -99,7 +84,7 @@ class ImageGenerationService:
99
  reference_image: Base64 encoded reference image (optional)
100
 
101
  Returns:
102
- dict: Result containing success status, message, and filename(s)
103
  """
104
  try:
105
  logger.info(f"Generating {n} image(s) with prompt: {prompt}")
@@ -108,7 +93,7 @@ class ImageGenerationService:
108
  # Use the newer responses API with image generation tools for reference images
109
  logger.info("Using reference image with responses API")
110
 
111
- generated_filenames = []
112
 
113
  # Generate multiple images by making multiple requests
114
  for i in range(n):
@@ -154,18 +139,9 @@ class ImageGenerationService:
154
  )
155
  continue
156
 
157
- # Decode base64 image
158
- image_bytes = base64.b64decode(image_data)
159
-
160
- # Generate unique filename and save
161
- filename = f"{uuid.uuid4()}.png"
162
- filepath = os.path.join(self.output_dir, filename)
163
-
164
- with open(filepath, "wb") as f:
165
- f.write(image_bytes)
166
-
167
- generated_filenames.append(filename)
168
- logger.info(f"Image {i+1} saved successfully: {filename}")
169
 
170
  except Exception as e:
171
  error_msg = str(e)
@@ -192,7 +168,7 @@ class ImageGenerationService:
192
 
193
  continue
194
 
195
- if not generated_filenames:
196
  # If responses API failed due to network restrictions, try fallback to regular DALL-E
197
  logger.warning(
198
  "Responses API failed, attempting fallback to regular DALL-E"
@@ -200,15 +176,14 @@ class ImageGenerationService:
200
  return await self._fallback_to_dalle(prompt, size, n, model)
201
 
202
  logger.info(
203
- f"Successfully generated {len(generated_filenames)}/{n} images"
204
  )
205
 
206
  return {
207
  "success": True,
208
- "message": f"Generated {len(generated_filenames)}/{n} images successfully",
209
- "filename": generated_filenames[0] if generated_filenames else None,
210
- "filenames": generated_filenames,
211
- "count": len(generated_filenames),
212
  }
213
 
214
  else:
@@ -223,32 +198,24 @@ class ImageGenerationService:
223
  response_format="b64_json",
224
  )
225
 
226
- generated_filenames = []
227
 
228
  for i, image_data in enumerate(response.data):
229
  try:
230
- image_bytes = base64.b64decode(image_data.b64_json)
231
-
232
- # Generate unique filename and save
233
- filename = f"{uuid.uuid4()}.png"
234
- filepath = os.path.join(self.output_dir, filename)
235
-
236
- with open(filepath, "wb") as f:
237
- f.write(image_bytes)
238
-
239
- generated_filenames.append(filename)
240
- logger.info(f"Image {i+1} saved successfully: {filename}")
241
 
242
  except Exception as e:
243
- logger.warning(f"Failed to save image {i+1}: {str(e)}")
244
  continue
245
 
246
  return {
247
  "success": True,
248
- "message": f"Generated {len(generated_filenames)}/{n} images successfully",
249
- "filename": generated_filenames[0] if generated_filenames else None,
250
- "filenames": generated_filenames,
251
- "count": len(generated_filenames),
252
  }
253
 
254
  except Exception as e:
@@ -256,8 +223,7 @@ class ImageGenerationService:
256
  return {
257
  "success": False,
258
  "message": f"Failed to generate image: {str(e)}",
259
- "filename": None,
260
- "filenames": [],
261
  "count": 0,
262
  }
263
 
 
19
  timeout=60.0, # Increase timeout for Hugging Face environment
20
  max_retries=2, # Reduce retries to fail faster
21
  )
 
 
 
 
 
 
 
22
 
23
  async def _fallback_to_dalle(
24
  self, prompt: str, size: str, n: int, model: str
 
38
  response_format="b64_json",
39
  )
40
 
41
+ generated_images = []
42
 
43
  for i, image_data in enumerate(response.data):
44
  try:
45
+ # Keep the base64 data instead of saving to file
46
+ image_base64 = image_data.b64_json
47
+ generated_images.append(image_base64)
48
+ logger.info(f"Fallback image {i+1} processed successfully")
 
 
 
 
 
 
 
49
 
50
  except Exception as e:
51
+ logger.warning(f"Failed to process fallback image {i+1}: {str(e)}")
52
  continue
53
 
54
+ if generated_images:
55
  return {
56
  "success": True,
57
+ "message": f"Generated {len(generated_images)}/{n} images using DALL-E fallback (reference image ignored due to network restrictions)",
58
+ "images": generated_images,
59
+ "count": len(generated_images),
 
60
  }
61
  else:
62
  raise Exception("Fallback also failed to generate any images")
 
84
  reference_image: Base64 encoded reference image (optional)
85
 
86
  Returns:
87
+ dict: Result containing success status, message, and base64 image data
88
  """
89
  try:
90
  logger.info(f"Generating {n} image(s) with prompt: {prompt}")
 
93
  # Use the newer responses API with image generation tools for reference images
94
  logger.info("Using reference image with responses API")
95
 
96
+ generated_images = []
97
 
98
  # Generate multiple images by making multiple requests
99
  for i in range(n):
 
139
  )
140
  continue
141
 
142
+ # Keep the base64 data instead of saving to file
143
+ generated_images.append(image_data)
144
+ logger.info(f"Image {i+1} processed successfully")
 
 
 
 
 
 
 
 
 
145
 
146
  except Exception as e:
147
  error_msg = str(e)
 
168
 
169
  continue
170
 
171
+ if not generated_images:
172
  # If responses API failed due to network restrictions, try fallback to regular DALL-E
173
  logger.warning(
174
  "Responses API failed, attempting fallback to regular DALL-E"
 
176
  return await self._fallback_to_dalle(prompt, size, n, model)
177
 
178
  logger.info(
179
+ f"Successfully generated {len(generated_images)}/{n} images"
180
  )
181
 
182
  return {
183
  "success": True,
184
+ "message": f"Generated {len(generated_images)}/{n} images successfully",
185
+ "images": generated_images,
186
+ "count": len(generated_images),
 
187
  }
188
 
189
  else:
 
198
  response_format="b64_json",
199
  )
200
 
201
+ generated_images = []
202
 
203
  for i, image_data in enumerate(response.data):
204
  try:
205
+ # Keep the base64 data instead of saving to file
206
+ image_base64 = image_data.b64_json
207
+ generated_images.append(image_base64)
208
+ logger.info(f"Image {i+1} processed successfully")
 
 
 
 
 
 
 
209
 
210
  except Exception as e:
211
+ logger.warning(f"Failed to process image {i+1}: {str(e)}")
212
  continue
213
 
214
  return {
215
  "success": True,
216
+ "message": f"Generated {len(generated_images)}/{n} images successfully",
217
+ "images": generated_images,
218
+ "count": len(generated_images),
 
219
  }
220
 
221
  except Exception as e:
 
223
  return {
224
  "success": False,
225
  "message": f"Failed to generate image: {str(e)}",
226
+ "images": [],
 
227
  "count": 0,
228
  }
229