Remove deepcopy usage in PromptQueue

This commit is contained in:
catboxanon 2025-02-10 15:12:54 -05:00
parent 095d867147
commit 9ab14711e5

View File

@ -1,5 +1,4 @@
import sys import sys
import copy
import logging import logging
import threading import threading
import heapq import heapq
@ -880,6 +879,16 @@ class PromptQueue:
self.flags = {} self.flags = {}
server.prompt_queue = self server.prompt_queue = self
def _copy(self, obj):
if isinstance(obj, dict):
return {k: self._copy(v) for k, v in obj.items()}
elif isinstance(obj, tuple):
return tuple(self._copy(v) for v in obj)
elif isinstance(obj, list):
return [self._copy(v) for v in obj]
else:
return obj
def put(self, item): def put(self, item):
with self.mutex: with self.mutex:
heapq.heappush(self.queue, item) heapq.heappush(self.queue, item)
@ -894,7 +903,7 @@ class PromptQueue:
return None return None
item = heapq.heappop(self.queue) item = heapq.heappop(self.queue)
i = self.task_counter i = self.task_counter
self.currently_running[i] = copy.deepcopy(item) self.currently_running[i] = self._copy(item)
self.task_counter += 1 self.task_counter += 1
self.server.queue_updated() self.server.queue_updated()
return (item, i) return (item, i)
@ -913,7 +922,7 @@ class PromptQueue:
status_dict: Optional[dict] = None status_dict: Optional[dict] = None
if status is not None: if status is not None:
status_dict = copy.deepcopy(status._asdict()) status_dict = self._copy(status._asdict())
self.history[prompt[1]] = { self.history[prompt[1]] = {
"prompt": prompt, "prompt": prompt,
@ -928,7 +937,7 @@ class PromptQueue:
out = [] out = []
for x in self.currently_running.values(): for x in self.currently_running.values():
out += [x] out += [x]
return (out, copy.deepcopy(self.queue)) return (out, self._copy(self.queue))
def get_tasks_remaining(self): def get_tasks_remaining(self):
with self.mutex: with self.mutex:
@ -967,7 +976,7 @@ class PromptQueue:
i += 1 i += 1
return out return out
elif prompt_id in self.history: elif prompt_id in self.history:
return {prompt_id: copy.deepcopy(self.history[prompt_id])} return {prompt_id: self._copy(self.history[prompt_id])}
else: else:
return {} return {}