Update app.py

app.py

@@ -3,8 +3,10 @@ import numpy as np
 import random
 import spaces
 import torch
-from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from diffusers import DiffusionPipeline
+from PIL import Image, ImageEnhance
+import io
+import zipfile
 
 # Initialize model and settings
 dtype = torch.bfloat16
@@ -15,36 +17,85 @@ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_d
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+# Helper function to apply image filters
+def apply_filters(image, brightness, contrast, saturation):
+    enhancer = ImageEnhance.Brightness(image)
+    image = enhancer.enhance(brightness)
+    
+    enhancer = ImageEnhance.Contrast(image)
+    image = enhancer.enhance(contrast)
+    
+    enhancer = ImageEnhance.Color(image)
+    image = enhancer.enhance(saturation)
+    
+    return image
+
+# Helper function to create a ZIP file of images
+def create_zip(images):
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED) as zip_file:
+        for i, img in enumerate(images):
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format="PNG")
+            zip_file.writestr(f"image_{i+1}.png", img_byte_arr.getvalue())
+    return zip_buffer.getvalue()
+
 @spaces.GPU(duration=190)
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, output_format="png", progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, num_outputs=1, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, output_format="png", brightness=1.0, contrast=1.0, saturation=1.0, style="None", progress=gr.Progress(track_tqdm=True)):
+    images = []
+    seeds = []
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    image = pipe(
-        prompt=prompt, 
-        width=width, 
-        height=height, 
-        num_inference_steps=num_inference_steps, 
-        generator=generator, 
-        guidance_scale=guidance_scale
-    ).images[0]
-
-    # Convert image to desired format
+    
+    for _ in range(num_outputs):
+        generator = torch.Generator(device=device).manual_seed(seed)
+        image = pipe(
+            prompt=prompt, 
+            width=width, 
+            height=height, 
+            num_inference_steps=num_inference_steps, 
+            generator=generator, 
+            guidance_scale=guidance_scale
+        ).images[0]
+
+        # Apply filters
+        image = apply_filters(image, brightness, contrast, saturation)
+
+        # Apply preset style if selected
+        if style != "None":
+            # Example of applying a filter based on a selected style
+            # You would expand this to include actual styles, e.g., cartoon, watercolor, etc.
+            if style == "Black & White":
+                image = image.convert("L").convert("RGB")
+            elif style == "Sepia":
+                sepia_filter = np.array(image)
+                sepia_filter = np.dot(sepia_filter[...,:3], [0.393, 0.769, 0.189])
+                sepia_filter = np.clip(sepia_filter, 0, 255)
+                image = Image.fromarray(sepia_filter.astype('uint8'))
+        
+        images.append(image)
+        seeds.append(seed)
+        seed += 1  # Increment seed for the next image
+
+    # Optionally convert image format
     if output_format.lower() != "png":
-        image = image.convert(output_format.upper())
+        images = [img.convert(output_format.upper()) for img in images]
     
-    return image, seed
+    return images, seeds, create_zip(images)
 
+# Example prompts for users to try
 examples = [
-    "a tiny astronaut hatching from an egg on the moon",
-    "a cat holding a sign that says hello world",
-    "an anime illustration of a wiener schnitzel",
+    ["a tiny astronaut hatching from an egg on the moon", 1],
+    ["a cat holding a sign that says hello world", 2],
+    ["an anime illustration of a wiener schnitzel", 3],
 ]
 
+# CSS styling for modern look
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 800px;
+    max-width: 1000px;
     padding: 20px;
     background-color: #f9f9f9;
     border-radius: 10px;
@@ -75,9 +126,9 @@ with gr.Blocks(css=css) as demo:
         
         # Title and Description
         gr.Markdown(f"""
-        <h1 id="title">FLUX.1 [dev] - Advanced Text-to-Image Generator</h1>
+        <h1 id="title">FLUX.1 [dev] - Feature-Rich Text-to-Image Generator</h1>
         <p style="text-align:center; color:#555;">
-        Experience the power of a 12B param rectified flow transformer. Customize your prompts and settings to generate unique images every time.
+        Unleash your creativity with this advanced text-to-image generator. Customize prompts, generate multiple images, apply filters, and more!
         </p>
         """)
         
@@ -90,12 +141,21 @@ with gr.Blocks(css=css) as demo:
                 container=False,
                 interactive=True
             )
-        
+
+        num_outputs = gr.Slider(
+            label="Number of Variations",
+            minimum=1,
+            maximum=10,
+            step=1,
+            value=1,
+            interactive=True
+        )
+
         run_button = gr.Button("Generate", elem_id="generate-button")
 
-        # Output image and settings
+        # Output image gallery and settings
         with gr.Row(elem_id="output-container"):
-            result = gr.Image(label="Generated Image", show_label=False).style(height=400)
+            gallery = gr.Gallery(label="Generated Images", show_label=False).style(grid=[4], height="auto")
             output_format = gr.Radio(
                 label="Output Format",
                 choices=["png", "jpeg", "bmp"],
@@ -151,21 +211,64 @@ with gr.Blocks(css=css) as demo:
                     interactive=True
                 )
 
+            # Image filter sliders
+            brightness = gr.Slider(
+                label="Brightness",
+                minimum=0.5,
+                maximum=2.0,
+                step=0.1,
+                value=1.0,
+                interactive=True
+            )
+            contrast = gr.Slider(
+                label="Contrast",
+                minimum=0.5,
+                maximum=2.0,
+                step=0.1,
+                value=1.0,
+                interactive=True
+            )
+            saturation = gr.Slider(
+                label="Saturation",
+                minimum=0.5,
+                maximum=2.0,
+                step=0.1,
+                value=1.0,
+                interactive=True
+            )
+
+            # Preset styles
+            style = gr.Dropdown(
+                label="Preset Styles",
+                choices=["None", "Black & White", "Sepia", "Vivid Colors"],
+                value="None",
+                interactive=True
+            )
+
         # Interactive Examples
         gr.Examples(
             examples=examples,
             fn=infer,
-            inputs=[prompt],
-            outputs=[result, seed],
+            inputs=[prompt, num_outputs],
+            outputs=[gallery, seed],
             cache_examples="lazy",
             label="Try these examples:"
         )
 
+        # Download all images as a ZIP file
+        download_button = gr.Button("Download All Images")
+        zip_file = gr.File(label="Download", show_label=False)
+        download_button.click(
+            fn=create_zip,
+            inputs=[gallery],
+            outputs=[zip_file]
+        )
+
         # Link button to trigger inference
         run_button.click(
             fn=infer,
-            inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, output_format],
-            outputs=[result, seed]
+            inputs=[prompt, num_outputs, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, output_format, brightness, contrast, saturation, style],
+            outputs=[gallery, seed, zip_file]
         )
 
 demo.launch()