Update app.py

app.py

@@ -1,87 +1,91 @@
+import os
 import torch
-from diffusers.utils import load_image
-from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel
 import gradio as gr
-import spaces
-
-# Ensure that you're using the appropriate data type for your GPU
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+import numpy as np
+from PIL import Image
+from einops import rearrange
+import requests
+from diffusers.utils import load_image
+from diffusers import FluxControlNetPipeline, FluxControlNetModel
+from gradio_imageslider import ImageSlider  # Import ImageSlider
 
+# Model and pipeline initialization
 base_model = 'black-forest-labs/FLUX.1-dev'
-controlnet_model_union = 'InstantX/FLUX.1-dev-Controlnet-Union'
-
-controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch_dtype)
-controlnet = FluxMultiControlNetModel([controlnet_union])
-
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch_dtype)
+controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
-# If you encounter issues with CUDA, you can run this on the CPU for debugging
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
 
-control_modes = [
-    "canny",
-    "tile",
-    "depth",
-    "blur",
-    "pose",
-    "gray",
-    "lq"
-]
+# Define control modes
+control_modes = {
+    "canny": 0,
+    "tile": 1,
+    "depth": 2,
+    "blur": 3,
+    "pose": 4,
+    "gray": 5,
+    "lq": 6
+}
 
-def adjust_dimensions(width, height):
-    adjusted_width = width - (width % 8)
-    adjusted_height = height - (height % 8)
-    return adjusted_width, adjusted_height
-
-@spaces.GPU
-def generate_image(prompt, control_image_depth, control_mode_depth_index, use_depth, control_image_canny, control_mode_canny_index):
-    control_images = []
-    control_modes = []
-    conditioning_scales = []
-
-    if use_depth:
-        control_images.append(control_image_depth)
-        control_modes.append(control_mode_depth_index)
-        conditioning_scales.append(0.2)
+# Preprocess image
+def preprocess_image(image, target_width, target_height, crop=True):
+    if crop:
+        original_width, original_height = image.size
+        scale = max(target_width / original_width, target_height / original_height)
+        resized_width = int(scale * original_width)
+        resized_height = int(scale * original_height)
+        image = image.resize((resized_width, resized_height), Image.LANCZOS)
+        left = (resized_width - target_width) // 2
+        top = (resized_height - target_height) // 2
+        image = image.crop((left, top, left + target_width, top + target_height))
+    else:
+        image = image.resize((target_width, target_height), Image.LANCZOS)
+    
+    return image
 
-    control_images.append(control_image_canny)
-    control_modes.append(control_mode_canny_index)
-    conditioning_scales.append(0.4)
+@gr.Interface
+def generate_image(prompt, control_image, control_mode, controlnet_conditioning_scale, num_steps=50, guidance=4, width=512, height=512, seed=42, random_seed=False):
+    if random_seed:
+        seed = np.random.randint(0, 10000)
+    
+    torch.manual_seed(seed)
 
-    width, height = control_image_canny.shape[:2]
-    adjusted_width, adjusted_height = adjust_dimensions(width, height)
+    control_mode_idx = control_modes.get(control_mode, 0)
+    
+    processed_input = preprocess_image(control_image, width, height)
 
-    try:
-        image = pipe(
-            prompt,
-            control_image=control_images,
-            control_mode=control_modes,
-            width=adjusted_width,
-            height=adjusted_height,
-            controlnet_conditioning_scale=conditioning_scales,
-            num_inference_steps=24,
-            guidance_scale=3.5,
-            generator=torch.manual_seed(42),
-        ).images[0]
-    except RuntimeError as e:
-        torch.cuda.empty_cache()
-        raise e
+    image = pipe(
+        prompt,
+        control_image=processed_input,
+        control_mode=control_mode_idx,
+        width=width,
+        height=height,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        num_inference_steps=num_steps,
+        guidance_scale=guidance,
+    ).images[0]
 
-    return image
+    return [processed_input, image]  # Return both the control image and the generated image
 
-iface = gr.Interface(
+interface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.Text(label="Prompt"),
-        gr.Image(label="Control Image (Depth)"),
-        gr.Dropdown(choices=control_modes, value=control_modes.index("depth"), label="Control Mode (Depth)"),
-        gr.Checkbox(label="Use Depth Control Image", value=True),
-        gr.Image(label="Control Image (Canny)"),
-        gr.Dropdown(choices=control_modes, value=control_modes.index("canny"), label="Control Mode (Canny)")
+        gr.Textbox(label="Prompt"),
+        gr.Image(type="pil", label="Control Image"),
+        gr.Dropdown(choices=list(control_modes.keys()), label="Control Mode", value="canny"),
+        gr.Slider(minimum=0.1, maximum=10.0, value=0.5, label="ControlNet Conditioning Scale"),
+        gr.Slider(step=1, minimum=1, maximum=64, value=28, label="Num Steps"),
+        gr.Slider(minimum=0.1, maximum=10, value=4, label="Guidance"),
+        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Width"),
+        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Height"),
+        gr.Number(value=42, label="Seed"),
+        gr.Checkbox(label="Random Seed")
     ],
-    outputs=gr.Image(label="Generated Image"),
-    title="FluxControlNet Image Generation",
-    description="Generate an image using FluxControlNet with depth and canny control images.",
+    outputs=ImageSlider(label="Before / After"),  # Use ImageSlider as the output
+    title="FLUX.1 Controlnet with Different Modes",
+    description="Generate images using ControlNet with selectable control modes and a text prompt."
 )
 
-iface.launch(share=True)
+if __name__ == "__main__":
+    interface.launch()