File size: 6,385 Bytes
7344bef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
from datetime import datetime
import time
import json
import gradio as gr

from shared.utils.plugins import WAN2GPPlugin
from shared.utils.process_locks import acquire_GPU_ressources, any_GPU_process_running, release_GPU_ressources

PlugIn_Name = "Sample Plugin"
PlugIn_Id = "SamplePlugin"


def acquire_GPU(state):
    GPU_process_running = any_GPU_process_running(state, PlugIn_Id)
    if GPU_process_running:
        gr.Error("Another PlugIn is using the GPU")
    acquire_GPU_ressources(state, PlugIn_Id, PlugIn_Name, gr=gr)


def release_GPU(state):
    release_GPU_ressources(state, PlugIn_Id)


class ConfigTabPlugin(WAN2GPPlugin):
    def setup_ui(self):
        self.request_global("get_current_model_settings")
        self.request_global("refresh_model_defs")        
        self.request_component("refresh_form_trigger")
        self.request_component("state")
        self.request_component("resolution")
        self.request_component("main_tabs")
        self.add_tab(tab_id=PlugIn_Id, label=PlugIn_Name, component_constructor=self.create_config_ui)

    @staticmethod
    def _demo_settings(_state) -> dict[str, object]:
        return {
            "model_type": "ltx2_22B_distilled",
            "prompt": "A reframed cinematic 10 second video shot: a woman walks slowly through a rainy neon street at night, camera gently orbiting, reflections shimmering on the wet pavement, natural motion, realistic lighting.",
            "resolution": "1280x720",
            "num_inference_steps": 8,
            "video_length": 241,
            "sliding_window_size": 481,
            "sliding_window_overlap": 17,
        }

    def on_tab_select(self, state: dict) -> None:
        settings = self.get_current_model_settings(state)
        prompt = settings["prompt"]
        return prompt

    def on_tab_deselect(self, state: dict) -> None:
        pass

    def on_model_change(self, state: dict, model_type) -> None:
        # print(f"new model selected is {model_type}")
        pass

    def create_config_ui(self, api_session):
        active_job = {"job": None}

        def update_prompt(state, text):
            settings = self.get_current_model_settings(state)
            settings["prompt"] = text
            return time.time()

        def big_process(state):
            acquire_GPU(state)
            gr.Info("Doing something important that Requires Full VRAM & GPU available")
            time.sleep(30)
            release_GPU(state)
            return "42"

        def generate_video(progress=gr.Progress(track_tqdm=False)):
            class DemoCallbacks:
                ratio = 0.0

                def on_status(self, status):
                    status = str(status or "").strip()
                    if status:
                        progress(self.ratio, desc=status)

                def on_progress(self, update):
                    self.ratio = max(0.0, min(1.0, float(getattr(update, "progress", 0)) / 100.0))
                    progress(self.ratio, desc=str(getattr(update, "status", "") or "Generating..."))

            job = api_session.submit_task(self._demo_settings(self.state.value), callbacks=DemoCallbacks())
            active_job["job"] = job
            try:
                result = job.result()
            finally:
                if active_job.get("job") is job:
                    active_job["job"] = None
            if result.success and result.generated_files:
                return result.generated_files[0]
            if result.cancelled:
                return gr.update()
            errors = list(result.errors or [])
            raise gr.Error(str(errors[0] if errors else "WanGP completed without returning an output file."))

        def cancel_demo():
            job = active_job.get("job")
            if job is not None and not job.done:
                job.cancel()

        def write_finetune():
            the_time= datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            finetune_def =  {
                "model": {
                        "name": f"Finetune of {the_time}",
                        "architecture": "ltx2_22B",
                        "description": "enjoy this finetune of the moment",
                "URLs": "ltx2_22B",
                "preload_URLs":"ltx2_22B",        
                "ltx2_pipeline": "distilled"
                },
                "prompt": f"A video that is all about {the_time}",
            }
            with open(f"finetunes/to_be_deleted{the_time}.json", "w", encoding="utf-8") as writer:
                writer.write(json.dumps(finetune_def, indent=4))
            self.refresh_model_defs()
            gr.Info(f"finetune {the_time} had been created")


        with gr.Column():
            state = self.state
            settings = self.get_current_model_settings(state.value)
            prompt = settings["prompt"]
            gr.HTML("<B><B>Sample Plugin that illustrates</B>:<BR>-How to get Settings from Main Form and then Modify them<BR>-How to suspend the Video Gen (and release VRAM) to execute your own GPU intensive process.<BR>-How to switch back automatically to the Main Tab<BR>-How to trigger a Video Gen from a plugin an track its progress<BR>-Add a new finetune on the fly")
            sample_text = gr.Text(label="Prompt Copy", value=prompt, lines=5)
            update_btn = gr.Button("Update Prompt On Main Page")
            gr.Markdown()
            process_btn = gr.Button("Use GPU To Do Something Important")
            process_output = gr.Text(label="Process Output", value="")
            goto_btn = gr.Button("Goto Video Tab")
            gr.Markdown("---")
            start_btn = gr.Button("Generate a LTX 2.3 Video")
            output_video = gr.Video(label="Output")
            abort_btn = gr.Button("Abort Generation")

            write_finetune_btn = gr.Button("Generate a Random LTX 2 Finetune")

        self.on_tab_outputs = [sample_text]

        update_btn.click(fn=update_prompt, inputs=[state, sample_text], outputs=[self.refresh_form_trigger])
        process_btn.click(fn=big_process, inputs=[state], outputs=[process_output])
        goto_btn.click(fn=self.goto_video_tab, inputs=[state], outputs=[self.main_tabs])
        start_btn.click(fn=generate_video, outputs=[output_video], queue=False)
        abort_btn.click(fn=cancel_demo, queue=False)
        write_finetune_btn.click(fn=write_finetune)