Update app.py

app.py

@@ -5,22 +5,16 @@ import spaces
 import torch
 import random
 from PIL import Image
-
-#from kontext_pipeline import FluxKontextPipeline
-from pipeline_flux_kontext import FluxKontextPipeline
+from diffusers import FluxKontextPipeline
 from diffusers import FluxTransformer2DModel
 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")
+pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
 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")
-
 @spaces.GPU
 def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
     
@@ -28,7 +22,7 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
         seed = random.randint(0, MAX_SEED)
         
     input_image = input_image.convert("RGB")
-    prompt_with_template = f"hange the lighting conditions in this image and add {prompt}. change the background details but maintain the forground. Lighting determines how bright or dark different parts of the image appear, where shadows fall, and how colors look. When you relight an image, you're simulating what the photo would look like if it were taken under different lighting conditions."
+    prompt_with_template = f"Change the lighting conditions in this image to {prompt}. Modify or replace the background as needed to create an environment that naturally produces these lighting conditions, while maintaining the foreground subject. Lighting determines how bright or dark different parts of the image appear, where shadows fall, and how colors look. When you relight an image, you're simulating what the photo would look like if it were taken under different lighting conditions with an appropriate background setting."
     
     image = pipe(
         image=input_image, 
@@ -50,8 +44,10 @@ css="""
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 Kontext [dev] Relight
+        gr.Markdown(f"""# FLUX.1 Kontext [dev] Relight 💡
         """)
+        gr.Markdown(f"""Kontext[dev] used for object relighting ✨
+         """)
 
         with gr.Row():
             with gr.Column():