[Dreambooth] flax fixes (huggingface#1765)

* Fail if there are less images than the effective batch size.

* Remove lr-scheduler arg as it's currently ignored.

* Make guidance_scale work for batch_size > 1.

Author: pcuenca

pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

@@ -337,7 +337,7 @@ def __call__(
             guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0])
             if len(prompt_ids.shape) > 2:
                 # Assume sharded
-                guidance_scale = guidance_scale.reshape(prompt_ids.shape[:2])
+                guidance_scale = guidance_scale[:, None]
 
         if jit:
             images = _p_generate(