File size: 5,350 Bytes
36333c5
 
 
 
 
 
 
eb808a5
36333c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eb808a5
 
 
 
 
 
 
 
 
 
 
 
 
 
36333c5
 
 
eb808a5
 
 
 
 
 
 
 
 
 
 
36333c5
 
 
 
 
 
 
 
eb808a5
 
 
 
 
 
 
 
 
 
 
 
36333c5
 
 
eb808a5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36333c5
 
 
eb808a5
 
 
 
 
 
 
 
 
 
 
36333c5
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
"""Full media workflow orchestration outside the HTTP layer."""

from __future__ import annotations

from pathlib import Path

from config import Settings
from core.executor import InferenceCommand
from core.manager import TaskManager
from core.schemas import WorkflowRequest
from utils.files import OutputManager


class WorkflowService:
    """Run image, voice, music, video, and subtitle generation in order."""

    def __init__(
        self, settings: Settings, tasks: TaskManager, outputs: OutputManager
    ) -> None:
        self.settings = settings
        self.tasks = tasks
        self.outputs = outputs

    async def run(self, payload: WorkflowRequest, request_id: str) -> dict[str, Path]:
        """Queue the entire pipeline as one non-interleavable operation."""
        width = payload.width if payload.width is not None else self.settings.image_width
        height = payload.height if payload.height is not None else self.settings.image_height
        voice = payload.voice or self.settings.kokoro_default_voice
        targets = {
            "image": self.outputs.allocate("images"),
            "voice": self.outputs.allocate("audio"),
            "music": self.outputs.allocate("music"),
            "video": self.outputs.allocate("videos"),
            "subtitle": self.outputs.allocate("subtitles"),
        }

        def execute() -> dict[str, Path]:
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="flux",
                    method_name="generate",
                    arguments={
                        "prompt": payload.image_prompt,
                        "output_path": targets["image"],
                        "width": width,
                        "height": height,
                        "steps": self.settings.image_steps,
                        "seed": payload.seed,
                        "guidance_scale": self.settings.flux_guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_flux_duration,
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="kokoro",
                    method_name="synthesize",
                    arguments={
                        "text": payload.script,
                        "voice": voice,
                        "speed": 1.0,
                        "output_path": targets["voice"],
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_kokoro_duration,
                ),
            )
            words = len(payload.script.split())
            music_duration = min(30.0, max(5.0, words / 2.5))
            music_prompt = payload.music_prompt or (
                f"cinematic instrumental background score for {payload.title}, no vocals"
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="musicgen",
                    method_name="generate",
                    arguments={
                        "prompt": music_prompt,
                        "duration": music_duration,
                        "guidance_scale": self.settings.music_guidance_scale,
                        "seed": payload.seed,
                        "output_path": targets["music"],
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_musicgen_duration,
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="wan",
                    method_name="generate",
                    arguments={
                        "image_path": targets["image"],
                        "prompt": payload.video_prompt,
                        "negative_prompt": "low quality, distorted, static",
                        "output_path": targets["video"],
                        "steps": self.settings.wan_steps,
                        "frames": self.settings.video_frames,
                        "fps": self.settings.video_fps,
                        "seed": payload.seed,
                        "guidance_scale": self.settings.wan_guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_wan_duration,
                    gpu_size="xlarge",
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="whisper",
                    method_name="transcribe",
                    arguments={
                        "source": targets["voice"],
                        "subtitle_path": targets["subtitle"],
                        "language": None,
                        "task": "transcribe",
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_whisper_duration,
                ),
            )
            return targets

        try:
            return await self.tasks.run_exclusive(request_id, "workflow", execute)
        except Exception:
            for target in targets.values():
                self.outputs.remove(target)
            raise