Update app.py

app.py

@@ -3,7 +3,7 @@ from PIL import Image
 import numpy as np
 import torch
 import os
-from diffusers import DiffusionPipeline  # Changed: Use DiffusionPipeline instead
+from diffusers import FluxKontextPipeline  # Fixed: Import the correct pipeline
 from diffusers.utils import load_image
 from huggingface_hub import HfApi, login
 from huggingface_hub.utils import HfHubHTTPError
@@ -35,15 +35,23 @@ else:
 # ───── Load FLUX-Kontext Pipeline ─────
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Fixed: Load the custom pipeline correctly
-pipe = DiffusionPipeline.from_pretrained(
-    "HighCWu/FLUX.1-Kontext-dev-bnb-hqq-4bit",
-    torch_dtype=torch.bfloat16,
-    custom_pipeline="FluxKontextPipeline",  # This loads the custom pipeline
-    trust_remote_code=True  # Added: Required for custom pipelines
-)
+# Try the quantized version first, fallback to official if needed
+try:
+    pipe = FluxKontextPipeline.from_pretrained(
+        "HighCWu/FLUX.1-Kontext-dev-bnb-hqq-4bit",
+        torch_dtype=torch.bfloat16,
+    )
+    print("✅ FLUX-Kontext quantized pipeline loaded")
+except Exception as e:
+    print(f"⚠️ Quantized model failed: {e}")
+    print("📥 Falling back to official model...")
+    pipe = FluxKontextPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-Kontext-dev",
+        torch_dtype=torch.bfloat16,
+    )
+    print("✅ FLUX-Kontext official pipeline loaded")
+
 pipe.to(device)
-print("✅ FLUX-Kontext pipeline loaded")
 
 # ───── Load YOLOv8 Segmentation Model ─────
 yolo = YOLO('yolov8s-seg.pt')  # can be replaced with your own weights