John Pollock
235cd267bf
This commit introduces a new block swapping mechanism specifically for WanVideo models to enable running them on GPUs with limited VRAM. A new `WanVideoBlockSwapManager` is implemented which uses a pre-allocation or "shell" strategy. Instead of moving entire blocks between CPU and GPU, this approach: 1. Pre-allocates a single "shell" block on the GPU, sized to match the largest block in the model. 2. Offloads designated model blocks to the CPU. 3. Patches the `forward` method of these offloaded blocks. 4. During inference, the patched method copies the weights (`state_dict`) from the CPU block into the GPU shell just before execution. This method avoids the overhead of allocating and deallocating GPU memory for each block, reducing memory fragmentation and potentially improving performance and or corruption copying potentially modified blocks back to the swap space.
24 lines
931 B
Python
24 lines
931 B
Python
def get_model_type(model_patcher):
|
|
"""
|
|
Identifies the model type using a multi-layered approach for robustness.
|
|
It first checks the diffusion model's class name, then falls back to the
|
|
model_type enum.
|
|
"""
|
|
if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'diffusion_model'):
|
|
class_name = type(model_patcher.model.diffusion_model).__name__
|
|
|
|
# Prioritize class name for accuracy
|
|
if "Flux" in class_name:
|
|
return "FLUX"
|
|
if "Qwen" in class_name:
|
|
return "QWEN"
|
|
if "WanModel" in class_name:
|
|
return "WANVIDEO"
|
|
|
|
# Fallback to the model_type enum for other cases
|
|
if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'model_type'):
|
|
# model_type is an enum, so we return its name as a string
|
|
return model_patcher.model.model_type.name
|
|
|
|
return "UNKNOWN"
|