stable-diffusion-webui/modules/torch_utils.py
cuba3 9e60cdbc3f Maintaining Project Compatibility for Python 3.9 Users Without Upgrade Requirements.
Sole usage of Python 3.10's match-case in the project hinders quick-start for beginners; consider replacing with if-else for improved accessibility.
2024-06-25 15:24:46 +08:00

26 lines
684 B
Python

from __future__ import annotations
import torch.nn
import torch
def get_param(model) -> torch.nn.Parameter:
"""
Find the first parameter in a model or module.
"""
if hasattr(model, "model") and hasattr(model.model, "parameters"):
# Unpeel a model descriptor to get at the actual Torch module.
model = model.model
for param in model.parameters():
return param
raise ValueError(f"No parameters found in model {model!r}")
def float64(t: torch.Tensor):
"""return torch.float64 if device is not mps or xpu, else return torch.float32"""
if t.device.type in ['mps', 'xpu']:
return torch.float32
return torch.float64