Update app.py

app.py

@@ -621,80 +621,129 @@ def handsome_chat_completions():
                 "https://api.siliconflow.cn/v1/images/generations",
                 headers=headers,
                 json=siliconflow_data,
-                timeout=120
+                timeout=120,
+                stream=data.get("stream", False)
             )
            
             if response.status_code == 429:
                 return jsonify(response.json()), 429
 
-            response.raise_for_status()
-            end_time = time.time()
-            response_json = response.json()
-            total_time = end_time - start_time
+            if data.get("stream", False):
+                def generate():
+                    first_chunk_time = None
+                    full_response_content = ""
+                    for chunk in response.iter_content(chunk_size=1024):
+                        if chunk:
+                            if first_chunk_time is None:
+                                first_chunk_time = time.time()
+                            try:
+                                chunk_json = json.loads(chunk.decode("utf-8"))
+                                if "images" in chunk_json and chunk_json["images"]:
+                                    image_url = chunk_json["images"][0].get("url", "")
+                                    if image_url:
+                                        full_response_content += json.dumps({"url":image_url})
+                                        yield f"data: {json.dumps({'url': image_url})}\n\n".encode('utf-8')
+                                    else:
+                                        full_response_content += json.dumps({"content":"Failed to generate image"})
+                                        yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                                else:
+                                    full_response_content += json.dumps({"content":"Failed to generate image"})
+                                    yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                            except json.JSONDecodeError:
+                                logging.error(f"Failed to decode chunk JSON: {chunk.decode('utf-8')}")
+                                full_response_content += json.dumps({"content":"Failed to generate image"})
+                                yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                                
+                    end_time = time.time()
+                    first_token_time = (
+                        first_chunk_time - start_time
+                        if first_chunk_time else 0
+                    )
+                    total_time = end_time - start_time
+                    
+                    
+                    logging.info(
+                        f"使用的key: {api_key}, "
+                        f"首字用时: {first_token_time:.4f}秒, "
+                        f"总共用时: {total_time:.4f}秒, "
+                        f"使用的模型: {model_name}"
+                    )
+
+                    with data_lock:
+                        request_timestamps.append(time.time())
+                        token_counts.append(0)  # Image generation doesn't use tokens
+                    
+                    yield "data: [DONE]\n\n".encode('utf-8')
+                return Response(stream_with_context(generate()), content_type='text/event-stream')
+            else:
+                response.raise_for_status()
+                end_time = time.time()
+                response_json = response.json()
+                total_time = end_time - start_time
             
-            try:
-                images = response_json.get("images", [])
-                
-                # Extract the first URL if available
-                image_url = ""
-                if images and isinstance(images[0], dict) and "url" in images[0]:
-                    image_url = images[0]["url"]
-                    logging.info(f"Extracted image URL: {image_url}")
-                elif images and isinstance(images[0], str):
-                  image_url = images[0]
-                  logging.info(f"Extracted image URL: {image_url}")
-                
-                # Construct the expected JSON output - Mimicking OpenAI
-                response_data = {
-                   "id": f"chatcmpl-{uuid.uuid4()}",
-                   "object": "chat.completion",
-                   "created": int(time.time()),
-                   "model": model_name,
-                   "choices": [
-                      {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": image_url if image_url else "Failed to generate image", # Directly return the URL in content
-                        },
-                        "finish_reason": "stop",
-                     }
-                   ],
-                }
+                try:
+                    images = response_json.get("images", [])
+                    
+                    # Extract the first URL if available
+                    image_url = ""
+                    if images and isinstance(images[0], dict) and "url" in images[0]:
+                        image_url = images[0]["url"]
+                        logging.info(f"Extracted image URL: {image_url}")
+                    elif images and isinstance(images[0], str):
+                        image_url = images[0]
+                        logging.info(f"Extracted image URL: {image_url}")
+                    
+                    # Construct the expected JSON output - Mimicking OpenAI
+                    response_data = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(time.time()),
+                        "model": model_name,
+                        "choices": [
+                            {
+                            "index": 0,
+                            "message": {
+                                "role": "assistant",
+                                "content": image_url if image_url else "Failed to generate image", # Directly return the URL in content
+                            },
+                            "finish_reason": "stop",
+                            }
+                        ],
+                    }
 
-            except (KeyError, ValueError, IndexError) as e:
-                logging.error(
-                    f"解析响应 JSON 失败: {e}, "
-                    f"完整内容: {response_json}"
-                )
-                response_data = {
-                   "id": f"chatcmpl-{uuid.uuid4()}",
-                   "object": "chat.completion",
-                   "created": int(time.time()),
-                   "model": model_name,
-                   "choices": [
-                      {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": "Failed to process image data",
-                        },
-                        "finish_reason": "stop",
-                      }
-                   ],
-                }
+                except (KeyError, ValueError, IndexError) as e:
+                    logging.error(
+                        f"解析响应 JSON 失败: {e}, "
+                        f"完整内容: {response_json}"
+                    )
+                    response_data = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(time.time()),
+                        "model": model_name,
+                        "choices": [
+                            {
+                            "index": 0,
+                            "message": {
+                                "role": "assistant",
+                                "content": "Failed to process image data",
+                            },
+                            "finish_reason": "stop",
+                            }
+                        ],
+                    }
 
-            logging.info(
-                f"使用的key: {api_key}, "
-                f"总共用时: {total_time:.4f}秒, "
-                f"使用的模型: {model_name}"
-            )
+                logging.info(
+                    f"使用的key: {api_key}, "
+                    f"总共用时: {total_time:.4f}秒, "
+                    f"使用的模型: {model_name}"
+                )
 
-            with data_lock:
-                request_timestamps.append(time.time())
-                token_counts.append(0)  # Image generation doesn't use tokens
+                with data_lock:
+                    request_timestamps.append(time.time())
+                    token_counts.append(0)  # Image generation doesn't use tokens
 
-            return jsonify(response_data)
+                return jsonify(response_data)
         except requests.exceptions.RequestException as e:
             logging.error(f"请求转发异常: {e}")
             return jsonify({"error": str(e)}), 500