Update app.py

app.py

@@ -9,9 +9,6 @@ import spaces
 from huggingface_hub import login
 from gradio_imageslider import ImageSlider  # Import ImageSlider
 
-# Login to Hugging Face
-login(token=os.getenv("HF_TOKEN"))
-
 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
@@ -24,7 +21,7 @@ if not os.path.exists(model_path):
     with open(model_path, 'wb') as f:
         f.write(response.content)
 
-# https://github.com/XLabs-AI/x-flux.git
+# Source: https://github.com/XLabs-AI/x-flux.git
 name = "flux-dev"
 device = torch.device("cuda")
 offload = False
@@ -45,9 +42,29 @@ def load_models():
 
 load_models()
 
-def preprocess_canny_image(image, width=1024, height=1024):
-    image = c_crop(image)
-    image = image.resize((width, height))
+def preprocess_image(image, target_width, target_height, crop=True):
+    if crop:
+        image = c_crop(image)  # Crop the image to square
+        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))
+    else:
+        image = image.resize((target_width, target_height), Image.LANCZOS)
+    
+    return image
+
+def preprocess_canny_image(image, target_width, target_height, crop=True):
+    image = preprocess_image(image, target_width, target_height, crop=crop)
     image = canny_processor(image)
     return image
 
@@ -55,7 +72,7 @@ def preprocess_canny_image(image, width=1024, height=1024):
 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/")
 
@@ -71,6 +88,7 @@ def generate_image(prompt, control_image, num_steps=50, guidance=4, width=512, h
     height = 16 * height // 16
     timesteps = get_schedule(num_steps, (width // 8) * (height // 8) // (16 * 16), shift=(not is_schnell))
     
+    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)
@@ -89,7 +107,7 @@ def generate_image(prompt, control_image, num_steps=50, guidance=4, width=512, h
     x1 = rearrange(x1[-1], "c h w -> h w c")
     output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
     
-    return [control_image, output_img]  # Return both images for slider
+    return [processed_input, output_img]  # Return both images for slider
 
 interface = gr.Interface(
     fn=generate_image,
@@ -104,7 +122,7 @@ interface = gr.Interface(
         gr.Checkbox(label="Random Seed")
     ],
     outputs=ImageSlider(label="Before / After"),  # Use ImageSlider as the output
-    title="FLUX.1 Controlnet Cany",
+    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)]"
 )