Update app.py

app.py

@@ -4,24 +4,47 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 from einops import rearrange
-from diffusers import FluxControlNetPipeline, FluxControlNetModel
-from diffusers.utils import load_image
+import requests
+import spaces
+from huggingface_hub import login
 from gradio_imageslider import ImageSlider  # Import ImageSlider
-import cv2  # Import OpenCV for Canny edge detection
 
-# Load the new ControlNet model
-base_model = 'black-forest-labs/FLUX.1-dev'
-controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
+from image_datasets.canny_dataset import canny_processor, c_crop
+from src.flux.sampling import denoise_controlnet, get_noise, get_schedule, prepare, unpack
+from src.flux.util import load_ae, load_clip, load_t5, load_flow_model, load_controlnet, load_safetensors
 
+# Download and load the ControlNet model
+model_url = "https://huggingface.co/XLabs-AI/flux-controlnet-canny-v3/resolve/main/flux-canny-controlnet-v3.safetensors?download=true"
+model_path = "./flux-canny-controlnet-v3.safetensors"
+if not os.path.exists(model_path):
+    response = requests.get(model_url)
+    with open(model_path, 'wb') as f:
+        f.write(response.content)
+
+# Source: https://github.com/XLabs-AI/x-flux.git
+name = "flux-dev"
 device = torch.device("cuda")
+offload = False
+is_schnell = name == "flux-schnell"
+
+model, ae, t5, clip, controlnet = None, None, None, None, None
+
+def load_models():
+    global model, ae, t5, clip, controlnet
+    t5 = load_t5(device, max_length=256 if is_schnell else 512)
+    clip = load_clip(device)
+    model = load_flow_model(name, device=device)
+    ae = load_ae(name, device=device)
+    controlnet = load_controlnet(name, device).to(device).to(torch.bfloat16)
 
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
-pipe.to(device)
+    checkpoint = load_safetensors(model_path)
+    controlnet.load_state_dict(checkpoint, strict=False)
+
+load_models()
 
 def preprocess_image(image, target_width, target_height, crop=True):
     if crop:
-        image = image.crop((0, 0, min(image.size), min(image.size)))  # Crop the image to square
+        image = c_crop(image)  # Crop the image to square
         original_width, original_height = image.size
 
         # Resize to match the target size without stretching
@@ -42,50 +65,61 @@ def preprocess_image(image, target_width, target_height, crop=True):
 
 def preprocess_canny_image(image, target_width, target_height, crop=True):
     image = preprocess_image(image, target_width, target_height, crop=crop)
-    image = np.array(image.convert('L'))  # Convert to grayscale for Canny processing
-    image = cv2.Canny(image, 100, 200)  # Apply Canny edge detection
-    image = Image.fromarray(image)
+    image = canny_processor(image)
     return image
 
-def generate_image(prompt, control_image, num_steps=24, guidance=3.5, width=512, height=512, seed=42, random_seed=False, control_mode=0):
+@spaces.GPU(duration=120)
+def generate_image(prompt, control_image, num_steps=50, guidance=4, width=512, height=512, seed=42, random_seed=False):
     if random_seed:
         seed = np.random.randint(0, 10000)
     
     if not os.path.isdir("./controlnet_results/"):
         os.makedirs("./controlnet_results/")
 
-    torch.manual_seed(seed)
+    torch_device = torch.device("cuda")
+    
+    model.to(torch_device)
+    t5.to(torch_device)
+    clip.to(torch_device)
+    ae.to(torch_device)
+    controlnet.to(torch_device)
 
-    control_image = preprocess_canny_image(control_image, width, height)  # Preprocess the control image for Canny mode
+    width = 16 * width // 16
+    height = 16 * height // 16
+    timesteps = get_schedule(num_steps, (width // 8) * (height // 8) // (16 * 16), shift=(not is_schnell))
     
-    controlnet_conditioning_scale = 0.5  # ControlNet conditioning scale
+    processed_input = preprocess_image(control_image, width, height)
+    canny_processed = preprocess_canny_image(control_image, width, height)
+    controlnet_cond = torch.from_numpy((np.array(canny_processed) / 127.5) - 1)
+    controlnet_cond = controlnet_cond.permute(2, 0, 1).unsqueeze(0).to(torch.bfloat16).to(torch_device)
+
+    torch.manual_seed(seed)
+    with torch.no_grad():
+        x = get_noise(1, height, width, device=torch_device, dtype=torch.bfloat16, seed=seed)
+        inp_cond = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
+
+        x = denoise_controlnet(model, **inp_cond, controlnet=controlnet, timesteps=timesteps, guidance=guidance, controlnet_cond=controlnet_cond)
+        
+        x = unpack(x.float(), height, width)
+        x = ae.decode(x)
+
+    x1 = x.clamp(-1, 1)
+    x1 = rearrange(x1[-1], "c h w -> h w c")
+    output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
     
-    # Generate the image using the pipeline
-    image = pipe(
-        prompt, 
-        control_image=control_image,
-        control_mode=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
+    return [processed_input, output_img]  # Return both images for slider
 
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
         gr.Textbox(label="Prompt"),
         gr.Image(type="pil", label="Control Image"),
-        gr.Slider(step=1, minimum=1, maximum=64, value=24, label="Num Steps"),
-        gr.Slider(minimum=0.1, maximum=10, value=3.5, label="Guidance"),
+        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"),
-        gr.Radio(choices=[0, 1, 2, 3, 4, 5, 6], value=0, label="Control Mode")
+        gr.Checkbox(label="Random Seed")
     ],
     outputs=ImageSlider(label="Before / After"),  # Use ImageSlider as the output
     title="FLUX.1 Controlnet Canny",
@@ -94,3 +128,4 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     interface.launch()
+