Update app.py

app.py

@@ -5,19 +5,28 @@ import numpy as np
 from PIL import Image
 from einops import rearrange
 import requests
+import spaces
+from huggingface_hub import login
+from gradio_imageslider import ImageSlider  # Import ImageSlider
 from diffusers.utils import load_image
 from diffusers import FluxControlNetPipeline, FluxControlNetModel
-from gradio_imageslider import ImageSlider  # Import ImageSlider
 
-# Model and pipeline initialization
+# Source: https://github.com/XLabs-AI/x-flux.git
+name = "flux-dev"
+device = torch.device("cuda")
+offload = False
+is_schnell = name == "flux-schnell"
+
 base_model = 'black-forest-labs/FLUX.1-dev'
 controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
+# Load the new ControlNet model and pipeline
 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")
+pipe.to(device)
+
+controlnet_conditioning_scale = 0.5
 
-# Define control modes
 control_modes = {
     "canny": 0,
     "tile": 1,
@@ -25,17 +34,21 @@ control_modes = {
     "blur": 3,
     "pose": 4,
     "gray": 5,
-    "lq": 6
+    "lq": 6,
 }
 
-# Preprocess image
 def preprocess_image(image, target_width, target_height, crop=True):
     if crop:
         original_width, original_height = image.size
+
+        # Resize to match the target size without stretching
         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)
+        
+        # Center crop to match the target dimensions
         left = (resized_width - target_width) // 2
         top = (resized_height - target_height) // 2
         image = image.crop((left, top, left + target_width, top + target_height))
@@ -44,29 +57,32 @@ def preprocess_image(image, target_width, target_height, crop=True):
     
     return image
 
-@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):
+@spaces.GPU(duration=120)
+def generate_image(prompt, control_image, control_mode, 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)
-
-    control_mode_idx = control_modes.get(control_mode, 0)
-    
-    processed_input = preprocess_image(control_image, width, height)
+    if not os.path.isdir("./controlnet_results/"):
+        os.makedirs("./controlnet_results/")
 
-    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]
+    torch_device = torch.device("cuda")
 
-    return [processed_input, image]  # Return both the control image and the generated image
+    control_image = preprocess_image(control_image, width, height)
+    
+    torch.manual_seed(seed)
+    with torch.no_grad():
+        image = pipe(
+            prompt,
+            control_image=control_image,
+            control_mode=control_modes[control_mode],
+            width=width,
+            height=height,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+        ).images[0]
+    
+    return [control_image, image]  # Return both images for slider
 
 interface = gr.Interface(
     fn=generate_image,
@@ -74,7 +90,6 @@ interface = gr.Interface(
         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"),
@@ -83,8 +98,8 @@ interface = gr.Interface(
         gr.Checkbox(label="Random Seed")
     ],
     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."
+    title="FLUX.1 Controlnet Canny",
+    description="Generate images using ControlNet and a text prompt.\n[[non-commercial license, Flux.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)]"
 )
 
 if __name__ == "__main__":