diff --git a/README.md b/README.md
index 394dcb5c..68020639 100644
--- a/README.md
+++ b/README.md
@@ -425,6 +425,9 @@ Everything else should work the same including layer targeting.
 
 Only larger updates are listed here. There are usually smaller daily updated that are omitted.
 
+### June 29, 2024
+- Fixed issue where Kontext forced sizes on sampling
+
 ### June 26, 2024
 - Added support for FLUX.1 Kontext training
 - added support for instruction dataset training
diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
index 439c9a22..49a79b9b 100644
--- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
+++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
@@ -217,6 +217,13 @@ class FluxKontextModel(BaseModel):
         else:
             control_img = Image.open(gen_config.ctrl_img)
             control_img = control_img.convert("RGB")
+            # resize to width and height
+            if control_img.size != (gen_config.width, gen_config.height):
+                control_img = control_img.resize(
+                    (gen_config.width, gen_config.height), Image.BILINEAR
+                )
+        gen_config.width = int(gen_config.width  // 16 * 16)
+        gen_config.height = int(gen_config.height // 16 * 16)
         img = pipeline(
             image=control_img,
             prompt_embeds=conditional_embeds.text_embeds,
@@ -227,6 +234,8 @@ class FluxKontextModel(BaseModel):
             guidance_scale=gen_config.guidance_scale,
             latents=gen_config.latents,
             generator=generator,
+            max_area=gen_config.height * gen_config.width,
+            _auto_resize=False,
             **extra
         ).images[0]
         return img
diff --git a/version.py b/version.py
index 85b16574..5396bd93 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.3.3"
\ No newline at end of file
+VERSION = "0.3.4"
\ No newline at end of file