Update app.py

app.py

@@ -6,8 +6,7 @@ import random
 import json
 import os
 from PIL import Image
-from kontext_pipeline import FluxKontextPipeline
-from diffusers import FluxTransformer2DModel
+from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
 from safetensors.torch import load_file
@@ -15,11 +14,9 @@ import requests
 import re
 
 # Load Kontext model
-kontext_path = hf_hub_download(repo_id="diffusers/kontext-v2", filename="dev-opt-2-a-3.safetensors")
 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")
 
 # Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)