Skip to content

Commit

Permalink
refactor: sd arch
Browse files Browse the repository at this point in the history
Signed-off-by: thxCode <[email protected]>
  • Loading branch information
thxCode committed Dec 28, 2024
1 parent e63e4a6 commit 3c4a412
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 11 deletions.
8 changes: 6 additions & 2 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ on:
- "**.mdx"
- "**.png"
- "**.jpg"
- "!.github/workflows/ci.yml"
- ".github/workflows/cmd.yml"
- ".github/workflows/prune.yml"
- ".github/workflows/sync.yml"
pull_request:
branches:
- 'main'
Expand All @@ -30,7 +32,9 @@ on:
- "**.mdx"
- "**.png"
- "**.jpg"
- "!.github/workflows/ci.yml"
- ".github/workflows/cmd.yml"
- ".github/workflows/prune.yml"
- ".github/workflows/sync.yml"

jobs:
ci:
Expand Down
4 changes: 3 additions & 1 deletion .github/workflows/cmd.yml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ on:
- "**.mdx"
- "**.png"
- "**.jpg"
- "!.github/workflows/ci.yml"
- ".github/workflows/ci.yml"
- ".github/workflows/prune.yml"
- ".github/workflows/sync.yml"
tags:
- "v*.*.*"

Expand Down
32 changes: 26 additions & 6 deletions file_architecture.go
Original file line number Diff line number Diff line change
Expand Up @@ -244,12 +244,16 @@ func (gf *GGUFFile) diffuserArchitecture() (ga GGUFArchitecture) {
const (
// Diffusion

sdKey = "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_v.weight" // SD 1.x/2.x
sdXlKey = "model.diffusion_model.output_blocks.5.1.transformer_blocks.1.attn1.to_v.weight" // SD XL
sdXlRefinerKey = "model.diffusion_model.output_blocks.8.1.transformer_blocks.1.attn1.to_v.weight" // SD XL Refiner
sd3Key = "model.diffusion_model.joint_blocks.23.x_block.attn.proj.weight" // SD 3.x
fluxKey = "model.diffusion_model.double_blocks.0.txt_attn.proj.weight" // FLUX.1
fluxKey2 = "double_blocks.0.txt_attn.proj.weight"
sdKey = "model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_v.weight" // SD 1.x/2.x
sdXlKey = "model.diffusion_model.output_blocks.5.1.transformer_blocks.1.attn1.to_v.weight" // SD XL
sdXlRefinerKey = "model.diffusion_model.output_blocks.8.1.transformer_blocks.1.attn1.to_v.weight" // SD XL Refiner
sd3Key = "model.diffusion_model.joint_blocks.23.x_block.attn.proj.weight" // SD 3.x
sdInPaintFeatureKey = "model.diffusion_model.input_blocks.0.0.weight" // SD in-paint feature

fluxKey = "model.diffusion_model.double_blocks.0.txt_attn.proj.weight" // FLUX.1
fluxKey2 = "double_blocks.0.txt_attn.proj.weight"
fluxFillFeatureKey = "model.diffusion_model.img_in.weight" // FLUX.1 Fill feature
fluxFillFeatureKey2 = "img_in.weight"

// Conditioner

Expand All @@ -265,8 +269,12 @@ func (gf *GGUFFile) diffuserArchitecture() (ga GGUFArchitecture) {
sdXlKey,
sdXlRefinerKey,
sd3Key,
sdInPaintFeatureKey,

fluxKey,
fluxKey2,
fluxFillFeatureKey,
fluxFillFeatureKey2,

openAiClipVitL14Key,
openClipVitH14Key,
Expand All @@ -283,21 +291,33 @@ func (gf *GGUFFile) diffuserArchitecture() (ga GGUFArchitecture) {
if ti.Dimensions[0] == 1024 {
ga.DiffusionArchitecture = "Stable Diffusion 2.x"
}
if ti, ok := tis[sdInPaintFeatureKey]; ok && ti.Dimensions[2] == 9 {
ga.DiffusionArchitecture += " InPaint"
}
} else if _, ok := tis[sdXlKey]; ok {
ga.DiffusionArchitecture = "Stable Diffusion XL"
if _, ok = tis[sdXlRefinerKey]; ok {
ga.DiffusionArchitecture = "Stable Diffusion XL Refiner"
}
if ti, ok := tis[sdInPaintFeatureKey]; ok && ti.Dimensions[2] == 9 {
ga.DiffusionArchitecture += " InPaint"
}
} else if _, ok := tis[sd3Key]; ok {
ga.DiffusionArchitecture = "Stable Diffusion 3.x"
ga.DiffusionTransformer = true
}
if _, ok := tis[fluxKey]; ok {
ga.DiffusionArchitecture = "FLUX.1"
ga.DiffusionTransformer = true
if ti, ok := tis[fluxFillFeatureKey]; ok && ti.Dimensions[0] == 384 {
ga.DiffusionArchitecture += " Fill"
}
} else if _, ok := tis[fluxKey2]; ok {
ga.DiffusionArchitecture = "FLUX.1"
ga.DiffusionTransformer = true
if ti, ok := tis[fluxFillFeatureKey2]; ok && ti.Dimensions[0] == 384 {
ga.DiffusionArchitecture += " Fill"
}
}

if ti, ok := tis[openAiClipVitL14Key]; ok {
Expand Down
4 changes: 2 additions & 2 deletions file_estimate__stablediffusioncpp.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ type (
// true for support.
NoMMap bool `json:"noMMap"`
// ImageOnly is the flag to indicate whether the model is used for generating image,
// true for embedding only.
// true for generating image only.
ImageOnly bool `json:"imageOnly"`
// Distributable is the flag to indicate whether the model is distributable,
// true for distributable.
Expand Down Expand Up @@ -240,7 +240,7 @@ func (gf *GGUFFile) EstimateStableDiffusionCppRun(opts ...GGUFRunEstimateOption)
//
{
usage := uint64(50 * 1024 * 1024)
usage += uint64(*o.SDCWidth) * uint64(*o.SDCHeight) * 3 /* output channels */ * 4 /* sizeof(float) */ * 2 /* include img2img*/
usage += uint64(*o.SDCWidth) * uint64(*o.SDCHeight) * 3 /* output channels */ * 4 /* sizeof(float) */ * 3 /* include img2img*/
e.Devices[0].Computation += GGUFBytesScalar(usage * uint64(ptr.Deref(o.ParallelSize, 1)) /* max batch */)
}

Expand Down

0 comments on commit 3c4a412

Please sign in to comment.