add num of images

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 from diffusers.utils import load_image
 from pipeline import FluxConditionalPipeline
 from transformer import FluxTransformer2DConditionalModel
-
+from recaption import enhance_prompt
 import os
 
 pipe = None
@@ -44,7 +44,8 @@ def generate_image(
     gemini_prompt: bool = True,
     guidance: float = 3.5,
     i_guidance: float = 1.0,
-    t_guidance: float = 1.0
+    t_guidance: float = 1.0,
+    num_images: int = 4,
 ):
     w, h, min_size = image.size[0], image.size[1], min(image.size)
     image = image.crop(
@@ -52,8 +53,13 @@ def generate_image(
     ).resize((512, 512))
 
     control_image = load_image(image)
+    text_list = []
+    for _ in range(num_images):
+        if gemini_prompt:
+            text = enhance_prompt(image, text.strip())
+        text_list.append(text.strip())
     result_image = pipe(
-        prompt=text.strip(),
+        prompt=text_list,
         negative_prompt="",
         num_inference_steps=28,
         height=512,
@@ -63,7 +69,7 @@ def generate_image(
         guidance_scale_real_i=i_guidance,
         guidance_scale_real_t=t_guidance,
         gemini_prompt=gemini_prompt,
-    ).images[0]
+    ).images
 
     return result_image
 
@@ -125,6 +131,10 @@ with demo:
             <a href="https://huggingface.co/datasets/primecai/dsd_data" target="_blank"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace%20-Data-yellow" style="display:inline-block;"></a>
             <a href="https://huggingface.co/primecai/dsd_model" target="_blank"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face%20-Model-green" style="display:inline-block;"></a>
             <a href="https://x.com/prime_cai?lang=en" target="_blank"><img src="https://img.shields.io/twitter/follow/prime_cai?style=social" style="display:inline-block;"></a>
+        
+            <div style="text-align: center;">
+                The model does have randomness because of both the Gemini prompt enhancement and the diffusion initial noises. Please give it a few tries to get the best results.
+            </div>
         </div>
         """
     )
@@ -132,14 +142,16 @@ with demo:
     iface = gr.Interface(
         fn=generate_image,
         inputs=[
-            gr.Image(type="pil", width=512),
+            gr.Image(type="pil", width=300),
             gr.Textbox(lines=2, label="text", info="Could be something as simple as 'this character playing soccer'."),
             gr.Checkbox(label="Gemini prompt", value=True, info="Use Gemini to enhance the prompt. This is recommended for most cases, unless you have a specific prompt similar to the examples in mind."),
             gr.Slider(minimum=1.0, maximum=6.0, step=0.5, value=3.5, label="guidance scale", info="Tip: start with 3.5, then gradually increase if the consistency is consistently off"),
             gr.Slider(minimum=1.0, maximum=2.0, step=0.05, value=1.5, label="real guidance scale for image", info="Tip: increase if the image is not consistent"),
             gr.Slider(minimum=1.0, maximum=2.0, step=0.05, value=1.0, label="real guidance scale for prompt", info="Tip: increase if the prompt is not consistent"),
+            gr.Slider(minimum=1, maximum=5, step=1, value=4, label="Number of images", info="Select how many images to generate"),
         ],
-        outputs=gr.Image(type="pil"),
+        # outputs=gr.Image(type="pil"),
+        outputs=gr.Gallery(label="Generated Images", height=544),
         # examples=get_samples(),
         live=False,
     )

pipeline.py

@@ -39,7 +39,7 @@ from diffusers.utils import (
 )
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from recaption import enhance_prompt
+# from recaption import enhance_prompt
 
 
 if is_torch_xla_available():
@@ -722,8 +722,8 @@ class FluxConditionalPipeline(DiffusionPipeline, SD3LoraLoaderMixin):
 
         device = self._execution_device
 
-        if gemini_prompt:
-            prompt = enhance_prompt(image, prompt)
+        # if gemini_prompt:
+        #     prompt = enhance_prompt(image, prompt)
         # if gemini_prompt:
         #     while True:
         #         try:
@@ -779,8 +779,8 @@ class FluxConditionalPipeline(DiffusionPipeline, SD3LoraLoaderMixin):
         # 3. Preprocess image
         image = self.image_processor.preprocess(image)
         # image = image[..., :512]
-        image = torch.nn.functional.interpolate(image, size=512)
-        black_image = torch.full((1, 3, 512, 512), -1.0)
+        image = torch.nn.functional.interpolate(image, size=512).repeat(batch_size, 1, 1, 1)
+        black_image = torch.full((batch_size, 3, 512, 512), -1.0)
         image = torch.cat([image, black_image], dim=3)
         latents_cond = self.vae.encode(image.to(self.vae.dtype).to(self.vae.device)).latent_dist.sample()
         latents_cond = (