@@ -131,7 +131,7 @@ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
if ( shape_only ) {
return ;
}
int range = 1000 ;
int range = 3 ;
for ( int i = 0 ; i < tensor - > ne [ 3 ] ; i + + ) {
if ( i > = range & & i + range < tensor - > ne [ 3 ] ) {
continue ;
@@ -335,7 +335,7 @@ void sd_image_to_tensor(const uint8_t* image_data,
}
}
float sd _mean( struct ggml_tensor * src ) {
float ggml_tensor _mean( struct ggml_tensor * src ) {
float mean = 0.0f ;
int64_t nelements = ggml_nelements ( src ) ;
float * data = ( float * ) src - > data ;
@@ -345,7 +345,18 @@ float sd_mean(struct ggml_tensor* src) {
return mean ;
}
void sd_scale ( struct ggml_tensor * src , float scale ) {
// a = a+b
void ggml_tensor_add ( struct ggml_tensor * a , struct ggml_tensor * b ) {
GGML_ASSERT ( ggml_nelements ( a ) = = ggml_nelements ( b ) ) ;
int64_t nelements = ggml_nelements ( a ) ;
float * vec_a = ( float * ) a - > data ;
float * vec_b = ( float * ) b - > data ;
for ( int i = 0 ; i < nelements ; i + + ) {
vec_a [ i ] = vec_a [ i ] + vec_b [ i ] ;
}
}
void ggml_tensor_scale ( struct ggml_tensor * src , float scale ) {
int64_t nelements = ggml_nelements ( src ) ;
float * data = ( float * ) src - > data ;
for ( int i = 0 ; i < nelements ; i + + ) {
@@ -353,7 +364,7 @@ void sd_scale(struct ggml_tensor* src, float scale) {
}
}
void sd _clamp( struct ggml_tensor * src , float min , float max ) {
void ggml_tensor _clamp( struct ggml_tensor * src , float min , float max ) {
int64_t nelements = ggml_nelements ( src ) ;
float * data = ( float * ) src - > data ;
for ( int i = 0 ; i < nelements ; i + + ) {
@@ -363,7 +374,7 @@ void sd_clamp(struct ggml_tensor* src, float min, float max) {
}
// convert values from [0, 1] to [-1, 1]
void sd_convert _input( struct ggml_tensor * src ) {
void ggml_tensor_scale _input( struct ggml_tensor * src ) {
int64_t nelements = ggml_nelements ( src ) ;
float * data = ( float * ) src - > data ;
for ( int i = 0 ; i < nelements ; i + + ) {
@@ -373,7 +384,7 @@ void sd_convert_input(struct ggml_tensor* src) {
}
// convert values from [-1, 1] to [0, 1]
void sd_convert _output( struct ggml_tensor * src ) {
void ggml_tensor_scale _output( struct ggml_tensor * src ) {
int64_t nelements = ggml_nelements ( src ) ;
float * data = ( float * ) src - > data ;
for ( int i = 0 ; i < nelements ; i + + ) {
@@ -387,6 +398,64 @@ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx,
return ggml_group_norm ( ctx , a , 32 ) ;
}
struct ggml_tensor * ggml_nn_linear ( struct ggml_context * ctx ,
struct ggml_tensor * x ,
struct ggml_tensor * w ,
struct ggml_tensor * b ) {
x = ggml_mul_mat ( ctx , w , x ) ;
x = ggml_add ( ctx , x , b ) ;
return x ;
}
// w: [OC, IC, KH, KW]
// x: [N, IC, IH, IW]
// b: [OC,]
// result: [N, OC, OH, OW]
struct ggml_tensor * ggml_nn_conv_2d ( struct ggml_context * ctx ,
struct ggml_tensor * x ,
struct ggml_tensor * w ,
struct ggml_tensor * b ,
int s0 = 1 ,
int s1 = 1 ,
int p0 = 0 ,
int p1 = 0 ,
int d0 = 1 ,
int d1 = 1 ) {
x = ggml_conv_2d ( ctx , w , x , s0 , s1 , p0 , p1 , d0 , d1 ) ;
if ( b ! = NULL ) {
b = ggml_reshape_4d ( ctx , b , 1 , 1 , b - > ne [ 0 ] , 1 ) ;
x = ggml_add ( ctx , x , b ) ;
}
return x ;
}
struct ggml_tensor * ggml_nn_layer_norm ( struct ggml_context * ctx ,
struct ggml_tensor * x ,
struct ggml_tensor * w ,
struct ggml_tensor * b ,
float eps = EPS ) {
x = ggml_norm ( ctx , x , eps ) ;
x = ggml_mul ( ctx , x , w ) ;
x = ggml_add ( ctx , x , b ) ;
return x ;
}
struct ggml_tensor * ggml_nn_group_norm ( struct ggml_context * ctx ,
struct ggml_tensor * x ,
struct ggml_tensor * w ,
struct ggml_tensor * b ,
int num_groups = 32 ) {
if ( x - > n_dims = = 4 ) {
w = ggml_reshape_4d ( ctx , w , 1 , 1 , w - > ne [ 0 ] , 1 ) ;
b = ggml_reshape_4d ( ctx , b , 1 , 1 , b - > ne [ 0 ] , 1 ) ;
}
x = ggml_group_norm ( ctx , x , num_groups ) ;
x = ggml_mul ( ctx , x , w ) ;
x = ggml_add ( ctx , x , b ) ;
return x ;
}
std : : pair < std : : unordered_map < std : : string , float > , std : : string > extract_and_remove_lora ( std : : string text ) {
std : : regex re ( " <lora:([^:]+):([^>]+)> " ) ;
std : : smatch matches ;
@@ -738,30 +807,21 @@ struct ResidualAttentionBlock {
struct ggml_tensor * r = x ;
// layer norm 1
{
x = ggml_norm ( ctx , x , EPS ) ;
x = ggml_add ( ctx ,
ggml_mul ( ctx , x , ln1_w ) ,
ln1_b ) ;
}
x = ggml_nn_layer_norm ( ctx , x , ln1_w , ln1_b ) ;
// self-attention
{
struct ggml_tensor * q = ggml_add ( ctx ,
ggml_mul_mat ( ctx , q_w , x ) ,
q_b ) ;
struct ggml_tensor * q = ggml_nn_linear ( ctx , x , q_w , q_b ) ;
q = ggml_scale_inplace ( ctx , q , attn_scale ) ;
q = ggml_reshape_4d ( ctx , q , d_model , n_head , n_token , N ) ; // [N, n_token, n_head, d_model]
q = ggml_cont ( ctx , ggml_permute ( ctx , q , 0 , 2 , 1 , 3 ) ) ; // [N, n_head, n_token, d_model]
q = ggml_reshape_3d ( ctx , q , d_model , n_token , n_head * N ) ; // [N * n_head, n_token, d_model]
struct ggml_tensor * k = ggml_add ( ctx ,
ggml_mul_mat ( ctx , k_w , x ) , k_b ) ;
struct ggml_tensor * k = ggml_nn_linear ( ctx , x , k_w , k_b ) ;
k = ggml_reshape_4d ( ctx , k , d_model , n_head , n_token , N ) ; // [N, n_token, n_head, d_model]
k = ggml_cont ( ctx , ggml_permute ( ctx , k , 0 , 2 , 1 , 3 ) ) ; // [N, n_head, n_token, d_model]
k = ggml_reshape_3d ( ctx , k , d_model , n_token , n_head ) ; // [N * n_head, n_token, d_model]
struct ggml_tensor * v = ggml_add ( ctx ,
ggml_mul_mat ( ctx , v_w , x ) , v_b ) ;
struct ggml_tensor * v = ggml_nn_linear ( ctx , x , v_w , v_b ) ;
v = ggml_reshape_4d ( ctx , v , d_model , n_head , n_token , N ) ; // [N, n_token, n_head, d_model]
v = ggml_cont ( ctx , ggml_permute ( ctx , v , 1 , 2 , 0 , 3 ) ) ; // [N, n_head, d_model, n_token]
v = ggml_reshape_3d ( ctx , v , n_token , d_model , n_head * N ) ; // [N * n_head, d_model, n_token]
@@ -779,24 +839,17 @@ struct ResidualAttentionBlock {
}
// attention output
x = ggml_mul_mat ( ctx , out_w , x ) ;
x = ggml_add ( ctx , x , out_b ) ;
x = ggml_nn_linear ( ctx , x , out_w , out_b ) ;
// residual
x = ggml_add ( ctx , x , r ) ;
r = x ;
// layer norm 2
{
x = ggml_norm ( ctx , x , EPS ) ;
x = ggml_add ( ctx , ggml_mul ( ctx , x , ln2_w ) ,
ln2_b ) ;
}
x = ggml_nn_layer_norm ( ctx , x , ln2_w , ln2_b ) ;
// mlp
x = ggml_mul_mat ( ctx , fc1_w , x ) ;
x = ggml_add ( ctx , x , fc1_b ) ;
x = ggml_nn_linear ( ctx , x , fc1_w , fc1_b ) ;
if ( hidden_size = = 1024 ) { // SD 2.x
x = ggml_gelu_inplace ( ctx , x ) ;
@@ -804,8 +857,7 @@ struct ResidualAttentionBlock {
x = ggml_gelu_quick_inplace ( ctx , x ) ;
}
x = ggml_mul_mat ( ctx , fc2_w , x ) ;
x = ggml_add ( ctx , x , fc2_b ) ;
x = ggml_nn_linear ( ctx , x , fc2_w , fc2_b ) ;
// residual 2
x = ggml_add ( ctx , x , r ) ;
@@ -993,12 +1045,7 @@ struct CLIPTextModel {
}
// final layer norm
{
x = ggml_norm ( ctx0 , x , EPS ) ;
x = ggml_add ( ctx0 , ggml_mul ( ctx0 , x , final_ln_w ) ,
final_ln_b ) ;
}
x = ggml_nn_layer_norm ( ctx0 , x , final_ln_w , final_ln_b ) ;
return x ; // [N, n_token, hidden_size]
}
@@ -1252,48 +1299,29 @@ struct ResBlock {
// emb: [N, emb_channels]
// in_layers
// group norm 32
auto h = ggml_group_norm_32 ( ctx , x ) ;
h = ggml_ad d ( ctx ,
ggml_mul ( ctx ,
h ,
ggml_reshape_4d ( ctx , in_layer_0_w , 1 , 1 , in_layer_0_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , in_layer_0_b , 1 , 1 , in_layer_0_b - > ne [ 0 ] , 1 ) ) ;
// silu
h = ggml_silu_inplace ( ctx , h ) ;
// conv2d
h = ggml_conv_2d ( ctx , in_layer_2_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h ,
ggml_reshape_4d ( ctx , in_layer_2_b , 1 , 1 , in_layer_2_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w]
auto h = ggml_nn_ group_ norm( ctx , x , in_layer_0_w , in_layer_0_b ) ;
h = ggml_silu_inplace ( ctx , h ) ;
h = ggml_nn_conv_2 d ( ctx , h , in_layer_2_w , in_layer_2_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h, w]
// emb_layers
auto emb_out = ggml_silu ( ctx , emb ) ;
emb_out = ggml_mul_mat ( ctx , emb_layer_1_w , emb_out ) ;
emb_out = ggml_add ( ctx , emb_out , emb_layer_1_b ) ; // [N, out_channels]
emb_out = ggml_nn_linear ( ctx , emb_out , emb_layer_1_w , emb_layer_1_b ) ; // [N, out_channels]
emb_out = ggml_reshape_4d ( ctx , emb_out , 1 , 1 , emb_out - > ne [ 0 ] , emb_out - > ne [ 1 ] ) ; // [N, out_channels, 1, 1]
// out_layers
h = ggml_add ( ctx , h , emb_out ) ;
// group norm 32
h = ggml_group_norm_inplace ( ctx , h , 32 ) ;
h = ggml_add ( ctx ,
ggml_mul ( ctx , h , ggml_reshape_4d ( ctx , out_layer_0_w , 1 , 1 , out_layer_0_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , out_layer_0_b , 1 , 1 , out_layer_0_b - > ne [ 0 ] , 1 ) ) ;
// silu
h = ggml_nn_ group_ norm( ctx , h , out_layer_0_w , out_layer_0_b ) ;
h = ggml_silu_inplace ( ctx , h ) ;
// dropout, skip for inference
// conv2d
h = ggml_conv_2d ( ctx , out_layer_3_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , out_layer_3_b , 1 , 1 , out_layer_3_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w
h = ggml_nn_ conv_2d ( ctx , h , out_layer_3_w , out_layer_3_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h, w]
// skip connection
if ( out_channels ! = channels ) {
x = ggml_conv_2d ( ctx , skip_w , x , 1 , 1 , 0 , 0 , 1 , 1 ) ;
x = ggml_add ( ctx ,
x , ggml_reshape_4d ( ctx , skip_b , 1 , 1 , skip_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w]
x = ggml_nn_ conv_2d ( ctx , x , skip_w , skip_b ) ; // [N, out_channels, h, w]
}
h = ggml_add ( ctx , h , x ) ;
return h ; // [N, out_channels, h, w]
}
@@ -1468,15 +1496,9 @@ struct SpatialTransformer {
// x: [N, in_channels, h, w]
// context: [N, max_position, hidden_size(aka context_dim)]
auto x_in = x ;
// group norm 32
x = ggml_group_norm_32 ( ctx , x ) ;
x = ggml_add ( ctx ,
ggml_mul ( ctx , x , ggml_reshape_4d ( ctx , norm_w , 1 , 1 , norm_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , norm_b , 1 , 1 , norm_b - > ne [ 0 ] , 1 ) ) ;
x = ggml_nn_group_norm ( ctx , x , norm_w , norm_b ) ;
// proj_in
x = ggml_conv_2d ( ctx , proj_in_w , x , 1 , 1 , 0 , 0 , 1 , 1 ) ;
x = ggml_add ( ctx ,
x , ggml_reshape_4d ( ctx , proj_in_b , 1 , 1 , proj_in_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
x = ggml_nn_ conv_2d ( ctx , x , proj_in_w , proj_in_b ) ; // [N, in_channels, h, w]
// transformer
const int64_t n = x - > ne [ 3 ] ;
@@ -1489,13 +1511,8 @@ struct SpatialTransformer {
{
auto r = x ;
// layer norm 1
{
x = ggml_reshape_2d ( ctx , x , c , w * h * n ) ;
x = ggml_norm ( ctx , x , EPS ) ;
x = ggml_add ( ctx ,
ggml_mul ( ctx , x , transformer . norm1_w ) ,
transformer . norm1_b ) ;
}
x = ggml_reshape_2d ( ctx , x , c , w * h * n ) ;
x = ggml_nn_layer_norm ( ctx , x , transformer . norm1_w , transformer . norm1_b ) ;
// self-attention
{
@@ -1533,7 +1550,7 @@ struct SpatialTransformer {
// x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n));
x = ggml_reshape_2d ( ctx , kqv , d_head * n_head , h * w * n ) ;
x = ggml_add ( ctx , ggml_mul_mat ( ct x, transformer . attn1_out_w , x ) , transformer . attn1_out_b ) ;
x = ggml_nn_linear ( ctx , x , transformer . attn1_out_w , transformer . attn1_out_b ) ;
x = ggml_reshape_4d ( ctx , x , c , w , h , n ) ;
}
@@ -1542,11 +1559,7 @@ struct SpatialTransformer {
r = x ;
// layer norm 2
{
x = ggml_norm ( ctx , x , EPS ) ;
x = ggml_add ( ctx ,
ggml_mul ( ctx , x , transformer . norm2_w ) , transformer . norm2_b ) ;
}
x = ggml_nn_layer_norm ( ctx , x , transformer . norm2_w , transformer . norm2_b ) ;
// cross-attention
{
@@ -1584,7 +1597,7 @@ struct SpatialTransformer {
// x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n)); // [N * h * w, in_channels]
x = ggml_reshape_2d ( ctx , kqv , d_head * n_head , h * w * n ) ; // [N * h * w, in_channels]
x = ggml_add ( ctx , ggml_mul_mat ( ct x, transformer . attn2_out_w , x ) , transformer . attn2_out_b ) ;
x = ggml_nn_linear ( ctx , x , transformer . attn2_out_w , transformer . attn2_out_b ) ;
x = ggml_reshape_4d ( ctx , x , c , w , h , n ) ;
}
@@ -1593,13 +1606,8 @@ struct SpatialTransformer {
r = x ;
// layer norm 3
{
x = ggml_reshape_2d ( ctx , x , c , h * w * n ) ; // [N * h * w, in_channels]
x = ggml_norm ( ctx , x , EPS ) ;
x = ggml_add ( ctx ,
ggml_mul ( ctx , x , transformer . norm3_w ) ,
transformer . norm3_b ) ;
}
x = ggml_reshape_2d ( ctx , x , c , h * w * n ) ; // [N * h * w, in_channels]
x = ggml_nn_layer_norm ( ctx , x , transformer . norm3_w , transformer . norm3_b ) ;
// ff
{
@@ -1626,17 +1634,14 @@ struct SpatialTransformer {
transformer . ff_0_proj_b - > nb [ 0 ] * transformer . ff_0_proj_b - > ne [ 0 ] / 2 ) ; // [in_channels * 4, ]
x = ggml_reshape_2d ( ctx , x , c , w * h * n ) ;
auto x_in = x ;
x = ggml_mul_mat ( ctx , x_w , x_in ) ; // [N * h * w, in_channels * 4]
x = ggml_add ( ctx , x , x_b ) ;
auto gate = ggml_mul_mat ( ctx , gate_w , x_in ) ; // [N * h * w, in_channels * 4]
gate = ggml_add ( ctx , gate , gate_b ) ;
x = ggml_nn_linear ( ctx , x_in , x_w , x_b ) ; // [N * h * w, in_channels * 4]
auto gate = ggml_nn_linear ( ctx , x_in , gate_w , gate_b ) ; // [N * h * w, in_channels * 4]
gate = ggml_gelu_inplace ( ctx , gate ) ;
x = ggml_mul ( ctx , x , gate ) ; // [N * h * w, in_channels * 4]
// fc
x = ggml_mul_mat ( ct x , transformer . ff_2_w , x ) ; // [N * h * w, in_channels]
x = ggml_add ( ctx , x , transformer . ff_2_b ) ;
x = ggml_nn_linear ( ctx , x , transformer . ff_2_w , transformer . ff_2_b ) ; // [N * h * w, in_channels]
}
x = ggml_reshape_4d ( ctx , x , c , w , h , n ) ; // [N, h, w, in_channels]
@@ -1644,12 +1649,11 @@ struct SpatialTransformer {
// residual
x = ggml_add ( ctx , x , r ) ;
}
x = ggml_cont ( ctx , ggml_permute ( ctx , x , 2 , 0 , 1 , 3 ) ) ; // // [N, in_channels, h, w]
x = ggml_cont ( ctx , ggml_permute ( ctx , x , 2 , 0 , 1 , 3 ) ) ; // [N, in_channels, h, w]
// proj_out
x = ggml_conv_2d ( ctx , proj_out_w , x , 1 , 1 , 0 , 0 , 1 , 1 ) ;
x = ggml_add ( ctx ,
x , ggml_reshape_4d ( ctx , proj_out_b , 1 , 1 , proj_out_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
x = ggml_nn_ conv_2d ( ctx , x , proj_out_w , proj_out_b ) ; // [N, in_channels, h, w]
x = ggml_add ( ctx , x , x_in ) ;
return x ;
}
@@ -1690,17 +1694,14 @@ struct DownSample {
struct ggml_tensor * forward ( struct ggml_context * ctx , struct ggml_tensor * x ) {
// x: [N, channels, h, w]
struct ggml_tensor * c = nullptr ;
struct ggml_tensor * c = NULL ;
if ( vae_downsample ) {
c = ggml_pad ( ctx , x , 1 , 1 , 0 , 0 ) ;
c = ggml_conv_2d ( ctx , op_w , c , 2 , 2 , 0 , 0 , 1 , 1 );
c = ggml_nn_ conv_2d ( ctx , c , op_w , op_b , 2 , 2 , 0 , 0 ) ;
} else {
c = ggml_conv_2d ( ctx , op_w , x , 2 , 2 , 1 , 1 , 1 , 1 ) ;
c = ggml_nn_ conv_2d ( ctx , x , op_w , op_b , 2 , 2 , 1 , 1 ) ;
}
c = ggml_add ( ctx ,
c ,
ggml_reshape_4d ( ctx , op_b , 1 , 1 , op_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h/2, w/2]
return c ;
return c ; // [N, out_channels, h/2, w/2]
}
} ;
@@ -1732,11 +1733,8 @@ struct UpSample {
struct ggml_tensor * forward ( struct ggml_context * ctx , struct ggml_tensor * x ) {
// x: [N, channels, h, w]
x = ggml_upscale ( ctx , x , 2 ) ; // [N, channels, h*2, w*2]
x = ggml_conv_2d ( ctx , conv_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ;
x = ggml_add ( ctx ,
x ,
ggml_reshape_4d ( ctx , conv_b , 1 , 1 , conv_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h*2, w*2]
x = ggml_upscale ( ctx , x , 2 ) ; // [N, channels, h*2, w*2]
x = ggml_nn_ conv_2d ( ctx , x , conv_w , conv_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h*2, w*2]
return x ;
}
} ;
@@ -2201,14 +2199,10 @@ struct UNetModel {
}
// time_embed = nn.Sequential
auto emb = ggml_nn_linear ( ctx0 , t_emb , time_embed_0_w , time_embed_0_b ) ;
emb = ggml_silu_inplace ( ctx0 , emb ) ;
// Linear
auto emb = ggml_mul_mat ( ctx0 , time_embed_0 _w , t_emb ) ;
emb = ggml_add ( ctx0 , emb , time_embed_0_b ) ;
// nn.SiLU()
emb = ggml_silu_inplace ( ctx0 , emb ) ;
// Linear
emb = ggml_mul_mat ( ctx0 , time_embed_2_w , emb ) ;
emb = ggml_add ( ctx0 , emb , time_embed_2_b ) ; // [N, time_embed_dim]
emb = ggml_nn_linear ( ctx0 , emb , time_embed_2 _w , time _embed_2_b ) ; // [N, time_embed_dim]
// SDXL
// label_emd = nn.Sequential
@@ -2216,13 +2210,9 @@ struct UNetModel {
// param y: an [N] Tensor of labels, if class-conditional. (clip g)
// if(y != NULL) {
// auto y_emb = ggml_mul_mat (ctx, label_embed_0_w, y );
// y_emb = ggml_add(ctx, y_emb, label_embed_0_b);
// // nn.SiLU()
// auto y_emb = ggml_nn_linear (ctx, y, label_embed_0_w, label_embed_0_b );
// y_emb = ggml_silu_inplace(ctx, y_emb);
// // Linear
// y_emb = ggml_mul_mat(ctx, label_embed_2_w, y_emb);
// y_emb = ggml_add(ctx, y_emb, label_embed_2_b);
// y_emb = ggml_nn_linear(ctx, y_emb, label_embed_2_w, label_embed_2_b);
// emb = ggml_add(ctx, emb, y_emb);
// }
@@ -2230,11 +2220,8 @@ struct UNetModel {
std : : vector < struct ggml_tensor * > hs ;
// input block 0
struct ggml_tensor * h = ggml_conv_2d ( ctx0 , input_block_0_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ; // [N, model_channels, h, w]
struct ggml_tensor * h = ggml_nn_ conv_2d ( ctx0 , x , input_block_0_w , input_block_0_b , 1 , 1 , 1 , 1 ) ; // [N, model_channels, h, w]
h = ggml_add ( ctx0 ,
h ,
ggml_reshape_4d ( ctx0 , input_block_0_b , 1 , 1 , input_block_0_b - > ne [ 0 ] , 1 ) ) ; // [N, model_channels, h, w]
ggml_set_name ( h , " bench-start " ) ;
hs . push_back ( h ) ;
// input block 1-11
@@ -2284,18 +2271,11 @@ struct UNetModel {
}
// out
// group norm 32
h = ggml_group_norm_32 ( ctx0 , h ) ;
h = ggml_add ( ctx0 ,
ggml_mul ( ctx0 , h , ggml_reshape_4d ( ctx0 , out_0_w , 1 , 1 , out_0_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx0 , out_0_b , 1 , 1 , out_0_b - > ne [ 0 ] , 1 ) ) ;
// silu
h = ggml_nn_ group_ norm( ctx0 , h , out_0_w , out_0_b ) ;
h = ggml_silu_inplace ( ctx0 , h ) ;
// conv2d
h = ggml_conv_2d ( ctx0 , out_2_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx0 ,
h , ggml_reshape_4d ( ctx0 , out_2_b , 1 , 1 , out_2_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w]
h = ggml_nn_ conv_2d ( ctx0 , h , out_2_w , out_2_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h, w]
ggml_set_name ( h , " bench-end " ) ;
return h ;
}
@@ -2492,38 +2472,19 @@ struct ResnetBlock {
struct ggml_tensor * forward ( struct ggml_context * ctx , struct ggml_tensor * z ) {
// z: [N, in_channels, h, w]
// group norm 32
auto h = ggml_group_norm_32 ( ctx , z ) ;
h = ggml_mul ( ctx ,
h , ggml_reshape_4d ( ctx , norm1 _w , 1 , 1 , norm1_w - > ne [ 0 ] , 1 ) ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , norm1_b , 1 , 1 , norm1_b - > ne [ 0 ] , 1 ) ) ;
// silu
h = ggml_silu_inplace ( ctx , h ) ;
// conv2d
h = ggml_conv_2d ( ctx , conv1_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv1_b , 1 , 1 , conv1_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w]
// group norm 32
h = ggml_group_norm_32 ( ctx , h ) ;
h = ggml_add ( ctx ,
ggml_mul ( ctx , h , ggml_reshape_4d ( ctx , norm2_w , 1 , 1 , norm2_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , norm2_b , 1 , 1 , norm2_b - > ne [ 0 ] , 1 ) ) ;
// silu
h = ggml_silu_inplace ( ctx , h ) ;
auto h = ggml_nn_group_norm ( ctx , z , norm1_w , norm1_b ) ;
h = ggml_silu_inplace ( ctx , h ) ;
h = ggml_nn_conv_2d ( ctx , h , conv1_w , conv1_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h, w]
h = ggml_nn_group_norm ( ctx , h , norm2 _w , norm2_b ) ;
h = ggml_silu_inplace ( ctx , h ) ;
// dropout, skip for inference
// conv2d
h = ggml_conv_2d ( ctx , conv2_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv2_b , 1 , 1 , conv2_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w
h = ggml_nn_ conv_ 2d( ctx , h , conv2_w , conv2_b , 1 , 1 , 1 , 1 ) ; // [N, out_channels, h, w]
// skip connection
if ( out_channels ! = in_channels ) {
z = ggml_conv_2d ( ctx , nin_shortcut_w , z , 1 , 1 , 0 , 0 , 1 , 1 ) ;
z = ggml_add ( ctx ,
z , ggml_reshape_4d ( ctx , nin_shortcut_b , 1 , 1 , nin_shortcut_b - > ne [ 0 ] , 1 ) ) ; // [N, out_channels, h, w]
z = ggml_nn_ conv_2d ( ctx , z , nin_shortcut_w , nin_shortcut_b ) ; // [N, out_channels, h, w]
}
h = ggml_add ( ctx , h , z ) ;
return h ; // [N, out_channels, h, w]
}
@@ -2593,30 +2554,16 @@ struct AttnBlock {
struct ggml_tensor * forward ( struct ggml_context * ctx , struct ggml_tensor * x ) {
// x: [N, in_channels, h, w]
// group norm 32
auto h_ = ggml_group_norm_32 ( ctx , x ) ;
h_ = ggml_add ( ctx ,
ggml_mul ( ctx , h_ , ggml_reshape_4d ( ctx , norm_w , 1 , 1 , norm_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , norm_b , 1 , 1 , norm_b - > ne [ 0 ] , 1 ) ) ;
auto h_ = ggml_nn_group_norm ( ctx , x , norm_w , norm_b ) ;
const int64_t n = h_ - > ne [ 3 ] ;
const int64_t c = h_ - > ne [ 2 ] ;
const int64_t h = h_ - > ne [ 1 ] ;
const int64_t w = h_ - > ne [ 0 ] ;
// q
auto q = ggml_conv_2d ( ctx , q_w , h_ , 1 , 1 , 0 , 0 , 1 , 1 ) ;
q = ggml_add ( ctx ,
q , ggml_reshape_4d ( ctx , q_b , 1 , 1 , q_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
// k
auto k = ggml_conv_2d ( ctx , k_w , h_ , 1 , 1 , 0 , 0 , 1 , 1 ) ;
k = ggml_ad d ( ctx ,
k , ggml_reshape_4d ( ctx , k_b , 1 , 1 , k_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
// v
auto v = ggml_conv_2d ( ctx , v_w , h_ , 1 , 1 , 0 , 0 , 1 , 1 ) ;
v = ggml_add ( ctx ,
v , ggml_reshape_4d ( ctx , v_b , 1 , 1 , v_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
auto q = ggml_nn_conv_2d ( ctx , h_ , q_w , q_b ) ; // [N, in_channels, h, w]
auto k = ggml_nn_ conv_2d ( ctx , h_ , k_w , k_b ) ; // [N, in_channels, h, w]
auto v = ggml_nn_conv_2 d ( ctx , h_ , v_w , v_b ) ; // [N, in_channels, h, w]
q = ggml_cont ( ctx , ggml_permute ( ctx , q , 1 , 2 , 0 , 3 ) ) ; // [N, h, w, in_channels]
q = ggml_reshape_3d ( ctx , q , c , h * w , n ) ; // [N, h * w, in_channels]
@@ -2634,9 +2581,8 @@ struct AttnBlock {
h_ = ggml_reshape_4d ( ctx , h_ , w , h , c , n ) ; // [N, in_channels, h, w]
// proj_out
h_ = ggml_conv_2d ( ctx , proj_out_w , h_ , 1 , 1 , 0 , 0 , 1 , 1 ) ;
h_ = ggml_add ( ctx ,
h_ , ggml_reshape_4d ( ctx , proj_out_b , 1 , 1 , proj_out_b - > ne [ 0 ] , 1 ) ) ; // [N, in_channels, h, w]
h_ = ggml_nn_ conv_2d ( ctx , h_ , proj_out_w , proj_out_b ) ; // [N, in_channels, h, w]
h_ = ggml_add ( ctx , h_ , x ) ;
return h_ ;
}
@@ -2803,9 +2749,7 @@ struct Encoder {
// x: [N, in_channels, h, w]
// conv_in
auto h = ggml_conv_2d ( ctx , conv_in_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv_in_b , 1 , 1 , conv_in_b - > ne [ 0 ] , 1 ) ) ; // [N, ch, h, w]
auto h = ggml_nn_ conv_2d ( ctx , x , conv_in_w , conv_in_b , 1 , 1 , 1 , 1 ) ; // [N, ch, h, w]
ggml_set_name ( h , " b-start " ) ;
int len_mults = sizeof ( ch_mult ) / sizeof ( int ) ;
for ( int i = 0 ; i < len_mults ; i + + ) {
@@ -2821,20 +2765,11 @@ struct Encoder {
h = mid . attn_1 . forward ( ctx , h ) ;
h = mid . block_2 . forward ( ctx , h ) ; // [N, block_in, h, w]
// group norm 32
h = ggml_group_norm_32 ( ctx , h ) ;
h = ggml_add ( ctx ,
ggml_mul ( ctx , h , ggml_reshape_4d ( ctx , norm_out_w , 1 , 1 , norm_out_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , norm_out_b , 1 , 1 , norm_out_b - > ne [ 0 ] , 1 ) ) ;
// silu
// silu
h = ggml_nn_group_norm ( ctx , h , norm_out_w , norm_out_b ) ;
h = ggml_silu_inplace ( ctx , h ) ;
// conv_out
h = ggml_conv_2d ( ctx , conv_out_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv_out_b , 1 , 1 , conv_out_b - > ne [ 0 ] , 1 ) ) ; // [N, z_channels*2, h, w]
h = ggml_nn_ conv_2d ( ctx , h , conv_out_w , conv_out_b , 1 , 1 , 1 , 1 ) ; // [N, z_channels*2, h, w]
return h ;
}
@@ -2996,9 +2931,7 @@ struct Decoder {
struct ggml_tensor * forward ( struct ggml_context * ctx , struct ggml_tensor * z ) {
// z: [N, z_channels, h, w]
// conv_in
auto h = ggml_conv_2d ( ctx , conv_in_w , z , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv_in_b , 1 , 1 , conv_in_b - > ne [ 0 ] , 1 ) ) ; // [N, block_in, h, w]
auto h = ggml_nn_ conv_2d ( ctx , z , conv_in_w , conv_in_b , 1 , 1 , 1 , 1 ) ; // [N, block_in, h, w]
h = mid . block_1 . forward ( ctx , h ) ;
h = mid . attn_1 . forward ( ctx , h ) ;
@@ -3015,19 +2948,11 @@ struct Decoder {
}
// group norm 32
h = ggml_group_norm_32 ( ctx , h ) ;
h = ggml_add ( ctx ,
ggml_mul ( ctx , h , ggml_reshape_4d ( ctx , norm_out_w , 1 , 1 , norm_out_w - > ne [ 0 ] , 1 ) ) ,
ggml_reshape_4d ( ctx , norm_out_b , 1 , 1 , norm_out_b - > ne [ 0 ] , 1 ) ) ;
// silu
// silu
h = ggml_nn_ group_norm ( ctx , h , norm_out_w , norm_out_b );
h = ggml_silu_inplace ( ctx , h ) ;
// conv_out
h = ggml_conv_2d ( ctx , conv_out_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx ,
h , ggml_reshape_4d ( ctx , conv_out_b , 1 , 1 , conv_out_b - > ne [ 0 ] , 1 ) ) ; // [N, out_ch, h, w]
h = ggml_nn_ conv_2d ( ctx , h , conv_out_w , conv_out_b , 1 , 1 , 1 , 1 ) ; // [N, out_ch, h, w]
return h ;
}
} ;
@@ -3176,9 +3101,7 @@ struct AutoEncoderKL {
struct ggml_tensor * decode ( struct ggml_context * ctx0 , struct ggml_tensor * z ) {
// z: [N, z_channels, h, w]
// post_quant_conv
auto h = ggml_conv_2d ( ctx0 , post_quant_conv_w , z , 1 , 1 , 0 , 0 , 1 , 1 ) ;
h = ggml_add ( ctx0 ,
h , ggml_reshape_4d ( ctx0 , post_quant_conv_b , 1 , 1 , post_quant_conv_b - > ne [ 0 ] , 1 ) ) ; // [N, z_channels, h, w]
auto h = ggml_nn_ conv_2d ( ctx0 , z , post_quant_conv_w , post_quant_conv_b ) ; // [N, z_channels, h, w]
ggml_set_name ( h , " bench-start " ) ;
h = decoder . forward ( ctx0 , h ) ;
ggml_set_name ( h , " bench-end " ) ;
@@ -3189,10 +3112,7 @@ struct AutoEncoderKL {
// x: [N, in_channels, h, w]
auto h = encoder . forward ( ctx0 , x ) ; // [N, 2*z_channels, h/8, w/8]
// quant_conv
h = ggml_conv_2d ( ctx0 , quant_conv_w , h , 1 , 1 , 0 , 0 , 1 , 1 ) ;
h = ggml_add ( ctx0 ,
h ,
ggml_reshape_4d ( ctx0 , quant_conv_b , 1 , 1 , quant_conv_b - > ne [ 0 ] , 1 ) ) ; // [N, 2*embed_dim, h/8, w/8]
h = ggml_nn_ conv_2d ( ctx0 , h , quant_conv_w , quant_conv_b ) ; // [N, 2*embed_dim, h/8, w/8]
ggml_set_name ( h , " b-end " ) ;
return h ;
}
@@ -3356,25 +3276,16 @@ struct TAEBlock {
ggml_tensor * forward ( ggml_context * ctx , ggml_tensor * x ) {
// conv(n_in, n_out)
ggml_tensor * h ;
h = ggml_conv_2d ( ctx , conv_0_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx , h , ggml_reshape_4d ( ctx , conv_0_b , 1 , 1 , conv_0_b - > ne [ 0 ] , 1 ) ) ;
// relu
h = ggml_nn_ conv_2d ( ctx , x , conv_0_w , conv_0_b , 1 , 1 , 1 , 1 ) ;
h = ggml_relu_inplace ( ctx , h ) ;
h = ggml_conv_2d ( ctx , conv_1_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx , h , ggml_reshape_4d ( ctx , conv_1_b , 1 , 1 , conv_1_b - > ne [ 0 ] , 1 ) ) ;
// relu
h = ggml_nn_conv_2d ( ctx , h , conv_1_w , conv_1_b , 1 , 1 , 1 , 1 ) ;
h = ggml_relu_inplace ( ctx , h ) ;
h = ggml_conv_2d ( ctx , conv_2_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx , h , ggml_reshape_4d ( ctx , conv_2_b , 1 , 1 , conv_2_b - > ne [ 0 ] , 1 ) ) ;
h = ggml_nn_conv_2d ( ctx , h , conv_2_w , conv_2_b , 1 , 1 , 1 , 1 ) ;
// skip connection
if ( in_channels ! = out_channels ) {
// skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
x = ggml_conv_2d ( ctx , conv_skip_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ;
x = ggml_nn_ conv_2d ( ctx , x , conv_skip_w , NULL , 1 , 1 , 1 , 1 ) ;
}
h = ggml_add ( ctx , h , x ) ;
@@ -3503,14 +3414,13 @@ struct TinyEncoder {
ggml_tensor * forward ( ggml_context * ctx , ggml_tensor * x ) {
// conv(3, 64)
auto z = ggml_conv_2d ( ctx , conv_input_w , x , 1 , 1 , 1 , 1 , 1 , 1 ) ;
z = ggml_add ( ctx , z , ggml_reshape_4d ( ctx , conv_input_b , 1 , 1 , conv_input_b - > ne [ 0 ] , 1 ) ) ;
auto z = ggml_nn_ conv_2d ( ctx , x , conv_input_w , conv_input_b , 1 , 1 , 1 , 1 ) ;
// Block(64, 64)
z = initial_block . forward ( ctx , z ) ;
// conv(64, 64, stride=2, bias=False)
z = ggml_conv_2d ( ctx , conv_1_w , z , 2 , 2 , 1 , 1 , 1 , 1 ) ;
z = ggml_nn_ conv_2d ( ctx , z , conv_1_w , NULL , 2 , 2 , 1 , 1 ) ;
// Block(64, 64), Block(64, 64), Block(64, 64)
for ( int i = 0 ; i < num_blocks ; i + + ) {
@@ -3518,7 +3428,7 @@ struct TinyEncoder {
}
// conv(64, 64, stride=2, bias=False)
z = ggml_conv_2d ( ctx , conv_2_w , z , 2 , 2 , 1 , 1 , 1 , 1 ) ;
z = ggml_nn_ conv_2d ( ctx , z , conv_2_w , NULL , 2 , 2 , 1 , 1 ) ;
// Block(64, 64), Block(64, 64), Block(64, 64)
for ( int i = 0 ; i < num_blocks ; i + + ) {
@@ -3526,7 +3436,7 @@ struct TinyEncoder {
}
// conv(64, 64, stride=2, bias=False)
z = ggml_conv_2d ( ctx , conv_3_w , z , 2 , 2 , 1 , 1 , 1 , 1 ) ;
z = ggml_nn_ conv_2d ( ctx , z , conv_3_w , NULL , 2 , 2 , 1 , 1 ) ;
// Block(64, 64), Block(64, 64), Block(64, 64)
for ( int i = 0 ; i < num_blocks ; i + + ) {
@@ -3534,8 +3444,7 @@ struct TinyEncoder {
}
// conv(64, 4)
z = ggml_conv_2d ( ctx , conv_final_w , z , 1 , 1 , 1 , 1 , 1 , 1 ) ;
z = ggml_add ( ctx , z , ggml_reshape_4d ( ctx , conv_final_b , 1 , 1 , conv_final_b - > ne [ 0 ] , 1 ) ) ;
z = ggml_nn_ conv_2d ( ctx , z , conv_final_w , conv_final_b , 1 , 1 , 1 , 1 ) ;
return z ;
}
} ;
@@ -3683,8 +3592,7 @@ struct TinyDecoder {
h = ggml_scale ( ctx , h , in_scale_3 ) ;
// conv(4, 64)
h = ggml_conv_2d ( ctx , conv_input_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx , h , ggml_reshape_4d ( ctx , conv_input_b , 1 , 1 , conv_input_b - > ne [ 0 ] , 1 ) ) ;
h = ggml_nn_ conv_2d ( ctx , h , conv_input_w , conv_input_b , 1 , 1 , 1 , 1 ) ;
// nn.ReLU()
h = ggml_relu_inplace ( ctx , h ) ;
@@ -3698,7 +3606,7 @@ struct TinyDecoder {
h = ggml_upscale ( ctx , h , 2 ) ;
// conv(64, 64, bias=False)
h = ggml_conv_2d ( ctx , conv_1_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_nn_ conv_2d ( ctx , h , conv_1_w , NULL , 1 , 1 , 1 , 1 ) ;
// Block(64, 64), Block(64, 64), Block(64, 64)
for ( int i = 0 ; i < num_blocks ; i + + ) {
@@ -3709,7 +3617,7 @@ struct TinyDecoder {
h = ggml_upscale ( ctx , h , 2 ) ;
// conv(64, 64, bias=False)
h = ggml_conv_2d ( ctx , conv_2_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_nn_ conv_2d ( ctx , h , conv_2_w , NULL , 1 , 1 , 1 , 1 ) ;
// Block(64, 64), Block(64, 64), Block(64, 64)
for ( int i = 0 ; i < num_blocks ; i + + ) {
@@ -3720,14 +3628,13 @@ struct TinyDecoder {
h = ggml_upscale ( ctx , h , 2 ) ;
// conv(64, 64, bias=False)
h = ggml_conv_2d ( ctx , conv_3_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_nn_ conv_2d ( ctx , h , conv_3_w , NULL , 1 , 1 , 1 , 1 ) ;
// Block(64, 64)
h = final_block . forward ( ctx , h ) ;
// conv(64, 3)
h = ggml_conv_2d ( ctx , conv_final_w , h , 1 , 1 , 1 , 1 , 1 , 1 ) ;
h = ggml_add ( ctx , h , ggml_reshape_4d ( ctx , conv_final_b , 1 , 1 , conv_final_b - > ne [ 0 ] , 1 ) ) ;
h = ggml_nn_ conv_2d ( ctx , h , conv_final_w , conv_final_b , 1 , 1 , 1 , 1 ) ;
return h ;
}
} ;
@@ -4724,7 +4631,7 @@ public:
LOG_DEBUG ( " computing condition graph completed, taking % " PRId64 " ms " , t1 - t0 ) ;
ggml_tensor * result = ggml_dup_tensor ( work_ctx , hidden_states ) ;
{
float original_mean = sd _mean( hidden_states ) ;
float original_mean = ggml_tensor _mean( hidden_states ) ;
for ( int i2 = 0 ; i2 < hidden_states - > ne [ 2 ] ; i2 + + ) {
for ( int i1 = 0 ; i1 < hidden_states - > ne [ 1 ] ; i1 + + ) {
for ( int i0 = 0 ; i0 < hidden_states - > ne [ 0 ] ; i0 + + ) {
@@ -4734,16 +4641,17 @@ public:
}
}
}
float new_mean = sd _mean( result ) ;
sd _scale( result , ( original_mean / new_mean ) ) ;
float new_mean = ggml_tensor _mean( result ) ;
ggml_tensor _scale( result , ( original_mean / new_mean ) ) ;
}
return result ; // [1, 77, 768]
}
ggml_tensor * sample ( ggml_context * work_ctx ,
ggml_tensor * x_t ,
ggml_tensor * positiv e,
ggml_tensor * negative ,
ggml_tensor * nois e,
ggml_tensor * c ,
ggml_tensor * uc ,
float cfg_scale ,
SampleMethod method ,
const std : : vector < float > & sigmas ) {
@@ -4756,12 +4664,18 @@ public:
struct ggml_tensor * noised_input = ggml_dup_tensor ( work_ctx , x_t ) ;
struct ggml_tensor * timesteps = ggml_new_tensor_1d ( work_ctx , GGML_TYPE_F32 , 1 ) ; // [N, ]
struct ggml_tensor * t_emb = new_timestep_embedding ( work_ctx , NULL , timesteps , diffusion_model . model_channels ) ; // [N, model_channels]
diffusion_model . begin ( noised_input , positive , t_emb ) ;
diffusion_model . begin ( noised_input , c , t_emb ) ;
bool has_unconditioned = cfg_scale ! = 1.0 & & negative ! = NULL ;
bool has_unconditioned = cfg_scale ! = 1.0 & & uc ! = NULL ;
// x = x * sigmas[0]
sd_scale ( x , sigmas [ 0 ] ) ;
if ( noise = = NULL ) {
// x = x * sigmas[0]
ggml_tensor_scale ( x , sigmas [ 0 ] ) ;
} else {
// xi = x + noise * sigma_sched[0]
ggml_tensor_scale ( noise , sigmas [ 0 ] ) ;
ggml_tensor_add ( x , noise ) ;
}
// denoise wrapper
struct ggml_tensor * out_cond = ggml_dup_tensor ( work_ctx , x ) ;
@@ -4797,15 +4711,15 @@ public:
copy_ggml_tensor ( noised_input , input ) ;
// noised_input = noised_input * c_in
sd _scale( noised_input , c_in ) ;
ggml_tensor _scale( noised_input , c_in ) ;
// cond
diffusion_model . compute ( out_cond , n_threads , noised_input , NULL , positive , t_emb ) ;
diffusion_model . compute ( out_cond , n_threads , noised_input , NULL , c , t_emb ) ;
float * negative_data = NULL ;
if ( has_unconditioned ) {
// uncond
diffusion_model . compute ( out_uncond , n_threads , noised_input , NULL , negative , t_emb ) ;
diffusion_model . compute ( out_uncond , n_threads , noised_input , NULL , uc , t_emb ) ;
negative_data = ( float * ) out_uncond - > data ;
}
float * vec_denoised = ( float * ) denoised - > data ;
@@ -5260,15 +5174,15 @@ public:
int64_t t0 = ggml_time_ms ( ) ;
if ( ! use_tiny_autoencoder ) {
if ( decode ) {
sd _scale( x , 1.0f / scale_factor ) ;
ggml_tensor _scale( x , 1.0f / scale_factor ) ;
} else {
sd_convert _input( x ) ;
ggml_tensor_scale _input( x ) ;
}
first_stage_model . begin ( x , decode ) ;
first_stage_model . compute ( result , n_threads , x , decode ) ;
first_stage_model . end ( ) ;
if ( decode ) {
sd_convert _output( result ) ;
ggml_tensor_scale _output( result ) ;
}
} else {
tae_first_stage . begin ( x , decode ) ;
@@ -5278,10 +5192,18 @@ public:
int64_t t1 = ggml_time_ms ( ) ;
LOG_DEBUG ( " computing vae [mode: %s] graph completed, taking %.2fs " , decode ? " DECODE " : " ENCODE " , ( t1 - t0 ) * 1.0f / 1000 ) ;
if ( decode ) {
sd _clamp( result , 0.0f , 1.0f ) ;
ggml_tensor _clamp( result , 0.0f , 1.0f ) ;
}
return result ;
}
ggml_tensor * encode_first_stage ( ggml_context * work_ctx , ggml_tensor * x ) {
return compute_first_stage ( work_ctx , x , false ) ;
}
ggml_tensor * decode_first_stage ( ggml_context * work_ctx , ggml_tensor * x ) {
return compute_first_stage ( work_ctx , x , true ) ;
}
} ;
/*================================================= StableDiffusion ==================================================*/
@@ -5358,11 +5280,11 @@ std::vector<uint8_t*> StableDiffusion::txt2img(std::string prompt,
seed = rand ( ) ;
}
t0 = ggml_time_ms ( ) ;
ggml_tensor * postive = sd - > get_learned_condition ( work_ctx , prompt ) ;
struct ggml_tensor * negative = NULL ;
t0 = ggml_time_ms ( ) ;
ggml_tensor * c = sd - > get_learned_condition ( work_ctx , prompt ) ;
struct ggml_tensor * uc = NULL ;
if ( cfg_scale ! = 1.0 ) {
negative = sd - > get_learned_condition ( work_ctx , negative_prompt ) ;
uc = sd - > get_learned_condition ( work_ctx , negative_prompt ) ;
}
t1 = ggml_time_ms ( ) ;
LOG_INFO ( " get_learned_condition completed, taking % " PRId64 " ms " , t1 - t0 ) ;
@@ -5387,7 +5309,7 @@ std::vector<uint8_t*> StableDiffusion::txt2img(std::string prompt,
std : : vector < float > sigmas = sd - > denoiser - > schedule - > get_sigmas ( sample_steps ) ;
struct ggml_tensor * x_0 = sd - > sample ( work_ctx , x_t , postive , negative , cfg_scale , sample_method , sigmas ) ;
struct ggml_tensor * x_0 = sd - > sample ( work_ctx , x_t , NULL , c , uc , cfg_scale , sample_method , sigmas ) ;
// struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
// print_ggml_tensor(x_0);
int64_t sampling_end = ggml_time_ms ( ) ;
@@ -5404,7 +5326,7 @@ std::vector<uint8_t*> StableDiffusion::txt2img(std::string prompt,
LOG_INFO ( " decoding %zu latents " , final_latents . size ( ) ) ;
for ( size_t i = 0 ; i < final_latents . size ( ) ; i + + ) {
t1 = ggml_time_ms ( ) ;
struct ggml_tensor * img = sd - > comput e_first_stage( work_ctx , final_latents [ i ] /* x_0 */ , true );
struct ggml_tensor * img = sd - > decod e_first_stage( work_ctx , final_latents [ i ] /* x_0 */ ) ;
if ( img ! = NULL ) {
results . push_back ( sd_tensor_to_image ( img ) ) ;
}
@@ -5483,10 +5405,10 @@ std::vector<uint8_t*> StableDiffusion::img2img(const uint8_t* init_img_data,
t0 = ggml_time_ms ( ) ;
ggml_tensor * init_latent = NULL ;
if ( ! sd - > use_tiny_autoencoder ) {
ggml_tensor * moments = sd - > comput e_first_stage( work_ctx , init_img , false );
ggml_tensor * moments = sd - > encod e_first_stage( work_ctx , init_img ) ;
init_latent = sd - > get_first_stage_encoding ( work_ctx , moments ) ;
} else {
init_latent = sd - > comput e_first_stage( work_ctx , init_img , false );
init_latent = sd - > encod e_first_stage( work_ctx , init_img ) ;
}
// print_ggml_tensor(init_latent);
t1 = ggml_time_ms ( ) ;
@@ -5507,8 +5429,12 @@ std::vector<uint8_t*> StableDiffusion::img2img(const uint8_t* init_img_data,
// requires encode_adm
// apply set_timestep_embedding with dim 256
sd - > rng - > manual_seed ( seed ) ;
struct ggml_tensor * noise = ggml_dup_tensor ( work_ctx , init_latent ) ;
ggml_tensor_set_f32_randn ( noise , sd - > rng ) ;
LOG_INFO ( " sampling using %s method " , sampling_methods_str [ sample_method ] ) ;
struct ggml_tensor * x_0 = sd - > sample ( work_ctx , init_latent , c , uc , cfg_scale , sample_method , sigma_sched ) ;
struct ggml_tensor * x_0 = sd - > sample ( work_ctx , init_latent , noise , c , uc , cfg_scale , sample_method , sigma_sched ) ;
// struct ggml_tensor *x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
// print_ggml_tensor(x_0);
int64_t t3 = ggml_time_ms ( ) ;
@@ -5517,7 +5443,7 @@ std::vector<uint8_t*> StableDiffusion::img2img(const uint8_t* init_img_data,
sd - > diffusion_model . destroy ( ) ;
}
struct ggml_tensor * img = sd - > comput e_first_stage( work_ctx , x_0 , true );
struct ggml_tensor * img = sd - > decod e_first_stage( work_ctx , x_0 ) ;
if ( img ! = NULL ) {
result . push_back ( sd_tensor_to_image ( img ) ) ;
}