Spaces:
Paused
Paused
| from __future__ import annotations | |
| import socket | |
| from app.configs.settings import get_settings | |
| from app.models.common import utc_now | |
| from app.models.domain import JobProgress | |
| from app.services.comfyui_service import ComfyUIService | |
| from app.services.container import AppContainer | |
| from app.services.huggingface_service import HuggingFaceService | |
| from app.services.workflow_service import WorkflowService | |
| from app.utils.enums import CreditReason, JobStatus | |
| class WorkerRuntime: | |
| def __init__(self) -> None: | |
| self._settings = get_settings() | |
| self._container = AppContainer(self._settings) | |
| self._comfyui = ComfyUIService(self._settings) | |
| self._workflow = WorkflowService() | |
| self._huggingface = ( | |
| HuggingFaceService(self._settings) | |
| if self._settings.features.enable_model_auto_sync | |
| else None | |
| ) | |
| async def startup(self) -> None: | |
| await self._container.startup() | |
| async def shutdown(self) -> None: | |
| await self._container.shutdown() | |
| async def process_job(self, job_id: str, celery_task_id: str) -> None: | |
| if self._container.repositories is None or self._container.services is None: | |
| raise RuntimeError("Worker container did not initialize.") | |
| repositories = self._container.repositories | |
| services = self._container.services | |
| job = await repositories.jobs.get_by_id(job_id) | |
| if job is None: | |
| raise ValueError(f"Job {job_id} not found.") | |
| user = await repositories.users.get_by_id(job.user_id) | |
| if user is None: | |
| raise ValueError(f"User {job.user_id} not found.") | |
| await repositories.jobs.update_status( | |
| job_id, | |
| status=JobStatus.RUNNING, | |
| progress=JobProgress(percent=2, stage="initializing", message="Preparing worker"), | |
| fields={ | |
| "celery_task_id": celery_task_id, | |
| "assigned_worker": socket.gethostname(), | |
| "started_at": utc_now(), | |
| }, | |
| ) | |
| model = services.models.get_model(job.options.model_key) | |
| if self._huggingface is not None: | |
| await self._huggingface.ensure_model_available(model) | |
| source_filename: str | None = None | |
| if job.options.source_image_id: | |
| source_image = await repositories.images.get_by_id(job.options.source_image_id) | |
| if source_image is None: | |
| raise ValueError("Source image not found.") | |
| source_bytes = await services.storage.read(source_image.storage_key) | |
| source_filename = f"{source_image.id}.png" | |
| await self._comfyui.upload_image(source_filename, source_bytes, source_image.mime_type) | |
| workflow = self._workflow.build(job, model, source_filename=source_filename) | |
| async def on_progress(percent: int, stage: str) -> None: | |
| await repositories.jobs.update_status( | |
| job_id, | |
| status=JobStatus.RUNNING, | |
| progress=JobProgress(percent=percent, stage=stage, message=stage.title()), | |
| ) | |
| execution = await self._comfyui.execute_workflow(workflow, progress_cb=on_progress) | |
| image_ids: list[str] = [] | |
| for image_bytes in execution.images: | |
| created_image = await services.images.store_generated_image(user.id or "", job_id, image_bytes) | |
| image_ids.append(created_image.id or "") | |
| await repositories.jobs.attach_images(job_id, image_ids) | |
| await repositories.jobs.update_status( | |
| job_id, | |
| status=JobStatus.SUCCEEDED, | |
| progress=JobProgress(percent=100, stage="completed", message="Completed"), | |
| fields={ | |
| "finished_at": utc_now(), | |
| "comfy_prompt_id": execution.prompt_id, | |
| "image_ids": image_ids, | |
| }, | |
| ) | |
| async def mark_retrying(self, job_id: str, error_message: str) -> None: | |
| if self._container.repositories is None: | |
| return | |
| await self._container.repositories.jobs.update_status( | |
| job_id, | |
| status=JobStatus.RETRYING, | |
| progress=JobProgress(percent=0, stage="retrying", message="Retrying"), | |
| error_message=error_message, | |
| ) | |
| async def mark_failed(self, job_id: str, error_message: str) -> None: | |
| if self._container.repositories is None or self._container.services is None: | |
| return | |
| repositories = self._container.repositories | |
| services = self._container.services | |
| job = await repositories.jobs.get_by_id(job_id) | |
| if job is None: | |
| return | |
| user = await repositories.users.get_by_id(job.user_id) | |
| if user is not None and job.charged_credits > 0: | |
| await services.credits.grant( | |
| user=user, | |
| amount=job.charged_credits, | |
| reason=CreditReason.REFUND, | |
| note=f"Refund for failed job {job_id}", | |
| ) | |
| await repositories.jobs.update_status( | |
| job_id, | |
| status=JobStatus.FAILED, | |
| progress=JobProgress(percent=0, stage="failed", message="Failed"), | |
| error_message=error_message, | |
| fields={"finished_at": utc_now()}, | |
| ) | |
| async def cleanup_expired_images(self) -> int: | |
| if self._container.repositories is None or self._container.services is None: | |
| return 0 | |
| repositories = self._container.repositories | |
| services = self._container.services | |
| expired_images = await repositories.images.list_expired(limit=500) | |
| deleted = 0 | |
| for image in expired_images: | |
| await services.storage.delete(image.storage_key) | |
| if image.id: | |
| await repositories.images.delete(image.id) | |
| deleted += 1 | |
| return deleted | |