File size: 22,953 Bytes
405c51e
 
 
58510b2
 
405c51e
 
 
 
 
 
 
 
 
ae27ee5
405c51e
 
 
 
 
 
8d7fd66
405c51e
 
 
 
 
 
 
58510b2
 
 
 
 
 
 
6677a32
405c51e
 
 
 
 
 
 
 
 
 
 
58510b2
405c51e
 
 
1c16c4d
405c51e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f5509c2
 
 
 
9fa898a
f5509c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9fa898a
f5509c2
 
 
 
 
 
 
 
 
 
 
405c51e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
73ac5f9
405c51e
73ac5f9
5f9671b
 
405c51e
5f9671b
405c51e
97b7208
73ac5f9
405c51e
 
73ac5f9
405c51e
 
 
9fa898a
405c51e
 
c399783
 
 
 
 
 
 
 
 
5f9671b
c399783
73ac5f9
 
 
 
 
405c51e
5f9671b
97b7208
5f9671b
97b7208
405c51e
 
 
5f9671b
 
405c51e
 
ed32595
 
 
 
 
 
 
 
 
98b4113
ed32595
 
 
 
 
 
98b4113
ed32595
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98b4113
ed32595
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
405c51e
 
 
 
ae27ee5
405c51e
58510b2
405c51e
 
58510b2
73ac5f9
405c51e
7f4f914
8d7fd66
7f4f914
405c51e
7f4f914
58510b2
 
 
405c51e
 
 
 
 
ae27ee5
 
 
 
 
 
 
405c51e
ae27ee5
405c51e
e4ae032
ae27ee5
1d1a8c3
 
ae27ee5
1d1a8c3
ae27ee5
1d1a8c3
ae27ee5
 
73ac5f9
ae27ee5
d36fcbc
f5509c2
d36fcbc
e4ae032
405c51e
 
dd8ce86
405c51e
dd8ce86
f5509c2
 
73ac5f9
f5509c2
14301a8
 
 
 
 
dd8ce86
73ac5f9
f5509c2
dd8ce86
 
 
405c51e
73ac5f9
97b7208
405c51e
 
dd8ce86
e4ae032
73ac5f9
2b02e16
 
 
405c51e
dd8ce86
405c51e
 
 
1c16c4d
 
58510b2
405c51e
 
 
 
 
 
 
 
73ac5f9
 
 
405c51e
 
1c16c4d
 
 
 
 
 
 
 
 
405c51e
 
9fa898a
4e26c94
73ac5f9
405c51e
 
73ac5f9
405c51e
1c16c4d
405c51e
 
 
 
73ac5f9
405c51e
1c16c4d
405c51e
 
 
73ac5f9
405c51e
 
 
 
 
 
 
 
4e26c94
 
 
 
 
405c51e
 
 
98b4113
 
 
 
 
 
 
2b02e16
 
 
405c51e
 
 
 
ae27ee5
405c51e
 
 
73ac5f9
405c51e
 
 
 
 
 
c6c3eb4
405c51e
 
 
 
 
73ac5f9
 
 
 
405c51e
 
 
 
 
 
73ac5f9
405c51e
 
97b7208
 
405c51e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c16c4d
405c51e
 
 
 
 
 
 
 
ae27ee5
97b7208
73ac5f9
58510b2
ae27ee5
c399783
97b7208
58510b2
97b7208
73ac5f9
405c51e
 
73ac5f9
 
14301a8
405c51e
97b7208
73ac5f9
405c51e
 
 
 
1c16c4d
 
405c51e
 
 
 
 
 
 
 
ae27ee5
405c51e
 
 
 
 
ae27ee5
97b7208
4b9047d
5f9671b
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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
# -*- coding: utf-8 -*-
"""
Hugging Face Space for Text-to-Video generation using the Wan 2.1 model,
enhanced with a base `FusionX` LoRA, dynamic user-selectable style LoRAs,
and an LLM-based prompt enhancer.
"""

# --- 1. Imports ---
import os
import re
import json
import random
import tempfile
import traceback
from functools import partial

import gradio as gr
import numpy as np
import torch
import spaces

from diffusers import DiffusionPipeline, AutoModel, AutoencoderKLWan
from diffusers.utils import export_to_video
from transformers import AutoTokenizer, AutoModelForCausalLM, UMT5EncoderModel, pipeline
from huggingface_hub import hf_hub_download, list_repo_files

# --- 2. Configuration & Constants ---

# --- Model & LoRA Identifiers ---
T2V_BASE_MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"

# Base LoRA (always applied)
FUSIONX_LORA_REPO = "vrgamedevgirl84/Wan14BT2VFusioniX"
FUSIONX_LORA_FILE = "FusionX_LoRa/Wan2.1_T2V_14B_FusionX_LoRA.safetensors"
FUSIONX_ADAPTER_NAME = "fusionx_t2v"
FUSIONX_LORA_WEIGHT = 0.75
ENHANCER_MODEL_ID = "Qwen/Qwen2-1.5B-Instruct" # Using a smaller model to save space

# Dynamic LoRAs (user selectable)
DYNAMIC_LORA_REPO_ID = "DeepBeepMeep/Wan2.1"
DYNAMIC_LORA_SUBFOLDER = "loras_i2v"
DYNAMIC_LORA_ADAPTER_NAME = "dynamic_lora"

# --- Generation Parameters ---
MOD_VALUE = 8
T2V_FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81
MAX_SEED = np.iinfo(np.int64).max

# --- UI Defaults ---
DEFAULT_H_SLIDER_VALUE = 480
DEFAULT_W_SLIDER_VALUE = 640
DEFAULT_PROMPT_T2V = "A majestic lion surveying its kingdom from a rocky outcrop at sunrise, cinematic lighting, hyperrealistic."
DEFAULT_NEGATIVE_PROMPT = "static image, no motion, watermark, text, signature, jpeg artifacts, ugly, incomplete, disfigured, low quality, worst quality, messy background"

# --- System Prompt for LLM Enhancer ---
T2V_CINEMATIC_PROMPT_SYSTEM = (
    "You are a prompt engineer for a generative AI video model. Your task is to rewrite user inputs into high-quality, "
    "detailed, and cinematic prompts. Focus on visual details, camera movements, lighting, and mood. "
    "Add natural motion attributes. The revised prompt should be around 80-100 words. "
    "Directly output the rewritten prompt in English without any conversational text or quotation marks."
)


# --- 3. Helper Functions ---

def sanitize_prompt_for_filename(prompt: str) -> str:
    """Creates a filesystem-safe filename from a prompt."""
    if not prompt:
        return "no_prompt"
    sanitized = re.sub(r'[^a-zA-Z0-9\s]', '', prompt)
    sanitized = re.sub(r'\s+', '_', sanitized).lower()
    return sanitized[:50]

def get_t2v_duration(
    prompt: str, height: int, width: int, negative_prompt: str,
    duration_seconds: float, steps: int, seed: int,
    randomize_seed: bool, selected_lora: str,
    lora_weight: float
) -> int:
    """
    Estimates GPU time for Text-to-Video generation.
    The logic is tiered and considers duration, steps, and resolution to prevent timeouts.
    """
    # Calculate a resolution multiplier. A higher resolution will significantly increase generation time.
    # Base resolution is considered 640x480 pixels.
    base_pixels = DEFAULT_W_SLIDER_VALUE * DEFAULT_H_SLIDER_VALUE
    current_pixels = width * height
    # Check if the current resolution is significantly larger than the base.
    is_high_res = current_pixels > (base_pixels * 1.5)

    # Tiered duration based on video length and number of inference steps.
    if steps > 10 or duration_seconds > 4:
        # Longest generations (e.g., high step count or long duration).
        base_duration = 600
    elif steps > 10 or duration_seconds > 3:
        # Medium-length generations.
        base_duration = 400
    else:
        # Shortest/quickest generations.
        base_duration = 250

    # Apply a multiplier for high-resolution videos.
    final_duration = base_duration * 2 if is_high_res else base_duration
    
    # Cap the duration at a maximum value (900s = 15 minutes) to comply with typical free-tier limits.
    final_duration = min(final_duration, 900)

    print(f"Requesting {final_duration}s of GPU time for {steps} steps, {duration_seconds:.1f}s duration, and {width}x{height} resolution.")
    return final_duration

def get_available_presets(repo_id, subfolder):
    """
    Fetches the list of available LoRA presets by looking for .lset files.
    """
    print(f"\nπŸ”Ž Discovering LoRA presets in {repo_id}...")
    try:
        all_files = list_repo_files(repo_id=repo_id, repo_type='model')
        subfolder_path = f"{subfolder}/"
        lset_files = [
            os.path.splitext(f.split('/')[-1])[0]
            for f in all_files
            if f.startswith(subfolder_path) and f.endswith('.lset')
        ]
        print(f"βœ… Discovered {len(lset_files)} LoRA presets.")
        return ["None"] + sorted(lset_files)
    except Exception as e:
        print(f"⚠️ Warning: Could not fetch LoRA presets from {repo_id}. LoRA selection will be disabled. Error: {e}")
        return ["None"]

def parse_lset_prompt(lset_prompt):
    """Parses a .lset prompt, resolving variables and highlighting them."""
    variables = dict(re.findall(r'! \{(\w+)\}="([^"]+)"', lset_prompt))
    prompt_template = re.sub(r'! \{\w+\}="[^"]+"\n?', '', lset_prompt).strip()
    resolved_prompt = prompt_template
    for key, value in variables.items():
        highlighted_value = f"__{value}__"
        resolved_prompt = resolved_prompt.replace(f"{{{key}}}", highlighted_value)
    return resolved_prompt

def handle_lora_selection_change(preset_name: str, current_prompt: str):
    """
    Appends the selected LoRA's trigger words to the current prompt text
    and controls the visibility of the weight slider. Ensures slider is only
    visible on success.
    """
    # If "None" is selected, hide the slider and return the prompt unchanged.
    if not preset_name or preset_name == "None":
        gr.Info("LoRA cleared.")
        return gr.update(value=current_prompt), gr.update(visible=False, interactive=False)

    try:
        # Fetch the trigger words from the LoRA's .lset file.
        lset_filename = f"{preset_name}.lset"
        lset_path = hf_hub_download(
            repo_id=DYNAMIC_LORA_REPO_ID,
            filename=lset_filename, subfolder=DYNAMIC_LORA_SUBFOLDER, repo_type='model'
        )
        with open(lset_path, 'r', encoding='utf-8') as f:
            lset_content = f.read()

        lset_prompt_raw = None
        try:
            lset_data = json.loads(lset_content)
            lset_prompt_raw = lset_data.get("prompt")
        except json.JSONDecodeError:
            lset_prompt_raw = lset_content

        # Only if we successfully get trigger words, we update the prompt and show the slider.
        if lset_prompt_raw:
            trigger_words = parse_lset_prompt(lset_prompt_raw)
            separator = ", " if current_prompt and not current_prompt.endswith((",", " ")) else ""
            new_prompt = f"{current_prompt}{separator}{trigger_words}".strip()
            gr.Info(f"βœ… Appended triggers from '{preset_name}'. You can now edit them.")
            return gr.update(value=new_prompt), gr.update(visible=True, interactive=True)
        else:
            # If the .lset file has no prompt, don't change the prompt and ensure the slider is hidden.
            gr.Info(f"ℹ️ No prompt found in '{preset_name}.lset'. Prompt unchanged.")
            return gr.update(value=current_prompt), gr.update(visible=False, interactive=False)

    except Exception as e:
        print(f"Info: Could not process .lset for '{preset_name}'. Reason: {e}")
        gr.Warning(f"⚠️ Could not load triggers for '{preset_name}'.")
        # On any error, don't change the prompt and ensure the slider is hidden.
        return gr.update(value=current_prompt), gr.update(visible=False, interactive=False)


def _manage_lora_state(pipe, selected_lora: str, lora_weight: float) -> bool:
    """
    Handles the loading, setting, and cleanup of dynamic LoRA adapters.

    Returns:
        bool: True if a dynamic LoRA was loaded, False otherwise.
    """
    # Pre-emptive cleanup of any previously loaded dynamic adapter.
    try:
        pipe.delete_adapters([DYNAMIC_LORA_ADAPTER_NAME])
        print("🧼 Pre-emptively unloaded previous dynamic LoRA.")
    except Exception:
        pass  # No dynamic lora was present, which is a clean state.

    if not selected_lora or selected_lora == "None":
        pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
        print("ℹ️ No dynamic LoRA selected. Using base LoRA only.")
        return False

    # --- DYNAMIC LORA HANDLING ---
    print(f"πŸš€ Processing preset: {selected_lora} with weight {lora_weight}")
    lora_filename = None
    try:
        lset_filename = f"{selected_lora}.lset"
        lset_path = hf_hub_download(
            repo_id=DYNAMIC_LORA_REPO_ID,
            filename=lset_filename, subfolder=DYNAMIC_LORA_SUBFOLDER, repo_type='model'
        )
        with open(lset_path, 'r', encoding='utf-8') as f:
            lset_content = f.read()
        try:
            lset_data = json.loads(lset_content)
            loras_list = lset_data.get("loras")
            if loras_list and isinstance(loras_list, list) and len(loras_list) > 0:
                lora_filename = loras_list[0]
                print(f"   - Found LoRA file in preset: {lora_filename}")
        except json.JSONDecodeError:
            print(f"   - Info: '{lset_filename}' is not JSON. Assuming filename matches preset name.")
    except Exception as e:
            print(f"   - Warning: Could not process .lset file for '{selected_lora}'. Assuming filename matches preset. Error: {e}")

    if not lora_filename:
        lora_filename = f"{selected_lora}.safetensors"

    pipe.load_lora_weights(
        DYNAMIC_LORA_REPO_ID, weight_name=lora_filename,
        subfolder=DYNAMIC_LORA_SUBFOLDER, adapter_name=DYNAMIC_LORA_ADAPTER_NAME,
    )
    pipe.set_adapters(
        [FUSIONX_ADAPTER_NAME, DYNAMIC_LORA_ADAPTER_NAME],
        adapter_weights=[FUSIONX_LORA_WEIGHT, lora_weight]
    )
    print("βœ… Dynamic LoRA activated alongside base LoRA.")
    return True


# --- 4. Pipeline Loading ---

def load_pipelines():
    """Loads and configures the T2V and LLM pipelines."""
    t2v_pipe, enhancer_pipe = None, None

    print("\nπŸš€ Loading T2V pipeline with base LoRA...")
    try:
        t2v_pipe = DiffusionPipeline.from_pretrained(
            T2V_BASE_MODEL_ID,
            torch_dtype=torch.bfloat16,
        )
        print("βœ… Base pipeline loaded. Overriding VAE with float32 version...")
        vae_fp32 = AutoencoderKLWan.from_pretrained(T2V_BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
        t2v_pipe.vae = vae_fp32
        t2v_pipe.to("cuda")
        print("βœ… Pipeline configured. Loading and activating base FusionX LoRA...")
        t2v_pipe.load_lora_weights(FUSIONX_LORA_REPO, weight_name=FUSIONX_LORA_FILE, adapter_name=FUSIONX_ADAPTER_NAME)
        t2v_pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
        print("βœ… T2V pipeline with base LoRA is ready.")
    except Exception as e:
        print(f"❌ CRITICAL ERROR: Failed to load T2V pipeline. T2V will be disabled. Reason: {e}")
        traceback.print_exc()
        t2v_pipe = None

    print("\nπŸ€– Loading LLM for Prompt Enhancement...")
    try:
        enhancer_pipe = pipeline("text-generation", model=ENHANCER_MODEL_ID, torch_dtype=torch.bfloat16, device="cpu")
        print("βœ… LLM Prompt Enhancer loaded successfully (on CPU).")
    except Exception as e:
        print(f"⚠️ WARNING: Could not load the LLM prompt enhancer. The feature will be disabled. Error: {e}")
        enhancer_pipe = None

    return t2v_pipe, enhancer_pipe


# --- 5. Core Generation & UI Logic ---

@spaces.GPU()
def enhance_prompt_with_llm(prompt: str, enhancer_pipeline):
    """
    Uses the loaded LLM to enhance a given prompt.
    """
    if enhancer_pipeline is None:
        gr.Warning("LLM enhancer is not available.")
        return prompt, gr.update(), gr.update()

    if enhancer_pipeline.model.device.type != 'cuda':
        print("Moving enhancer model to CUDA for on-demand GPU execution...")
        enhancer_pipeline.model.to("cuda")

    messages = [{"role": "system", "content": T2V_CINEMATIC_PROMPT_SYSTEM}, {"role": "user", "content": prompt}]
    print(f"Enhancing prompt: '{prompt}'")

    try:
        tokenizer = enhancer_pipeline.tokenizer
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenized_inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
        
        if isinstance(tokenized_inputs, torch.Tensor):
            inputs_on_cuda = {"input_ids": tokenized_inputs.to("cuda")}
            inputs_on_cuda["attention_mask"] = torch.ones_like(inputs_on_cuda["input_ids"])
        else:
            inputs_on_cuda = {k: v.to("cuda") for k, v in tokenized_inputs.items()}

        generated_ids = enhancer_pipeline.model.generate(**inputs_on_cuda, max_new_tokens=256, do_sample=True, temperature=0.7, top_p=0.95)
        input_token_length = inputs_on_cuda['input_ids'].shape[1]
        newly_generated_ids = generated_ids[:, input_token_length:]
        final_answer = tokenizer.decode(newly_generated_ids[0], skip_special_tokens=True)

        print(f"Enhanced prompt: '{final_answer.strip()}'")
        # The enhanced prompt overwrites the textbox. The LoRA selection is reset.
        return final_answer.strip(), "None", gr.update(visible=False, interactive=False)
    except Exception as e:
        print(f"❌ Error during prompt enhancement: {e}")
        traceback.print_exc()
        gr.Warning(f"An error occurred during prompt enhancement. See console for details.")
        return prompt, gr.update(), gr.update()
    finally:
        print("🧹 Clearing CUDA cache after prompt enhancement...")
        torch.cuda.empty_cache()


@spaces.GPU(duration_from_args=get_t2v_duration)
def generate_t2v_video(
    prompt: str, height: int, width: int, negative_prompt: str,
    duration_seconds: float, steps: int, seed: int,
    randomize_seed: bool, selected_lora: str,
    lora_weight: float,
    progress=gr.Progress(track_tqdm=True)
):
    """Main function to generate a video from a text prompt."""
    if t2v_pipe is None:
        raise gr.Error("Text-to-Video pipeline is not available due to a loading error.")
    if not prompt:
        raise gr.Error("Please enter a prompt for Text-to-Video generation.")

    # --- The prompt from the textbox is now the final prompt. No more combining. ---
    final_prompt = prompt

    target_h = max(MOD_VALUE, (height // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (width // MOD_VALUE) * MOD_VALUE)

    requested_frames = int(round(duration_seconds * T2V_FIXED_FPS))
    frames_minus_one = requested_frames - 1
    valid_frames_minus_one = round(frames_minus_one / 4.0) * 4
    num_frames = int(valid_frames_minus_one) + 1
    if num_frames != requested_frames:
        print(f"Info: Adjusted number of frames from {requested_frames} to {num_frames} to meet model constraints.")
    
    num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    lora_loaded = False

    try:
        lora_loaded = _manage_lora_state(pipe=t2v_pipe, selected_lora=selected_lora, lora_weight=lora_weight)

        print("\n--- Starting T2V Generation ---")
        print(f"Final Prompt: {final_prompt}")
        print(f"Resolution: {target_w}x{target_h}, Frames: {num_frames}, Seed: {current_seed}")
        print(f"Steps: {steps}, Guidance: 1.0 (fixed for FusionX)")
        print("---------------------------------")

        with torch.inference_mode():
            output_frames_list = t2v_pipe(
                prompt=final_prompt, negative_prompt=negative_prompt,
                height=target_h, width=target_w, num_frames=num_frames,
                guidance_scale=1.0, num_inference_steps=int(steps),
                generator=torch.Generator(device="cuda").manual_seed(current_seed)
            ).frames[0]

        sanitized_prompt = sanitize_prompt_for_filename(final_prompt)
        filename = f"t2v_{sanitized_prompt}_{current_seed}.mp4"
        temp_dir = tempfile.mkdtemp()
        video_path = os.path.join(temp_dir, filename)
        export_to_video(output_frames_list, video_path, fps=T2V_FIXED_FPS)

        print(f"βœ… Video saved to: {video_path}")
        download_label = f"πŸ“₯ Download: {filename}"
        return video_path, current_seed, gr.File(value=video_path, visible=True, label=download_label)
    
    except Exception as e:
        print(f"❌ An error occurred during video generation: {e}")
        traceback.print_exc()
        raise gr.Error("Video generation failed. Please check the logs for details.")

    finally:
        if lora_loaded:
            print(f"🧼 Cleaning up dynamic LoRA: {selected_lora}")
            try:
                t2v_pipe.delete_adapters([DYNAMIC_LORA_ADAPTER_NAME])
                t2v_pipe.set_adapters([FUSIONX_ADAPTER_NAME], adapter_weights=[FUSIONX_LORA_WEIGHT])
                print("βœ… Cleanup complete. Pipeline reset to base LoRA state.")
            except Exception as e:
                print(f"⚠️ Error during LoRA cleanup: {e}. State may be inconsistent.")
        
        print("🧹 Clearing CUDA cache after video generation...")
        torch.cuda.empty_cache()


# --- 6. Gradio UI Layout ---

def build_ui(t2v_pipe, enhancer_pipe, available_loras):
    """Creates and configures the Gradio UI."""
    with gr.Blocks(theme=gr.themes.Soft(), css=".main-container { max-width: 1080px; margin: auto; }") as demo:
        gr.Markdown("# ✨ Wan 2.1 Text-to-Video Suite with Dynamic LoRAs")
        gr.Markdown("Generate videos from text. Edit the prompt below. Selecting a LoRA will append its triggers to your prompt.")

        with gr.Tabs():
            with gr.TabItem("✍️ Text-to-Video", id="t2v_tab", interactive=t2v_pipe is not None):
                if t2v_pipe is None:
                    gr.Markdown("<h3 style='color: #ff9999; text-align: center;'>⚠️ T2V Pipeline Failed to Load. Tab disabled.</h3>")
                else:
                    with gr.Row():
                        with gr.Column(scale=2):
                            t2v_prompt = gr.Textbox(
                                label="✏️ Prompt", value=DEFAULT_PROMPT_T2V, lines=4,
                                placeholder="e.g., A cinematic drone shot flying over a futuristic city at night..."
                            )
                            t2v_enhance_btn = gr.Button(
                                "πŸ€– Enhance Prompt with AI",
                                interactive=enhancer_pipe is not None
                            )

                            with gr.Group():
                                t2v_lora_preset = gr.Dropdown(
                                    label="🎨 Dynamic Style LoRA (Optional)",
                                    choices=available_loras,
                                    value="None",
                                    info="Appends style triggers to the prompt text above."
                                )
                                t2v_lora_weight = gr.Slider(
                                    label="πŸ’ͺ LoRA Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.8, 
                                    interactive=False, visible=False
                                )

                            t2v_duration = gr.Slider(
                                minimum=round(MIN_FRAMES_MODEL / T2V_FIXED_FPS, 1),
                                maximum=round(MAX_FRAMES_MODEL / T2V_FIXED_FPS, 1),
                                step=0.1, value=2.0, label="⏱️ Duration (seconds)"
                            )
                            with gr.Accordion("βš™οΈ Advanced Settings", open=False):
                                t2v_neg_prompt = gr.Textbox(label="❌ Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=3)
                                with gr.Row():
                                    t2v_seed = gr.Slider(label="🎲 Seed", minimum=0, maximum=MAX_SEED, step=1, value=1234)
                                    t2v_rand_seed = gr.Checkbox(label="πŸ”€ Randomize", value=True)
                                with gr.Row():
                                    t2v_height = gr.Slider(minimum=256, maximum=896, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label="πŸ“ Height")
                                    t2v_width = gr.Slider(minimum=256, maximum=896, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label="πŸ“ Width")
                                t2v_steps = gr.Slider(minimum=5, maximum=40, step=1, value=8, label="πŸš€ Inference Steps")

                            t2v_generate_btn = gr.Button("🎬 Generate Video", variant="primary")

                        with gr.Column(scale=3):
                            t2v_output_video = gr.Video(label="πŸŽ₯ Generated Video", autoplay=True, interactive=False)
                            t2v_download = gr.File(label="πŸ“₯ Download Video", visible=False)

        if t2v_pipe is not None:
            enhance_fn = partial(enhance_prompt_with_llm, enhancer_pipeline=enhancer_pipe)
            
            # 1. When the user enhances the prompt with the LLM.
            t2v_enhance_btn.click(
                fn=enhance_fn,
                inputs=[t2v_prompt],
                outputs=[t2v_prompt, t2v_lora_preset, t2v_lora_weight]
            )

            # 2. When the user selects a LoRA from the dropdown.
            t2v_lora_preset.change(
                fn=handle_lora_selection_change,
                # Pass the current prompt text in, get the new text back out.
                inputs=[t2v_lora_preset, t2v_prompt],
                outputs=[t2v_prompt, t2v_lora_weight]
            )
            
            # 3. When the user clicks the final generate button.
            t2v_generate_btn.click(
                fn=generate_t2v_video,
                inputs=[
                    t2v_prompt, t2v_height, t2v_width, t2v_neg_prompt,
                    t2v_duration, t2v_steps, t2v_seed,
                    t2v_rand_seed, t2v_lora_preset, t2v_lora_weight
                ],
                outputs=[t2v_output_video, t2v_seed, t2v_download]
            )
    return demo


# --- 7. Main Execution ---
if __name__ == "__main__":
    t2v_pipe, enhancer_pipe = load_pipelines()
    
    available_loras = []
    if t2v_pipe:
        available_loras = get_available_presets(DYNAMIC_LORA_REPO_ID, DYNAMIC_LORA_SUBFOLDER)

    app_ui = build_ui(t2v_pipe, enhancer_pipe, available_loras)
    app_ui.queue(max_size=10).launch()