Merge pull request #37 from storytold/feature/image_generation

Feature/image generation

image_generation/Storyteller-Studio-Filter/main.py

@@ -131,14 +131,16 @@ def main():
     parser.add_argument("--prompt", type=str, required=True, help="Text input called prompt")
     parser.add_argument("--negative-prompt", type=str, required=True, help="Text input called negative prompt")
     parser.add_argument("--number-of-samples",type=int, required=True)
-    parser.add_argument("--sampler", type=str, required=False, help="Text input called samplers", choices=_available_samplers)
+    #parser.add_argument("--sampler", type=str, required=False, help="Text input called samplers", choices=_available_samplers)
+    parser.add_argument("--samplers", type=str, required=False, help="Text input called samplers", choices=list(_au111_to_diffusers_samplers.keys()))
+       
     # ZDisket: We shoud really keep these static and instead use an upsampler afterwards instead
     parser.add_argument("--width", type=InputCheckers.check_power_of_2, required=False, help="Width value (power of 2, between 256 and 4096)")
     parser.add_argument("--height", type=InputCheckers.check_power_of_2, required=False, help="Height value (power of 2, between 256 and 4096)")
-
+    
     parser.add_argument("--cfg-scale", type=InputCheckers.check_cfg_scale, required=True, help="Integer input called CFG scale (range from 1-30)")
     parser.add_argument("--seed", type=int, required=False, help="Seed value (integer of any value)")
-    parser.add_argument("--loRA-path", type=InputCheckers.valid_path, required=True, help="Path object to the loRA")
+    parser.add_argument("--loRA-path", type=InputCheckers.valid_path, required=False, help="Path object to the loRA")
     parser.add_argument("--check-point", type=InputCheckers.valid_path, required=True, help="Path object to the Check Point")
     parser.add_argument("--vae", type=InputCheckers.valid_path, required=False, help="Path to the Variational auto encoder")
     parser.add_argument("--openpose", type=InputCheckers.valid_path, required=False,
@@ -163,7 +165,7 @@ def main():
     print(f"prompt: {args.prompt}")
     print(f"negative_prompt: {args.negative_prompt}")
     print(f"number of samples: {args.number_of_samples}")
-    print(f"samplers: {args.sampler}")
+    print(f"samplers: {args.samplers}")
     print(f"width: {args.width}")
     print(f"height: {args.height}")
     print(f"cfg_scale: {args.cfg_scale}")
@@ -189,13 +191,14 @@ def main():
     sd_fe.cfg_scale = args.cfg_scale
 
     # Load LoRA
-    lora_path = args.loRA_path
-    sd_fe.load_lora_weights(lora_path=lora_path)
+    if args.loRA_path is not None:
+        lora_path = args.loRA_path
+        sd_fe.load_lora_weights(lora_path=lora_path)
 
     # Replace sampler if specified
-    if args.sampler is not None:
-        sd_fe.replace_scheduler(args.sampler)
-
+    if args.samplers is not None:
+        sd_fe.replace_scheduler(_au111_to_diffusers_samplers[args.samplers])
+        
     if args.seed is not None:
         sd_fe.manual_seed(args.seed)
 
@@ -221,8 +224,6 @@ def main():
         if openpose_auxnet_name is not None:
             sd_fe.make_openpose(openpose_auxnet_name)
 
-
-
     # Set width and height to dimensions. If any or both are None, it will result in the pipeline doing its default.
     dimensions = (args.width, args.height)
 

image_generation/Storyteller-Studio-Filter/requirements.txt

@@ -18,13 +18,23 @@ ruff==0.1.3
 safetensors==0.4.0
 sympy==1.12
 tokenizers==0.13.3
-torch==2.1.0+cu118
-torchaudio==2.1.0+cu118
-torchvision==0.16.0+cu118
 tqdm==4.66.1
 transformers==4.30.2
 triton==2.1.0
 typing_extensions==4.4.0
 urllib3==1.26.13
 xformers==0.0.22.post7+cu118
-diffusers
+diffusers==0.23.1
+accelerate==0.26.1
+omegaconf==2.3.0
+opencv-python==4.9.0.80
+psutil==5.9.7
+torchsde==0.2.6
+trampoline==0.1.2
+imageio==2.33.1
+importlib-metadata==7.0.1
+zipp==3.17.0
+--extra-index-url https://download.pytorch.org/whl/cu118 --trusted-host download.pytorch.org
+torch==2.1.0+cu118
+torchaudio==2.1.0+cu118
+torchvision==0.16.0+cu118

image_generation/Storyteller-Studio-Filter/stablediffusionfe.py

@@ -155,7 +155,8 @@ class StableDiffusionFE:
             raise ValueError("Model must be loaded before loading LoRA weights.")
 
         self.lora_path = lora_path
-        self.pipeline.load_lora_weights(lora_path)
+    
+        self.pipeline.load_lora_weights(lora_path,low_cpu_mem_usage=False, ignore_mismatched_sizes=True)
 
     def load_vae(self, vae_path: str) -> None:
         """

image_generation/Storyteller-Studio-Filter/test.sh

@@ -1,14 +1,15 @@
 python3.10 main.py \
---prompt "This is a prompt" \
+--prompt "raiden mei, pantyhose, boots, single glove, earrings, dress, choker, tank top" \
 --negative-prompt "This is a negative prompt" \
---number-of-samples 20 \
---samplers "DPM++ 2M Karras" \
---width 512 \
---height 256 \
---cfg-scale 10 \
---seed 1234 \
---loRA-path "/path/to/loRA" \
---check-point "/path/to/checkpoint" \
---vae "/path/to/vae" \
+--number-of-samples 128 \
+--samplers "DPM++ 2M SDE Heun Karras" \
+--width 1024 \
+--height 1024 \
+--cfg-scale 8 \
+--seed 1 \
+--loRA-path "xs4vow.safetensors" \
+--check-point "waifuDiffusionBeta03_beta3.safetensors" \
+--vae "vaeFtEma560000Ema_original.pt" \
 --batch-size 1 \
---batch-count 1
+--batch-count 10
+