Add LoRA optimization to the SD training example

training/stable_diffusion/README.md

@@ -32,13 +32,13 @@ Make sure to customize the training parameters in the script to suit your specif
 
 ## Inference
 
-For inference, you can use the `inf-loop.py` Python code. Follow these steps:
+For inference, you can use the `inf_txt2img_loop.py` Python code. Follow these steps:
 
 1. Provide your desired prompts as input in the script.
 2. Run the `inf_txt2img_loop.py` script.
 
 Here's an example command to run the inference script:
 
 <pre>
-deepspeed inf_txt2img_loop.py
+deepspeed inf_txt2img_loop.py --out_dir out_images/
 </pre>

training/stable_diffusion/inf_txt2img_loop.py

@@ -2,13 +2,15 @@
 import torch
 import os
 from local_pipeline_stable_diffusion import StableDiffusionPipeline
+from diffusers import DPMSolverMultistepScheduler
 from diffusers import StableDiffusionPipeline as StableDiffusionPipelineBaseline
 import argparse
 
+
 seed = 123450011
 parser = argparse.ArgumentParser()
-parser.add_argument("--ft_model", default="new_sd-distill-v21-10k-1e", type=str, help="Path to the fine-tuned model")
-parser.add_argument("--b_model", default="stabilityai/stable-diffusion-2-1-base", type=str, help="Path to the baseline model")
+parser.add_argument("--finetuned_model", default="./sd-distill-lora-multi-50k-50", type=str, help="Path to the fine-tuned model")
+parser.add_argument("--base_model", default="stabilityai/stable-diffusion-2-1-base", type=str, help="Path to the baseline model")
 parser.add_argument("--out_dir", default="image_out/", type=str, help="Path to the generated images")
 parser.add_argument('--guidance_scale', type=float, default=7.5, help='Guidance Scale')
 parser.add_argument("--use_local_pipe", action='store_true', help="Use local SD pipeline")
@@ -40,17 +42,27 @@
            "A person holding a cat"]
 
 
+# Load the pipelines
+pipe_new = StableDiffusionPipeline.from_pretrained(args.base_model, torch_dtype=torch.float16).to("cuda")
+pipe_baseline = StableDiffusionPipelineBaseline.from_pretrained(args.base_model, torch_dtype=torch.float16).to("cuda")
+
+pipe_new.scheduler = DPMSolverMultistepScheduler.from_config(pipe_new.scheduler.config)
+pipe_baseline.scheduler = DPMSolverMultistepScheduler.from_config(pipe_baseline.scheduler.config)
+
+# Load the Lora weights
+pipe_new.unet.load_attn_procs(args.finetuned_model)
+
+pipe_new = deepspeed.init_inference(pipe_new, mp_size=world_size, dtype=torch.half)
+pipe_baseline = deepspeed.init_inference(pipe_baseline, mp_size=world_size, dtype=torch.half)
+
+# Generate the images
 for prompt in prompts:
-    #--- new image
-    pipe_new = StableDiffusionPipeline.from_pretrained(args.ft_model, torch_dtype=torch.float16).to("cuda")
-    generator = torch.Generator("cuda").manual_seed(seed)
-    pipe_new = deepspeed.init_inference(pipe_new, mp_size=world_size, dtype=torch.half)
-    image_new = pipe_new(prompt, num_inference_steps=50, guidance_scale=args.guidance_scale, generator=generator).images[0]
-    image_new.save(args.out_dir+"/NEW__seed_"+str(seed)+"_"+prompt[0:100]+".png")
-
-    #--- baseline image
-    pipe_baseline = StableDiffusionPipelineBaseline.from_pretrained(args.b_model, torch_dtype=torch.float16).to("cuda")
-    generator = torch.Generator("cuda").manual_seed(seed)                                              
-    pipe_baseline = deepspeed.init_inference(pipe_baseline, mp_size=world_size, dtype=torch.half)
-    image_baseline = pipe_baseline(prompt, num_inference_steps=50, guidance_scale=args.guidance_scale, generator=generator).images[0]
-    image_baseline.save(args.out_dir+"/BASELINE_seed_"+str(seed)+"_"+prompt[0:100]+".png")
+        #--- baseline image
+        generator = torch.Generator("cuda").manual_seed(seed)
+        image_baseline = pipe_baseline(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
+        image_baseline.save(args.out_dir+"BASELINE_seed_"+str(seed)+"_"+prompt[0:100]+".png")
+
+        #--- new image
+        generator = torch.Generator("cuda").manual_seed(seed)
+        image_new = pipe_new(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
+        image_new.save(args.out_dir+"NEW_seed_"+str(seed)+"_"+prompt[0:100]+".png")

training/stable_diffusion/local_pipeline_stable_diffusion.py

@@ -29,11 +29,11 @@
     is_accelerate_available,
     is_accelerate_version,
     logging,
-    randn_tensor,
     replace_example_docstring,
 )
 
-from diffusers.pipeline_utils import DiffusionPipeline
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 

training/stable_diffusion/mytrainbash.sh

@@ -1,21 +1,21 @@
 export MODEL_NAME="stabilityai/stable-diffusion-2-1-base"
-export OUTPUT_DIR="./sd-distill-v21"
+export OUTPUT_DIR="./sd-distill-lora-multi-50k-50"
 
 if [ ! -d "$OUTPUT_DIR" ]; then
     mkdir "$OUTPUT_DIR"
     echo "Folder '$OUTPUT_DIR' created"
 else
     echo "Folder '$OUTPUT_DIR' already exists"
 fi
-
 
-accelerate launch train_sd_distil_lora.py \
+
+accelerate launch train_sd_distill_lora.py \
 	   --pretrained_model_name_or_path=$MODEL_NAME  \
 	   --output_dir=$OUTPUT_DIR \
-	   --default_prompt="A man dancing" \
 	   --resolution=512 \
 	   --train_batch_size=1 \
 	   --gradient_accumulation_steps=1 \
 	   --learning_rate=5e-6 \
 	   --lr_scheduler="constant" \
-	   --lr_warmup_steps=0
+	   --lr_warmup_steps=0 \
+	   --default_prompt="A man dancing"

training/stable_diffusion/requirements.txt

@@ -1,4 +1,5 @@
 accelerate>=0.16.0
+diffusers==0.23.1
 torchvision
 transformers>=4.25.1
 ftfy