jbilcke-hf HF Staff commited on
Commit
0431fa9
·
1 Parent(s): de828f6

improving the preview_tab

Browse files
vms/ui/project/services/previewing.py CHANGED
@@ -85,10 +85,17 @@ class PreviewingService:
85
  # Return updated log string for UI updates
86
  return "\n".join(log_messages)
87
 
88
- # Find latest LoRA weights
89
- lora_path = self.find_latest_lora_weights()
90
- if not lora_path:
91
- return None, "Error: No LoRA weights found", log("Error: No LoRA weights found in output directory")
 
 
 
 
 
 
 
92
 
93
  # Add prefix to prompt
94
  if prompt_prefix and not prompt.startswith(prompt_prefix):
@@ -143,12 +150,13 @@ class PreviewingService:
143
 
144
  log(f"Generating video with model type: {internal_model_type}")
145
  log(f"Using model version: {model_version}")
146
- log(f"Using LoRA weights from: {lora_path}")
 
147
  log(f"Resolution: {width}x{height}, Frames: {num_frames}, FPS: {fps}")
148
- log(f"Guidance Scale: {guidance_scale}, Flow Shift: {flow_shift}, LoRA Weight: {lora_weight}")
149
  log(f"Generation Seed: {seed}")
150
- log(f"Prompt: {full_prompt}")
151
- log(f"Negative Prompt: {negative_prompt}")
152
 
153
  # Import required components based on model type
154
  if internal_model_type == "wan":
@@ -246,11 +254,14 @@ class PreviewingService:
246
  log_fn("Enabling model CPU offload...")
247
  pipe.enable_model_cpu_offload()
248
 
249
- log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
250
- pipe.load_lora_weights(lora_path)
251
-
252
- # TODO: Set the lora scale directly instead of using fuse_lora
253
- #pipe._lora_scale = lora_weight
 
 
 
254
 
255
  # Create temporary file for the output
256
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
@@ -371,9 +382,13 @@ class PreviewingService:
371
  log_fn("Enabling model CPU offload...")
372
  pipe.enable_model_cpu_offload()
373
 
374
- log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
375
- pipe.load_lora_weights(lora_path)
376
- pipe.fuse_lora(lora_weight)
 
 
 
 
377
 
378
  # Create temporary file for the output
379
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
@@ -488,9 +503,13 @@ class PreviewingService:
488
  log_fn("Enabling model CPU offload...")
489
  pipe.enable_model_cpu_offload()
490
 
491
- log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
492
- pipe.load_lora_weights(lora_path)
493
- pipe.fuse_lora(lora_weight)
 
 
 
 
494
 
495
  # Create temporary file for the output
496
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
 
85
  # Return updated log string for UI updates
86
  return "\n".join(log_messages)
87
 
88
+ # Find latest LoRA weights if lora_weight > 0
89
+ lora_path = None
90
+ using_lora = lora_weight > 0
91
+
92
+ if using_lora:
93
+ lora_path = self.find_latest_lora_weights()
94
+ if not lora_path:
95
+ return None, "Error: No LoRA weights found", log("Error: No LoRA weights found in output directory")
96
+ log(f"Using LoRA weights with weight {lora_weight}")
97
+ else:
98
+ log("Using original model without LoRA weights")
99
 
100
  # Add prefix to prompt
101
  if prompt_prefix and not prompt.startswith(prompt_prefix):
 
150
 
151
  log(f"Generating video with model type: {internal_model_type}")
152
  log(f"Using model version: {model_version}")
153
+ if using_lora and lora_path:
154
+ log(f"Using LoRA weights from: {lora_path}")
155
  log(f"Resolution: {width}x{height}, Frames: {num_frames}, FPS: {fps}")
156
+ log(f"Guidance Scale: {guidance_scale}, Flow Shift: {flow_shift}, LoRA Weight: {lora_weight if using_lora else 0}")
157
  log(f"Generation Seed: {seed}")
158
+ #log(f"Prompt: {full_prompt}")
159
+ #log(f"Negative Prompt: {negative_prompt}")
160
 
161
  # Import required components based on model type
162
  if internal_model_type == "wan":
 
254
  log_fn("Enabling model CPU offload...")
255
  pipe.enable_model_cpu_offload()
256
 
257
+ # Apply LoRA weights if using them
258
+ if lora_weight > 0 and lora_path:
259
+ log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
260
+ pipe.load_lora_weights(lora_path)
261
+ # TODO: Set the lora scale directly instead of using fuse_lora
262
+ #pipe._lora_scale = lora_weight
263
+ else:
264
+ log_fn("Using base model without LoRA weights")
265
 
266
  # Create temporary file for the output
267
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
 
382
  log_fn("Enabling model CPU offload...")
383
  pipe.enable_model_cpu_offload()
384
 
385
+ # Apply LoRA weights if using them
386
+ if lora_weight > 0 and lora_path:
387
+ log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
388
+ pipe.load_lora_weights(lora_path)
389
+ pipe.fuse_lora(lora_weight)
390
+ else:
391
+ log_fn("Using base model without LoRA weights")
392
 
393
  # Create temporary file for the output
394
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
 
503
  log_fn("Enabling model CPU offload...")
504
  pipe.enable_model_cpu_offload()
505
 
506
+ # Apply LoRA weights if using them
507
+ if lora_weight > 0 and lora_path:
508
+ log_fn(f"Loading LoRA weights from {lora_path} with weight {lora_weight}...")
509
+ pipe.load_lora_weights(lora_path)
510
+ pipe.fuse_lora(lora_weight)
511
+ else:
512
+ log_fn("Using base model without LoRA weights")
513
 
514
  # Create temporary file for the output
515
  with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
vms/ui/project/tabs/preview_tab.py CHANGED
@@ -52,13 +52,15 @@ class PreviewTab(BaseTab):
52
  value=DEFAULT_PROMPT_PREFIX
53
  )
54
 
 
55
  self.components["seed"] = gr.Slider(
56
  label="Generation Seed (-1 for random)",
57
  minimum=-1,
58
  maximum=2147483647, # 2^31 - 1
59
  step=1,
60
  value=-1,
61
- info="Set to -1 for random seed or specific value for reproducible results"
 
62
  )
63
 
64
  with gr.Row():
@@ -81,6 +83,24 @@ class PreviewTab(BaseTab):
81
  value=self.get_default_model_version(default_model)
82
  )
83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  # Add image input for image-to-video models
85
  self.components["conditioning_image"] = gr.Image(
86
  label="Conditioning Image (for Image-to-Video models)",
@@ -148,7 +168,8 @@ class PreviewTab(BaseTab):
148
  minimum=0.0,
149
  maximum=1.0,
150
  step=0.01,
151
- value=0.7
 
152
  )
153
 
154
  self.components["inference_steps"] = gr.Slider(
@@ -156,12 +177,12 @@ class PreviewTab(BaseTab):
156
  minimum=1,
157
  maximum=100,
158
  step=1,
159
- value=30
160
  )
161
 
162
  self.components["enable_cpu_offload"] = gr.Checkbox(
163
  label="Enable Model CPU Offload (for low-VRAM GPUs)",
164
- value=True
165
  )
166
 
167
  self.components["generate_btn"] = gr.Button(
@@ -184,10 +205,37 @@ class PreviewTab(BaseTab):
184
  self.components["log"] = gr.TextArea(
185
  label="Generation Log",
186
  interactive=False,
187
- lines=15
188
  )
189
 
190
  return tab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
191
 
192
  def get_model_version_choices(self, model_type: str) -> List[str]:
193
  """Get model version choices based on model type"""
@@ -325,6 +373,13 @@ class PreviewTab(BaseTab):
325
  ]
326
  )
327
 
 
 
 
 
 
 
 
328
  # Load preview UI state when the tab is selected
329
  if hasattr(self.app, 'tabs_component') and self.app.tabs_component is not None:
330
  self.app.tabs_component.select(
@@ -343,7 +398,9 @@ class PreviewTab(BaseTab):
343
  self.components["lora_weight"],
344
  self.components["inference_steps"],
345
  self.components["enable_cpu_offload"],
346
- self.components["model_version"]
 
 
347
  ]
348
  )
349
 
@@ -351,7 +408,7 @@ class PreviewTab(BaseTab):
351
  for component_name in [
352
  "prompt", "negative_prompt", "prompt_prefix", "model_version", "resolution_preset",
353
  "width", "height", "num_frames", "fps", "guidance_scale", "flow_shift",
354
- "lora_weight", "inference_steps", "enable_cpu_offload"
355
  ]:
356
  if component_name in self.components:
357
  self.components[component_name].change(
@@ -378,7 +435,9 @@ class PreviewTab(BaseTab):
378
  self.components["inference_steps"],
379
  self.components["enable_cpu_offload"],
380
  self.components["fps"],
381
- self.components["conditioning_image"]
 
 
382
  ],
383
  outputs=[
384
  self.components["preview_video"],
@@ -444,6 +503,26 @@ class PreviewTab(BaseTab):
444
  if model_version not in model_version_choices and model_version_choices:
445
  model_version = model_version_choices[0]
446
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
447
  return (
448
  preview_state.get("prompt", ""),
449
  preview_state.get("negative_prompt", "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background"),
@@ -457,7 +536,9 @@ class PreviewTab(BaseTab):
457
  preview_state.get("lora_weight", 0.7),
458
  preview_state.get("inference_steps", 30),
459
  preview_state.get("enable_cpu_offload", True),
460
- model_version
 
 
461
  )
462
  except Exception as e:
463
  logger.error(f"Error loading preview state: {e}")
@@ -467,7 +548,9 @@ class PreviewTab(BaseTab):
467
  "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background",
468
  DEFAULT_PROMPT_PREFIX,
469
  832, 480, 49, 16, 5.0, 3.0, 0.7, 30, True,
470
- self.get_default_model_version(self.get_default_model_type())
 
 
471
  )
472
 
473
  def save_preview_state_value(self, value: Any) -> None:
@@ -522,7 +605,9 @@ class PreviewTab(BaseTab):
522
  inference_steps: int,
523
  enable_cpu_offload: bool,
524
  fps: int,
525
- conditioning_image: Optional[str] = None
 
 
526
  ) -> Tuple[Optional[str], str, str]:
527
  """Handler for generate button click, delegates to preview service"""
528
  # Save all the parameters to preview state before generating
@@ -550,7 +635,9 @@ class PreviewTab(BaseTab):
550
  "flow_shift": flow_shift,
551
  "lora_weight": lora_weight,
552
  "inference_steps": inference_steps,
553
- "enable_cpu_offload": enable_cpu_offload
 
 
554
  }
555
 
556
  state["preview"] = preview_state
@@ -564,7 +651,13 @@ class PreviewTab(BaseTab):
564
  # Initial UI update
565
  video_path, status, log = None, "Initializing generation...", "Starting video generation process..."
566
 
 
 
 
567
  # Start actual generation
 
 
 
568
  result = self.app.previewing.generate_video(
569
  model_type=model_type,
570
  model_version=model_version_id,
@@ -576,11 +669,12 @@ class PreviewTab(BaseTab):
576
  num_frames=num_frames,
577
  guidance_scale=guidance_scale,
578
  flow_shift=flow_shift,
579
- lora_weight=lora_weight,
580
  inference_steps=inference_steps,
581
  enable_cpu_offload=enable_cpu_offload,
582
  fps=fps,
583
- conditioning_image=conditioning_image
 
584
  )
585
 
586
  # Return final result
 
52
  value=DEFAULT_PROMPT_PREFIX
53
  )
54
 
55
+ # Ensure seed is interactive with a slider
56
  self.components["seed"] = gr.Slider(
57
  label="Generation Seed (-1 for random)",
58
  minimum=-1,
59
  maximum=2147483647, # 2^31 - 1
60
  step=1,
61
  value=-1,
62
+ info="Set to -1 for random seed or specific value for reproducible results",
63
+ interactive=True
64
  )
65
 
66
  with gr.Row():
 
83
  value=self.get_default_model_version(default_model)
84
  )
85
 
86
+ # Add dropdown to choose between LoRA and original model
87
+ has_lora = self.check_lora_model_exists()
88
+ lora_choices = []
89
+ default_lora_choice = ""
90
+
91
+ if has_lora:
92
+ lora_choices = ["Use LoRA model", "Use original model"]
93
+ default_lora_choice = "Use LoRA model"
94
+ else:
95
+ lora_choices = ["Cannot find LoRA model", "Use original model"]
96
+ default_lora_choice = "Use original model"
97
+
98
+ self.components["use_lora"] = gr.Dropdown(
99
+ choices=lora_choices,
100
+ label="Model Selection",
101
+ value=default_lora_choice
102
+ )
103
+
104
  # Add image input for image-to-video models
105
  self.components["conditioning_image"] = gr.Image(
106
  label="Conditioning Image (for Image-to-Video models)",
 
168
  minimum=0.0,
169
  maximum=1.0,
170
  step=0.01,
171
+ value=0.7,
172
+ visible=has_lora # Only visible if using LoRA
173
  )
174
 
175
  self.components["inference_steps"] = gr.Slider(
 
177
  minimum=1,
178
  maximum=100,
179
  step=1,
180
+ value=20
181
  )
182
 
183
  self.components["enable_cpu_offload"] = gr.Checkbox(
184
  label="Enable Model CPU Offload (for low-VRAM GPUs)",
185
+ value=False # let's assume user is using a video model training rig with a good GPU
186
  )
187
 
188
  self.components["generate_btn"] = gr.Button(
 
205
  self.components["log"] = gr.TextArea(
206
  label="Generation Log",
207
  interactive=False,
208
+ lines=20
209
  )
210
 
211
  return tab
212
+
213
+ def check_lora_model_exists(self) -> bool:
214
+ """Check if any LoRA model files exist in the output directory"""
215
+ # Look for the standard LoRA weights file
216
+ lora_path = OUTPUT_PATH / "pytorch_lora_weights.safetensors"
217
+ if lora_path.exists():
218
+ return True
219
+
220
+ # If not found in the expected location, try to find in checkpoints
221
+ checkpoints = list(OUTPUT_PATH.glob("checkpoint-*"))
222
+ if not checkpoints:
223
+ return False
224
+
225
+ for checkpoint in checkpoints:
226
+ lora_path = checkpoint / "pytorch_lora_weights.safetensors"
227
+ if lora_path.exists():
228
+ return True
229
+
230
+ return False
231
+
232
+ def update_lora_ui(self, use_lora_value: str) -> Dict[str, Any]:
233
+ """Update UI based on LoRA selection"""
234
+ is_using_lora = "Use LoRA model" in use_lora_value
235
+
236
+ return {
237
+ self.components["lora_weight"]: gr.Slider(visible=is_using_lora)
238
+ }
239
 
240
  def get_model_version_choices(self, model_type: str) -> List[str]:
241
  """Get model version choices based on model type"""
 
373
  ]
374
  )
375
 
376
+ # Connect LoRA selection dropdown to update LoRA weight visibility
377
+ self.components["use_lora"].change(
378
+ fn=self.update_lora_ui,
379
+ inputs=[self.components["use_lora"]],
380
+ outputs=[self.components["lora_weight"]]
381
+ )
382
+
383
  # Load preview UI state when the tab is selected
384
  if hasattr(self.app, 'tabs_component') and self.app.tabs_component is not None:
385
  self.app.tabs_component.select(
 
398
  self.components["lora_weight"],
399
  self.components["inference_steps"],
400
  self.components["enable_cpu_offload"],
401
+ self.components["model_version"],
402
+ self.components["seed"],
403
+ self.components["use_lora"]
404
  ]
405
  )
406
 
 
408
  for component_name in [
409
  "prompt", "negative_prompt", "prompt_prefix", "model_version", "resolution_preset",
410
  "width", "height", "num_frames", "fps", "guidance_scale", "flow_shift",
411
+ "lora_weight", "inference_steps", "enable_cpu_offload", "seed", "use_lora"
412
  ]:
413
  if component_name in self.components:
414
  self.components[component_name].change(
 
435
  self.components["inference_steps"],
436
  self.components["enable_cpu_offload"],
437
  self.components["fps"],
438
+ self.components["conditioning_image"],
439
+ self.components["seed"],
440
+ self.components["use_lora"]
441
  ],
442
  outputs=[
443
  self.components["preview_video"],
 
503
  if model_version not in model_version_choices and model_version_choices:
504
  model_version = model_version_choices[0]
505
 
506
+ # Check if LoRA exists and set appropriate dropdown options
507
+ has_lora = self.check_lora_model_exists()
508
+ use_lora = preview_state.get("use_lora", "")
509
+
510
+ # Validate use_lora value against current state
511
+ if has_lora:
512
+ valid_choices = ["Use LoRA model", "Use original model"]
513
+ if use_lora not in valid_choices:
514
+ use_lora = "Use LoRA model" # Default when LoRA exists
515
+ else:
516
+ valid_choices = ["Cannot find LoRA model", "Use original model"]
517
+ if use_lora not in valid_choices:
518
+ use_lora = "Use original model" # Default when no LoRA
519
+
520
+ # Update the dropdown choices in the UI
521
+ try:
522
+ self.components["use_lora"].choices = valid_choices
523
+ except Exception as e:
524
+ logger.error(f"Failed to update use_lora choices: {e}")
525
+
526
  return (
527
  preview_state.get("prompt", ""),
528
  preview_state.get("negative_prompt", "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background"),
 
536
  preview_state.get("lora_weight", 0.7),
537
  preview_state.get("inference_steps", 30),
538
  preview_state.get("enable_cpu_offload", True),
539
+ model_version,
540
+ preview_state.get("seed", -1),
541
+ use_lora
542
  )
543
  except Exception as e:
544
  logger.error(f"Error loading preview state: {e}")
 
548
  "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background",
549
  DEFAULT_PROMPT_PREFIX,
550
  832, 480, 49, 16, 5.0, 3.0, 0.7, 30, True,
551
+ self.get_default_model_version(self.get_default_model_type()),
552
+ -1,
553
+ "Use original model" if not self.check_lora_model_exists() else "Use LoRA model"
554
  )
555
 
556
  def save_preview_state_value(self, value: Any) -> None:
 
605
  inference_steps: int,
606
  enable_cpu_offload: bool,
607
  fps: int,
608
+ conditioning_image: Optional[str] = None,
609
+ seed: int = -1,
610
+ use_lora: str = "Use LoRA model"
611
  ) -> Tuple[Optional[str], str, str]:
612
  """Handler for generate button click, delegates to preview service"""
613
  # Save all the parameters to preview state before generating
 
635
  "flow_shift": flow_shift,
636
  "lora_weight": lora_weight,
637
  "inference_steps": inference_steps,
638
+ "enable_cpu_offload": enable_cpu_offload,
639
+ "seed": seed,
640
+ "use_lora": use_lora
641
  }
642
 
643
  state["preview"] = preview_state
 
651
  # Initial UI update
652
  video_path, status, log = None, "Initializing generation...", "Starting video generation process..."
653
 
654
+ # Set lora_path to None if not using LoRA
655
+ use_lora_model = use_lora == "Use LoRA model"
656
+
657
  # Start actual generation
658
+ # If not using LoRA, set lora_weight to 0 to disable it
659
+ effective_lora_weight = lora_weight if use_lora_model else 0.0
660
+
661
  result = self.app.previewing.generate_video(
662
  model_type=model_type,
663
  model_version=model_version_id,
 
669
  num_frames=num_frames,
670
  guidance_scale=guidance_scale,
671
  flow_shift=flow_shift,
672
+ lora_weight=effective_lora_weight, # Use 0.0 if not using LoRA
673
  inference_steps=inference_steps,
674
  enable_cpu_offload=enable_cpu_offload,
675
  fps=fps,
676
+ conditioning_image=conditioning_image,
677
+ seed=seed
678
  )
679
 
680
  # Return final result