feat: update workspace paths and enhance gitignore
- Updated stablediffusion crate path from "../stable-diffusion-burn" to "./crates/stable-diffusion-burn" for proper workspace resolution - Enhanced .gitignore to include generated model files (.mpk, .pt, .bin, .safetensors, .ckpt) and user_data directory - Added Cargo.lock to gitignore with appropriate comment - Reorganized IDE files section in gitignore for better clarity - Added newline at end of file for proper formatting
This commit is contained in:
100
crates/stable-diffusion-burn/python/save.py
Normal file
100
crates/stable-diffusion-burn/python/save.py
Normal file
@@ -0,0 +1,100 @@
|
||||
import pathlib
|
||||
import numpy as np
|
||||
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
def save_scalar(s, name, path):
|
||||
s = np.array([1.0, float(s)]).astype(np.float32)
|
||||
np.save(pathlib.Path(path, f'{name}.npy'), s)
|
||||
|
||||
def save_tensor(tensor, name, path):
|
||||
tensor_numpy = tensor.numpy()
|
||||
tensor_dims = np.array(tensor_numpy.shape)
|
||||
tensor_values = tensor_numpy.flatten()
|
||||
tensor_to_save = np.concatenate((tensor_dims, tensor_values)).astype(np.float32)
|
||||
np.save(pathlib.Path(path, f'{name}.npy'), tensor_to_save)
|
||||
|
||||
def save_linear(linear, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
save_tensor(linear.weight.transpose(), 'weight', path) # PyTorch and Tinygrad strangely transpose linear weights so reverse that
|
||||
if linear.bias is not None:
|
||||
save_tensor(linear.bias, 'bias', path)
|
||||
|
||||
def save_layer_norm(layer_norm, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
save_tensor(layer_norm.weight, 'weight', path)
|
||||
save_tensor(layer_norm.bias, 'bias', path)
|
||||
save_scalar(layer_norm.eps, 'eps', path)
|
||||
|
||||
def save_group_norm(layer_norm, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
if layer_norm.weight is not None:
|
||||
save_tensor(layer_norm.weight, 'weight', path)
|
||||
if layer_norm.bias is not None:
|
||||
save_tensor(layer_norm.bias, 'bias', path)
|
||||
save_scalar(layer_norm.eps, 'eps', path)
|
||||
save_scalar(layer_norm.num_groups, 'n_group', path)
|
||||
save_scalar(layer_norm.num_channels, 'n_channel', path)
|
||||
|
||||
def to_tuple_tensor(val):
|
||||
if isinstance(val, tuple):
|
||||
# Convert tuple to Tensor
|
||||
if len(val) == 1:
|
||||
return Tensor([val[0], val[0]])
|
||||
elif len(val) == 2:
|
||||
return Tensor([val[0], val[1]])
|
||||
else:
|
||||
raise ValueError('Tuple should be of length 1 or 2 only.')
|
||||
else:
|
||||
# Treat as scalar and convert to Tensor
|
||||
return Tensor([val, val])
|
||||
|
||||
def save_conv2d(conv2d, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
save_tensor(conv2d.weight, 'weight', path)
|
||||
if conv2d.bias is not None:
|
||||
save_tensor(conv2d.bias, 'bias', path)
|
||||
save_tensor(to_tuple_tensor(conv2d.stride), 'stride', path)
|
||||
save_tensor(to_tuple_tensor(conv2d.padding), 'padding', path)
|
||||
save_tensor(to_tuple_tensor(conv2d.dilation), 'dilation', path)
|
||||
save_scalar(conv2d.groups, "n_group", path)
|
||||
save_tensor(to_tuple_tensor(conv2d.kernel_size), 'kernel_size', path)
|
||||
|
||||
assert conv2d.groups == 1
|
||||
in_channels = conv2d.weight.shape[1]
|
||||
out_channels = conv2d.weight.shape[0]
|
||||
save_scalar(in_channels, "n_channels_in", path)
|
||||
save_scalar(out_channels, "n_channels_out", path)
|
||||
|
||||
def save_padded_conv2d(padded_conv2d, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Store conv2d layer weights
|
||||
orig_padding = padded_conv2d.padding
|
||||
padded_conv2d.padding = (0, 0)
|
||||
save_conv2d(padded_conv2d, f"{path}/conv")
|
||||
padded_conv2d.padding = orig_padding
|
||||
|
||||
# Dimensions: in-channels and out-channels
|
||||
assert padded_conv2d.groups == 1
|
||||
channels = (padded_conv2d.weight.shape[1], padded_conv2d.weight.shape[0])
|
||||
save_tensor(to_tuple_tensor(channels), 'channels', path)
|
||||
|
||||
assert len(padded_conv2d.kernel_size) == 1 or padded_conv2d.kernel_size[0] == padded_conv2d.kernel_size[1]
|
||||
save_scalar(padded_conv2d.kernel_size[0], 'kernel_size', path)
|
||||
|
||||
# Stride
|
||||
assert not isinstance(padded_conv2d.stride, tuple) or len(padded_conv2d.stride) == 1
|
||||
save_scalar(padded_conv2d.stride, 'stride', path)
|
||||
|
||||
# Padding
|
||||
padding = [padded_conv2d.padding[0], padded_conv2d.padding[1],
|
||||
padded_conv2d.padding[2], padded_conv2d.padding[3]]
|
||||
save_tensor(Tensor(padding), 'padding', path)
|
||||
|
||||
|
||||
def save_embedding(embedding, path):
|
||||
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
|
||||
save_tensor(embedding.weight, 'weight', path)
|
||||
|
||||
Reference in New Issue
Block a user