server save img to memory

[black-eleven]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces an in-memory image result handling mechanism, allowing the synchronous API to return PNG bytes directly instead of streaming from disk. Key changes include adding a prefer_memory_result flag to task requests, implementing a PNG encoding utility for various pipeline output formats, and updating the task manager to store and retrieve in-memory results. Feedback highlights the need to clear image bytes from memory after retrieval to prevent potential Out-Of-Memory (OOM) errors and suggests improving the robustness of the image encoding logic to handle edge cases like empty lists or single image objects.

[gemini-code-assist]

The result_png field stores image bytes in memory. Since TaskManager retains up to 1000 tasks by default, frequent use of the sync API could lead to significant memory consumption and potential Out-Of-Memory (OOM) errors. It is recommended to clear the result_png data once it has been retrieved by the client.

Suggested change

	def get_task_result_png(self, task_id: str) -> Optional[bytes]:
	with self._lock:
	task = self._tasks.get(task_id)
	if not task:
	return None
	return task.result_png
	def get_task_result_png(self, task_id: str) -> Optional[bytes]:
	with self._lock:
	task = self._tasks.get(task_id)
	if not task:
	return None
	res = task.result_png
	task.result_png = None # Clear memory after retrieval
	return res

[gemini-code-assist]

The logic for handling dictionary returns from the pipeline assumes that images is either a torch.Tensor or a non-empty list of PIL images. If images is an empty list or a single PIL.Image object, it will raise an IndexError or TypeError. The suggested change adds robustness for these cases.

Suggested change

	if isinstance(images, torch.Tensor):
	return _tensor_to_png_bytes(images)
	return _pil_images_structure_to_png(images)
	if isinstance(pipeline_return, list) and len(pipeline_return) > 0:
	if isinstance(pipeline_return[0], torch.Tensor):
	return _tensor_to_png_bytes(pipeline_return[0])
	return _pil_images_structure_to_png(pipeline_return)
	if isinstance(pipeline_return, dict):
	images = pipeline_return.get("images")
	if images is None:
	return None
	if isinstance(images, torch.Tensor):
	return _tensor_to_png_bytes(images)
	if isinstance(images, (list, tuple)):
	return _pil_images_structure_to_png(images) if images else None
	if isinstance(images, Image.Image):
	return _pil_to_png_bytes(images)
	return None