update DPM++ and juggernaut-lightning

Fanghua-Yu · Fanghua-Yu · commit 98558f9143d0 · 2024-03-08T22:44:40.000+08:00
diff --git a/README.md b/README.md
@@ -39,6 +39,10 @@ For users who can connect to huggingface, please setting `LLAVA_CLIP_PATH, SDXL_
 * [SDXL base 1.0_0.9vae](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors)
 * [LLaVA CLIP](https://huggingface.co/openai/clip-vit-large-patch14-336)
 * [LLaVA v1.5 13B](https://huggingface.co/liuhaotian/llava-v1.5-13b)
+* (optional) [Juggernaut-XL_v9_RunDiffusionPhoto_v2](https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/blob/main/Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors)
+  * Replacement of `SDXL base 1.0_0.9vae` for Photo Realistic
+* (optional) [Juggernaut_RunDiffusionPhoto2_Lightning_4Steps](https://huggingface.co/RunDiffusion/Juggernaut-XL-Lightning/blob/main/Juggernaut_RunDiffusionPhoto2_Lightning_4Steps.safetensors)
+  * Distilling model used in `SUPIR_v0_Juggernautv9_lightning.yaml`
 
 
 #### Models we provided:
@@ -109,6 +113,9 @@ CUDA_VISIBLE_DEVICES=0,1 python test.py --img_dir '/opt/data/private/LV_Dataset/
 ```Shell
 CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history
 
+# Juggernaut_RunDiffusionPhoto2_Lightning_4Steps and DPM++ M2 SDE Karras for fast sampling
+CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history --opt options/SUPIR_v0_Juggernautv9_lightning.yaml
+
 # less VRAM & slower (12G for Diffusion, 16G for LLaVA)
 CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history --loading_half_params --use_tile_vae --load_8bit_llava
 ```
diff --git a/SUPIR/util.py b/SUPIR/util.py
@@ -31,7 +31,7 @@ def create_model(config_path):
     return model
 
 
-def create_SUPIR_model(config_path, SUPIR_sign=None):
+def create_SUPIR_model(config_path, SUPIR_sign=None, load_default_setting=False):
     config = OmegaConf.load(config_path)
     model = instantiate_from_config(config.model).cpu()
     print(f'Loaded model config from [{config_path}]')
@@ -45,6 +45,9 @@ def create_SUPIR_model(config_path, SUPIR_sign=None):
             model.load_state_dict(load_state_dict(config.SUPIR_CKPT_F), strict=False)
         elif SUPIR_sign == 'Q':
             model.load_state_dict(load_state_dict(config.SUPIR_CKPT_Q), strict=False)
+    if load_default_setting:
+        default_setting = config.default_setting
+        return model, default_setting
     return model
 
 def load_QF_ckpt(config_path):
diff --git a/gradio_demo.py b/gradio_demo.py
@@ -15,6 +15,7 @@
 import time
 
 parser = argparse.ArgumentParser()
+parser.add_argument("--opt", type=str, default='options/SUPIR_v0.yaml')
 parser.add_argument("--ip", type=str, default='127.0.0.1')
 parser.add_argument("--port", type=int, default='6688')
 parser.add_argument("--no_llava", action='store_true', default=False)
@@ -40,15 +41,15 @@
     raise ValueError('Currently support CUDA only.')
 
 # load SUPIR
-model = create_SUPIR_model('options/SUPIR_v0.yaml', SUPIR_sign='Q')
+model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
 if args.loading_half_params:
     model = model.half()
 if args.use_tile_vae:
     model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
 model = model.to(SUPIR_device)
 model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
 model.current_model = 'v0-Q'
-ckpt_Q, ckpt_F = load_QF_ckpt('options/SUPIR_v0.yaml')
+ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
 
 # load LLaVA
 if use_llava:
@@ -144,7 +145,7 @@ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale
 
 
 def load_and_reset(param_setting):
-    edm_steps = 50
+    edm_steps = default_setting.edm_steps
     s_stage2 = 1.0
     s_stage1 = -1.0
     s_churn = 5
@@ -160,11 +161,11 @@ def load_and_reset(param_setting):
     linear_s_stage2 = False
     linear_CFG = True
     if param_setting == "Quality":
-        s_cfg = 7.5
-        spt_linear_CFG = 4.0
+        s_cfg = default_setting.s_cfg_Quality
+        spt_linear_CFG = default_setting.spt_linear_CFG_Quality
     elif param_setting == "Fidelity":
-        s_cfg = 4.0
-        spt_linear_CFG = 1.0
+        s_cfg = default_setting.s_cfg_Fidelity
+        spt_linear_CFG = default_setting.spt_linear_CFG_Fidelity
     else:
         raise NotImplementedError
     return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
@@ -230,8 +231,9 @@ def submit_feedback(event_id, fb_score, fb_text):
                 num_samples = gr.Slider(label="Num Samples", minimum=1, maximum=4 if not args.use_image_slider else 1
                                         , value=1, step=1)
                 upscale = gr.Slider(label="Upscale", minimum=1, maximum=8, value=1, step=1)
-                edm_steps = gr.Slider(label="Steps", minimum=20, maximum=200, value=50, step=1)
-                s_cfg = gr.Slider(label="Text Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.1)
+                edm_steps = gr.Slider(label="Steps", minimum=1, maximum=200, value=default_setting.edm_steps, step=1)
+                s_cfg = gr.Slider(label="Text Guidance Scale", minimum=1.0, maximum=15.0,
+                                  value=default_setting.s_cfg_Quality, step=0.1)
                 s_stage2 = gr.Slider(label="Stage2 Guidance Strength", minimum=0., maximum=1., value=1., step=0.05)
                 s_stage1 = gr.Slider(label="Stage1 Guidance Strength", minimum=-1.0, maximum=6.0, value=-1.0, step=1.0)
                 seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
@@ -251,7 +253,7 @@ def submit_feedback(event_id, fb_score, fb_text):
                     with gr.Column():
                         linear_CFG = gr.Checkbox(label="Linear CFG", value=True)
                         spt_linear_CFG = gr.Slider(label="CFG Start", minimum=1.0,
-                                                        maximum=9.0, value=4.0, step=0.5)
+                                                        maximum=9.0, value=default_setting.spt_linear_CFG_Quality, step=0.5)
                     with gr.Column():
                         linear_s_stage2 = gr.Checkbox(label="Linear Stage2 Guidance", value=False)
                         spt_linear_s_stage2 = gr.Slider(label="Guidance Start", minimum=0.,
diff --git a/options/SUPIR_v0.yaml b/options/SUPIR_v0.yaml
@@ -154,3 +154,9 @@ SUPIR_CKPT_F: /opt/data/private/AIGC_pretrain/SUPIR_cache/SUPIR-v0F.ckpt
 SUPIR_CKPT_Q: /opt/data/private/AIGC_pretrain/SUPIR_cache/SUPIR-v0Q.ckpt
 SUPIR_CKPT: ~
 
+default_setting:
+  s_cfg_Quality: 7.5
+  spt_linear_CFG_Quality: 4.0
+  s_cfg_Fidelity: 4.0
+  spt_linear_CFG_Fidelity: 1.0
+  edm_steps: 50
diff --git a/options/SUPIR_v0_Juggernautv9_lightning.yaml b/options/SUPIR_v0_Juggernautv9_lightning.yaml
@@ -0,0 +1,162 @@
+model:
+  target: SUPIR.models.SUPIR_model.SUPIRModel
+  params:
+    ae_dtype: bf16
+    diffusion_dtype: fp16
+    scale_factor: 0.13025
+    disable_first_stage_autocast: True
+    network_wrapper: sgm.modules.diffusionmodules.wrappers.ControlWrapper
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiserWithControl
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+    control_stage_config:
+      target: SUPIR.modules.SUPIR_v0.GLVControl
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4]
+        num_head_channels: 64
+        use_spatial_transformer: True
+        use_linear_in_transformer: True
+        transformer_depth: [1, 2, 10]  # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+#        transformer_depth: [1, 1, 4]
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        legacy: False
+        input_upscale: 1
+
+    network_config:
+      target: SUPIR.modules.SUPIR_v0.LightGLVUNet
+      params:
+        mode: XL-base
+        project_type: ZeroSFT
+        project_channel_scale: 2
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4]
+        num_head_channels: 64
+        use_spatial_transformer: True
+        use_linear_in_transformer: True
+        transformer_depth: [1, 2, 10]  # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        legacy: False
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditionerWithControl
+      params:
+        emb_models:
+          # crossattn cond
+          - is_trainable: False
+            input_key: txt
+            target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+            params:
+              layer: hidden
+              layer_idx: 11
+          # crossattn and vector cond
+          - is_trainable: False
+            input_key: txt
+            target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2
+            params:
+              arch: ViT-bigG-14
+              version: laion2b_s39b_b160k
+              freeze: True
+              layer: penultimate
+              always_return_pooled: True
+              legacy: False
+          # vector cond
+          - is_trainable: False
+            input_key: original_size_as_tuple
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+          # vector cond
+          - is_trainable: False
+            input_key: crop_coords_top_left
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+          # vector cond
+          - is_trainable: False
+            input_key: target_size_as_tuple
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: ~
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1, 2, 4, 4 ]
+          num_res_blocks: 2
+          attn_resolutions: [ ]
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.RestoreDPMPP2MSampler
+      params:
+        num_steps: 100
+        restore_cfg: 4.0
+        s_churn: 0
+        s_noise: 1.003
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.LinearCFG
+          params:
+            scale: 7.5
+            scale_min: 4.0
+
+    p_p:
+        'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, 
+             hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, 
+             skin pore detailing, hyper sharpness, perfect without deformations.'
+    n_p:
+        'painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, CG Style, 3D render, 
+        unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, signature, 
+        jpeg artifacts, deformed, lowres, over-smooth'
+
+SDXL_CKPT: /opt/data/private/AIGC_pretrain/SDXL_lightning_cache/Juggernaut_RunDiffusionPhoto2_Lightning_4Steps.safetensors
+SUPIR_CKPT_F: /opt/data/private/AIGC_pretrain/SUPIR_cache/SUPIR-v0F.ckpt
+SUPIR_CKPT_Q: /opt/data/private/AIGC_pretrain/SUPIR_cache/SUPIR-v0Q.ckpt
+SUPIR_CKPT: ~
+
+default_setting:
+  s_cfg_Quality: 2.0
+  spt_linear_CFG_Quality: 2.0
+  s_cfg_Fidelity: 1.5
+  spt_linear_CFG_Fidelity: 1.5
+  edm_steps: 8
diff --git a/requirements.txt b/requirements.txt
@@ -38,3 +38,5 @@ urllib3==1.26.15
 webdataset==0.2.48
 xformers>=0.0.20
 facexlib==0.3.0
+k-diffusion==0.1.1.post1
+diffusers==0.16.1
diff --git a/sgm/modules/diffusionmodules/sampling.py b/sgm/modules/diffusionmodules/sampling.py
