Make model.generate() result unpacking more robust to handle different return types

parent 72d01acd
...@@ -152,7 +152,12 @@ def analyze_single_image(image_path, prompt, model): ...@@ -152,7 +152,12 @@ def analyze_single_image(image_path, prompt, model):
} }
] ]
result, tokens_used = model.generate({"messages": messages}, max_new_tokens=128) gen_result = model.generate({"messages": messages}, max_new_tokens=128)
if isinstance(gen_result, tuple):
result, tokens_used = gen_result
else:
result = gen_result
tokens_used = 0
return result return result
# For now, estimate tokens (could be improved with actual token counting) # For now, estimate tokens (could be improved with actual token counting)
estimated_tokens = len(result.split()) + len(prompt.split()) estimated_tokens = len(result.split()) + len(prompt.split())
...@@ -270,10 +275,20 @@ def analyze_media(media_path, prompt, model_path, interval=10, job_id=None, comm ...@@ -270,10 +275,20 @@ def analyze_media(media_path, prompt, model_path, interval=10, job_id=None, comm
# Use vision model for summary # Use vision model for summary
summary_prompt = f"Summarize the video based on frame descriptions: {' '.join(descriptions)}" summary_prompt = f"Summarize the video based on frame descriptions: {' '.join(descriptions)}"
messages = [{"role": "user", "content": [{"type": "text", "text": summary_prompt}]}] messages = [{"role": "user", "content": [{"type": "text", "text": summary_prompt}]}]
summary, summary_tokens = model.generate({"messages": messages}, max_new_tokens=256) gen_result = model.generate({"messages": messages}, max_new_tokens=256)
if isinstance(gen_result, tuple):
summary, summary_tokens = gen_result
else:
summary = gen_result
summary_tokens = 0
else: else:
# Use text-only model for summary # Use text-only model for summary
summary, summary_tokens = model.generate(f"Summarize the video based on frame descriptions: {' '.join(descriptions)}", max_new_tokens=256) gen_result = model.generate(f"Summarize the video based on frame descriptions: {' '.join(descriptions)}", max_new_tokens=256)
if isinstance(gen_result, tuple):
summary, summary_tokens = gen_result
else:
summary = gen_result
summary_tokens = 0
total_tokens += summary_tokens total_tokens += summary_tokens
print(f"DEBUG: Summary generated for job {job_id}") print(f"DEBUG: Summary generated for job {job_id}")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment