align image resize with VAE sample size

app.py

@@ -16,12 +16,12 @@ MAX_IMAGE_SIZE = 2048
 # Load the diffusion pipeline
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
-def preprocess_image(image):
+def preprocess_image(image, image_size):
     # Preprocess the image for the VAE
     preprocess = transforms.Compose([
-        transforms.Resize((512, 512)),  # Adjust the size as needed
+        transforms.Resize((image_size, image_size)),  # Use model-specific size
         transforms.ToTensor(),
-        transforms.Normalize([0.5], [0.5])
+        transforms.Normalize([0.5], [0.5])  # Ensure this matches the VAE's training normalization
     ])
     image = preprocess(image).unsqueeze(0).to(device, dtype=dtype)
     return image
@@ -37,17 +37,17 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
+
+    # Get the expected image size for the VAE
+    vae_image_size = pipe.vae.config.sample_size
     
     if init_image is not None:
         # Process img2img
         init_image = init_image.convert("RGB")
-        init_image = preprocess_image(init_image)
+        init_image = preprocess_image(init_image, vae_image_size)
         latents = encode_image(init_image, pipe.vae)
         # Ensure latents are correctly shaped and adjusted
         latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
-        latents = latents * 0.18215  # Adjust latent scaling factor if necessary
-        
-        # Ensure latents are reshaped to match the expected input dimensions of the model
         latents = latents.view(1, -1, height // 8, width // 8)
         
         image = pipe(
@@ -72,6 +72,8 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     
     return image, seed
 
+
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",