diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 2233e073..2b504e21 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -708,7 +708,11 @@ class SDTrainer(BaseSDTrainProcess): unconditional_target = unconditional_target * alpha target = unconditional_target + guidance_scale * (target - unconditional_target) - + + if self.train_config.do_differential_guidance: + with torch.no_grad(): + guidance_scale = self.train_config.differential_guidance_scale + target = noise_pred + guidance_scale * (target - noise_pred) if target is None: target = noise diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 22a78189..96617620 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -545,7 +545,10 @@ class TrainConfig: self.unconditional_prompt: str = kwargs.get('unconditional_prompt', '') if isinstance(self.guidance_loss_target, tuple): self.guidance_loss_target = list(self.guidance_loss_target) - + + self.do_differential_guidance = kwargs.get('do_differential_guidance', False) + self.differential_guidance_scale = kwargs.get('differential_guidance_scale', 3.0) + # for multi stage models, how often to switch the boundary self.switch_boundary_every: int = kwargs.get('switch_boundary_every', 1) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 11bae8b1..ab9a57f5 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -2907,7 +2907,7 @@ class StableDiffusion: try: te_has_grad = encoder.text_model.final_layer_norm.weight.requires_grad except: - te_has_grad = encoder.encoder.block[0].layer[0].SelfAttention.q.weight.requires_grad + te_has_grad = False self.device_state['text_encoder'].append({ 'training': encoder.training, 'device': encoder.device, diff --git a/ui/public/imgs/diff_guidance.png b/ui/public/imgs/diff_guidance.png new file mode 100644 index 00000000..9697c978 Binary files /dev/null and b/ui/public/imgs/diff_guidance.png differ diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index fa9d532a..093481ed 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -669,6 +669,45 @@ export default function SimpleJob({ +
+ +
+
+ { + let newValue = value == false ? undefined : value; + setJobConfig(newValue, 'config.process[0].train.do_differential_guidance'); + if (!newValue) { + setJobConfig(undefined, 'config.process[0].train.differential_guidance_scale'); + } else if ( + jobConfig.config.process[0].train.differential_guidance_scale === undefined || + jobConfig.config.process[0].train.differential_guidance_scale === null + ) { + // set default differential guidance scale to 3.0 + setJobConfig(3.0, 'config.process[0].train.differential_guidance_scale'); + } + }} + /> + {jobConfig.config.process[0].train.differential_guidance_scale && ( + <> + setJobConfig(value, 'config.process[0].train.differential_guidance_scale')} + placeholder="eg. 3.0" + min={0} + /> + + )} +
+
+ +
<> diff --git a/ui/src/components/Card.tsx b/ui/src/components/Card.tsx index 13c7409b..14875181 100644 --- a/ui/src/components/Card.tsx +++ b/ui/src/components/Card.tsx @@ -1,13 +1,41 @@ +import { Disclosure, DisclosureButton, DisclosurePanel } from '@headlessui/react'; +import { FaChevronDown } from 'react-icons/fa'; +import classNames from 'classnames'; + interface CardProps { title?: string; children?: React.ReactNode; + collapsible?: boolean; + defaultOpen?: boolean; } -const Card: React.FC = ({ title, children }) => { +const Card: React.FC = ({ title, children, collapsible, defaultOpen }) => { + if (collapsible) { + return ( + + {({ open }) => ( + <> + +
+ {title && ( +

+ {title} +

+ )} +
+ + + {children ?? null} + {open &&
} + + )} + + ); + } return (
{title &&

{title}

} - {children ? children : null} + {children ?? null}
); }; diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx index c7545793..291d9ee3 100644 --- a/ui/src/docs.tsx +++ b/ui/src/docs.tsx @@ -258,6 +258,25 @@ const docs: { [key: string]: ConfigDoc } = { ), }, + 'train.do_differential_guidance': { + title: 'Differential Guidance', + description: ( + <> + Differential Guidance will amplify the difference of the model prediction and the target during training to make + a new target. Differential Guidance Scale will be the multiplier for the difference. This is still experimental, + but in my tests, it makes the model train faster, and learns details better in every scenario I have tried with + it. +
+
+ The idea is that normal training inches closer to the target but never actually gets there, because it is + limited by the learning rate. With differential guidance, we amplify the difference for a new target beyond the + actual target, this would make the model learn to hit or overshoot the target instead of falling short. +
+
+ Differential Guidance Diagram + + ), + }, }; export const getDoc = (key: string | null | undefined): ConfigDoc | null => { diff --git a/ui/src/types.ts b/ui/src/types.ts index d01bd89f..80f0b782 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -139,6 +139,8 @@ export interface TrainConfig { blank_prompt_preservation_multiplier?: number; switch_boundary_every: number; loss_type: 'mse' | 'mae' | 'wavelet' | 'stepped'; + do_differential_guidance?: boolean; + differential_guidance_scale?: number; } export interface QuantizeKwargsConfig { diff --git a/version.py b/version.py index 48879f5b..9b8720f5 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.7.3" \ No newline at end of file +VERSION = "0.7.4" \ No newline at end of file