Docfile commited on
Commit
6822891
·
verified ·
1 Parent(s): c8b9113

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -52,10 +52,13 @@ def generate():
52
  image.save(img_byte_arr, format='PNG')
53
  img_byte_arr = img_byte_arr.getvalue()
54
 
55
- # Stocker la tâche de génération dans le dictionnaire
 
 
 
56
  pending_responses[request_id] = {
57
  'status': 'processing',
58
- 'response': model.generate_content([mm, {"mime_type": "image/png", "data": img_byte_arr}], stream=True)
59
  }
60
 
61
  return jsonify({'request_id': request_id})
@@ -64,13 +67,14 @@ def generate():
64
  finally:
65
  # Nettoyer le fichier temporaire
66
  os.unlink(temp_file.name)
67
-
68
  @app.route('/stream/<request_id>')
69
  def stream(request_id):
70
  def generate_stream():
71
  while request_id in pending_responses and pending_responses[request_id]['status'] == 'processing':
72
  try:
73
- chunk = next(pending_responses[request_id]['response'])
 
74
  yield f"data: {chunk.text}\n\n"
75
  except StopIteration:
76
  pending_responses[request_id]['status'] = 'completed'
 
52
  image.save(img_byte_arr, format='PNG')
53
  img_byte_arr = img_byte_arr.getvalue()
54
 
55
+ # Générer le contenu en streaming et obtenir l'itérateur
56
+ response_iterator = model.generate_content([mm, {"mime_type": "image/png", "data": img_byte_arr}], stream=True)
57
+
58
+ # Stocker l'itérateur dans le dictionnaire
59
  pending_responses[request_id] = {
60
  'status': 'processing',
61
+ 'iterator': response_iterator
62
  }
63
 
64
  return jsonify({'request_id': request_id})
 
67
  finally:
68
  # Nettoyer le fichier temporaire
69
  os.unlink(temp_file.name)
70
+
71
  @app.route('/stream/<request_id>')
72
  def stream(request_id):
73
  def generate_stream():
74
  while request_id in pending_responses and pending_responses[request_id]['status'] == 'processing':
75
  try:
76
+ # Utiliser l'itérateur stocké pour obtenir le prochain chunk
77
+ chunk = next(pending_responses[request_id]['iterator'])
78
  yield f"data: {chunk.text}\n\n"
79
  except StopIteration:
80
  pending_responses[request_id]['status'] = 'completed'