hzruo commited on
Commit
1f3daf2
·
verified ·
1 Parent(s): 7f3fc58

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +24 -15
main.py CHANGED
@@ -407,12 +407,26 @@ async def process_image_generation(msg_data: str, session: requests.Session, hea
407
  job_id, prompt, negative = match.groups()
408
  print(f"Starting image generation process for job_id: {job_id}")
409
 
410
- # 构建思考消息的开始部分
411
- think_msg = "<think>\n"
412
- think_msg += "🎨 Generating image...\n\n"
413
- think_msg += f"Prompt: {prompt}\n"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
414
 
415
- # 记录开始时间(在检查图片状态前)
416
  start_time = time.time()
417
 
418
  # 检查图片状态和上传
@@ -421,27 +435,22 @@ async def process_image_generation(msg_data: str, session: requests.Session, hea
421
  # 计算实际经过的时间
422
  elapsed_time = time.time() - start_time
423
 
424
- # 完成思考部分
425
- think_msg += f"\n🤔 Thinking for {elapsed_time:.1f}s...\n"
426
- think_msg += "</think>"
427
-
428
- # 返回两个独立的消息块
429
- messages = []
430
 
431
- # 第一个消息块:思考过程
432
  messages.append({
433
- "id": f"chatcmpl-{chat_id}-think",
434
  "object": "chat.completion.chunk",
435
  "created": int(time.time()),
436
  "model": "AkashGen",
437
  "choices": [{
438
- "delta": {"content": think_msg},
439
  "index": 0,
440
  "finish_reason": None
441
  }]
442
  })
443
 
444
- # 第二个消息块:图片结果
445
  if result:
446
  image_msg = f"\n\n![Generated Image]({result})"
447
  messages.append({
 
407
  job_id, prompt, negative = match.groups()
408
  print(f"Starting image generation process for job_id: {job_id}")
409
 
410
+ # 返回多个消息块
411
+ messages = []
412
+
413
+ # 第一个消息块:开始思考
414
+ start_msg = "<think>\n🎨 Generating image...\n\n"
415
+ start_msg += f"Prompt: {prompt}\n"
416
+
417
+ messages.append({
418
+ "id": f"chatcmpl-{chat_id}-start",
419
+ "object": "chat.completion.chunk",
420
+ "created": int(time.time()),
421
+ "model": "AkashGen",
422
+ "choices": [{
423
+ "delta": {"content": start_msg},
424
+ "index": 0,
425
+ "finish_reason": None
426
+ }]
427
+ })
428
 
429
+ # 记录开始时间
430
  start_time = time.time()
431
 
432
  # 检查图片状态和上传
 
435
  # 计算实际经过的时间
436
  elapsed_time = time.time() - start_time
437
 
438
+ # 第二个消息块:结束思考
439
+ end_think_msg = f"\n🤔 Thinking for {elapsed_time:.1f}s...\n</think>"
 
 
 
 
440
 
 
441
  messages.append({
442
+ "id": f"chatcmpl-{chat_id}-end-think",
443
  "object": "chat.completion.chunk",
444
  "created": int(time.time()),
445
  "model": "AkashGen",
446
  "choices": [{
447
+ "delta": {"content": end_think_msg},
448
  "index": 0,
449
  "finish_reason": None
450
  }]
451
  })
452
 
453
+ # 第三个消息块:图片结果
454
  if result:
455
  image_msg = f"\n\n![Generated Image]({result})"
456
  messages.append({