Remove some test information

This commit is contained in:
Song
2023-09-09 22:07:31 +08:00
parent 3eac598c97
commit 0d793bd12d
2 changed files with 4 additions and 8 deletions
+2 -2
View File
@@ -23,7 +23,7 @@ class ComputeTaskType(Enum):
class ComputeTask:
def __init__(self) -> None:
self.task_type = "llm_completion"
self.task_type = ComputeTaskType.NONE
self.create_time = None
self.task_id: str = None
@@ -37,7 +37,7 @@ class ComputeTask:
self.error_str = None
def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None):
self.task_type = "llm_completion"
self.task_type = ComputeTaskType.LLM_COMPLETION
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
+2 -6
View File
@@ -31,21 +31,17 @@ class Stability_ComputeNode(ComputeNode):
self.is_start = True
self.node_id = "stability_node"
self.api_key = "" # "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0"
self.engine = "" # stable-diffusion-512-v2-0
self.api_key = ""
self.engine = "stable-diffusion-512-v2-1"
self.task_queue = Queue()
if os.getenv("STABILITY_API_KEY") is not None:
self.api_key = os.getenv("STABILITY_API_KEY")
else:
self.api_key = "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0"
# Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine
if os.getenv("STABILITY_ENGINE") is not None:
self.engine = os.getenv("STABILITY_ENGINE")
else:
self.engine = "stable-diffusion-512-v2-1"
self.client = client.StabilityInference(
key=self.api_key,