Stable diffusion: Make guidance modifiable (#2077)

examples/stable_diffusion.py

@@ -568,6 +568,7 @@ if __name__ == "__main__":
   parser.add_argument('--fp16', action='store_true', help="Cast the weights to float16")
   parser.add_argument('--timing', action='store_true', help="Print timing per step")
   parser.add_argument('--seed', type=int, help="Set the random latent seed")
+  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
   args = parser.parse_args()
 
   Tensor.no_grad = True
@@ -594,12 +595,11 @@ if __name__ == "__main__":
   # done with clip model
   del model.cond_stage_model
 
-  def get_model_output(latent, timestep):
+  def get_model_output(latent, timestep, unconditional_guidance_scale):
     # put into diffuser
     latents = model.model.diffusion_model(latent.expand(2, *latent.shape[1:]), timestep, unconditional_context.cat(context, dim=0))
     unconditional_latent, latent = latents[0:1], latents[1:2]
 
-    unconditional_guidance_scale = 7.5
     e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
     return e_t
 
@@ -624,8 +624,8 @@ if __name__ == "__main__":
     return x_prev, pred_x0
 
   @TinyJit
-  def do_step(latent, timestep, index):
-    e_t = get_model_output(latent, timestep)
+  def do_step(latent, timestep, index, guidance):
+    e_t = get_model_output(latent, timestep, guidance)
     x_prev, _ = get_x_prev_and_pred_x0(latent, e_t, index)
     #e_t_next = get_model_output(x_prev)
     #e_t_prime = (e_t + e_t_next) / 2
@@ -641,7 +641,7 @@ if __name__ == "__main__":
     GlobalCounters.reset()
     t.set_description("%3d %3d" % (index, timestep))
     with Timing("step in ", enabled=args.timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
-      latent = do_step(latent, Tensor([timestep]), Tensor([index]))
+      latent = do_step(latent, Tensor([timestep]), Tensor([index]), Tensor([args.guidance]))
       if args.timing: Device[Device.DEFAULT].synchronize()
   del do_step