Update app.py

app.py

@@ -5,15 +5,13 @@ import torch
 import random
 from PIL import Image
 #from kontext_pipeline import FluxKontextPipeline
-from pipeline_flux_kontext import FluxKontextPipeline
-from diffusers import FluxTransformer2DModel
+from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download
 
-kontext_path = hf_hub_download(repo_id="diffusers/kontext-v2", filename="dev-opt-2-a-3.safetensors")
+# Load Kontext model
 MAX_SEED = np.iinfo(np.int32).max
-transformer = FluxTransformer2DModel.from_single_file(kontext_path, torch_dtype=torch.bfloat16)
-pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
+
+pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
 def concatenate_images(images, direction="horizontal"):
     """