Update app.py

app.py

@@ -5,6 +5,7 @@ import time
 import random
 import logging
 import sys
+import asyncio
 
 app = Flask(__name__)
 
@@ -133,60 +134,115 @@ def handle_request():
             
             response.raise_for_status()
             response_body = response.json()
-            image_url = response_body['images'][0]['url']
+            
+            if 'images' in response_body and response_body['images'] and 'url' in response_body['images'][0]:
+                image_url = response_body['images'][0]['url']
+                logger.info(f"Successfully retrieved image URL: {image_url}")
+            else:
+                logger.error(f"Unexpected response structure: {response_body}")
+                return jsonify({"error": "Unexpected response structure from image generation API"}), 500
         except requests.exceptions.RequestException as e:
             logger.error(f"Error in image generation request: {str(e)}")
             return jsonify({"error": "Failed to generate image"}), 500
+        except (KeyError, IndexError, ValueError) as e:
+            logger.error(f"Error parsing image generation response: {str(e)}")
+            return jsonify({"error": "Failed to parse image generation response"}), 500
         
         unique_id = int(time.time() * 1000)
         current_timestamp = unique_id // 1000
+        system_fingerprint = "fp_" + ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=9))
+        
+        image_data = {'data': [{'url': image_url}]}
         
         if stream:
-            response_payload = {
-                "id": unique_id,
-                "object": "chat.completion.chunk",
-                "created": current_timestamp,
-                "model": model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "delta": {
-                            "content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({image_url})"
-                        },
-                        "finish_reason": "stop"
-                    }
-                ]
-            }
-            data_string = json.dumps(response_payload)
-            return Response(f"data: {data_string}\n\n", content_type='text/event-stream')
+            return stream_response(request, unique_id, image_data, prompt, enhanced_prompt, "1024x1024", current_timestamp, model, system_fingerprint)
         else:
-            response_payload = {
-                "id": unique_id,
-                "object": "chat.completion",
-                "created": current_timestamp,
-                "model": model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({image_url})"
-                        },
-                        "logprobs": None,
-                        "finish_reason": "length"
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": len(prompt),
-                    "completion_tokens": len(enhanced_prompt) + len(image_url),
-                    "total_tokens": len(prompt) + len(enhanced_prompt) + len(image_url)
-                }
-            }
-            data_string = json.dumps(response_payload)
-            return Response(f"{data_string}\n\n", content_type='text/event-stream')
+            return non_stream_response(unique_id, image_data, prompt, enhanced_prompt, "1024x1024", current_timestamp, model, system_fingerprint)
     except Exception as e:
         logger.error(f"Unexpected error in handle_request: {str(e)}")
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
+def stream_response(request, unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint):
+    logger.debug("Starting stream response")
+    response = Response(stream_with_context(generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint)), content_type='text/event-stream')
+    logger.debug("Stream response completed")
+    return response
+
+def generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint):
+    chunks = [
+        f"原始提示词：\n{original_prompt}\n",
+        f"翻译后的提示词：\n{translated_prompt}\n",
+        f"图像规格：{size}\n",
+        "正在根据提示词生成图像...\n",
+        "图像正在处理中...\n",
+        "即将完成...\n",
+        f"生成成功！\n图像生成完毕，以下是结果：\n\n![生成的图像]({image_data['data'][0]['url']})"
+    ]
+
+    for i, chunk in enumerate(chunks):
+        json_chunk = json.dumps({
+            "id": unique_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model,
+            "system_fingerprint": system_fingerprint,
+            "choices": [{
+                "index": 0,
+                "delta": {"content": chunk},
+                "logprobs": None,
+                "finish_reason": None
+            }]
+        })
+        yield f"data: {json_chunk}\n\n"
+        time.sleep(0.5)  # 模拟生成时间
+
+    final_chunk = json.dumps({
+        "id": unique_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "system_fingerprint": system_fingerprint,
+        "choices": [{
+            "index": 0,
+            "delta": {},
+            "logprobs": None,
+            "finish_reason": "stop"
+        }]
+    })
+    yield f"data: {final_chunk}\n\n"
+
+def non_stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint):
+    content = (
+        f"原始提示词：{original_prompt}\n"
+        f"翻译后的提示词：{translated_prompt}\n"
+        f"图像规格：{size}\n"
+        f"图像生成成功！\n"
+        f"以下是结果：\n\n"
+        f"![生成的图像]({image_data['data'][0]['url']})"
+    )
+
+    response = {
+        'id': unique_id,
+        'object': "chat.completion",
+        'created': created,
+        'model': model,
+        'system_fingerprint': system_fingerprint,
+        'choices': [{
+            'index': 0,
+            'message': {
+                'role': "assistant",
+                'content': content
+            },
+            'finish_reason': "stop"
+        }],
+        'usage': {
+            'prompt_tokens': len(original_prompt),
+            'completion_tokens': len(content),
+            'total_tokens': len(original_prompt) + len(content)
+        }
+    }
+
+    return jsonify(response)
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=8000)