hzruo commited on
Commit
86175fb
·
verified ·
1 Parent(s): 1f3daf2

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +48 -16
main.py CHANGED
@@ -410,17 +410,40 @@ async def process_image_generation(msg_data: str, session: requests.Session, hea
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
  }]
@@ -435,44 +458,53 @@ async def process_image_generation(msg_data: str, session: requests.Session, hea
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({
457
- "id": f"chatcmpl-{chat_id}-image",
458
  "object": "chat.completion.chunk",
459
  "created": int(time.time()),
460
  "model": "AkashGen",
461
  "choices": [{
462
- "delta": {"content": image_msg},
463
  "index": 0,
464
  "finish_reason": None
465
  }]
466
  })
467
  else:
468
- fail_msg = "\n\n*Image generation or upload failed.*"
469
  messages.append({
470
- "id": f"chatcmpl-{chat_id}-fail",
471
  "object": "chat.completion.chunk",
472
  "created": int(time.time()),
473
  "model": "AkashGen",
474
  "choices": [{
475
- "delta": {"content": fail_msg},
476
  "index": 0,
477
  "finish_reason": None
478
  }]
 
410
  # 返回多个消息块
411
  messages = []
412
 
413
+ # 发送思考开始标记
414
+ messages.append({
415
+ "id": f"chatcmpl-{chat_id}-1",
416
+ "object": "chat.completion.chunk",
417
+ "created": int(time.time()),
418
+ "model": "AkashGen",
419
+ "choices": [{
420
+ "delta": {"content": "<think>\n"},
421
+ "index": 0,
422
+ "finish_reason": None
423
+ }]
424
+ })
425
+
426
+ # 发送生成图片的提示
427
+ messages.append({
428
+ "id": f"chatcmpl-{chat_id}-2",
429
+ "object": "chat.completion.chunk",
430
+ "created": int(time.time()),
431
+ "model": "AkashGen",
432
+ "choices": [{
433
+ "delta": {"content": "🎨 Generating image...\n\n"},
434
+ "index": 0,
435
+ "finish_reason": None
436
+ }]
437
+ })
438
 
439
+ # 发送提示词
440
  messages.append({
441
+ "id": f"chatcmpl-{chat_id}-3",
442
  "object": "chat.completion.chunk",
443
  "created": int(time.time()),
444
  "model": "AkashGen",
445
  "choices": [{
446
+ "delta": {"content": f"Prompt: {prompt}\n\n"},
447
  "index": 0,
448
  "finish_reason": None
449
  }]
 
458
  # 计算实际经过的时间
459
  elapsed_time = time.time() - start_time
460
 
461
+ # 发送思考时间
462
+ messages.append({
463
+ "id": f"chatcmpl-{chat_id}-4",
464
+ "object": "chat.completion.chunk",
465
+ "created": int(time.time()),
466
+ "model": "AkashGen",
467
+ "choices": [{
468
+ "delta": {"content": f"🤔 Thinking for {elapsed_time:.1f}s...\n"},
469
+ "index": 0,
470
+ "finish_reason": None
471
+ }]
472
+ })
473
 
474
+ # 发送思考结束标记
475
  messages.append({
476
+ "id": f"chatcmpl-{chat_id}-5",
477
  "object": "chat.completion.chunk",
478
  "created": int(time.time()),
479
  "model": "AkashGen",
480
  "choices": [{
481
+ "delta": {"content": "</think>\n\n"},
482
  "index": 0,
483
  "finish_reason": None
484
  }]
485
  })
486
 
487
+ # 发送图片结果
488
  if result:
 
489
  messages.append({
490
+ "id": f"chatcmpl-{chat_id}-6",
491
  "object": "chat.completion.chunk",
492
  "created": int(time.time()),
493
  "model": "AkashGen",
494
  "choices": [{
495
+ "delta": {"content": f"![Generated Image]({result})"},
496
  "index": 0,
497
  "finish_reason": None
498
  }]
499
  })
500
  else:
 
501
  messages.append({
502
+ "id": f"chatcmpl-{chat_id}-6",
503
  "object": "chat.completion.chunk",
504
  "created": int(time.time()),
505
  "model": "AkashGen",
506
  "choices": [{
507
+ "delta": {"content": "*Image generation or upload failed.*"},
508
  "index": 0,
509
  "finish_reason": None
510
  }]