Update app.py

app.py

@@ -11,9 +11,30 @@ from diffusers.utils import load_image
 
 from huggingface_hub import hf_hub_download
 
-pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
-pipe.load_lora_weights("kontext-community/relighting-kontext-dev-lora-v3", weight_name="relighting-kontext-dev-lora-v3.safetensors", adapter_name="lora")
-pipe.set_adapters(["lora"], adapter_weights=[0.75])
+# --- Lazy pipeline init to avoid CUDA errors on Zero GPU ---
+PIPE = None
+
+def get_pipe():
+    """Create the pipeline lazily inside a GPU context when available.
+    On Hugging Face Zero GPU, CUDA is only present inside @spaces.GPU functions.
+    """
+    global PIPE
+    if PIPE is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.float16 if device == "cuda" else torch.bfloat16
+        pipe = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=dtype
+        )
+        pipe.load_lora_weights(
+            "kontext-community/relighting-kontext-dev-lora-v3",
+            weight_name="relighting-kontext-dev-lora-v3.safetensors",
+            adapter_name="lora",
+        )
+        pipe.set_adapters(["lora"], adapter_weights=[0.75])
+        pipe.to(device)
+        PIPE = pipe
+    return PIPE
+# ----------------------------------------------------------
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -286,14 +307,18 @@ def infer(input_image, prompt, illumination_dropdown, direction_dropdown, seed=4
     prompt_with_template = f"Relight the image{prompt_prefix}. {prompt} Maintain the identity of the foreground subjects."
     
     print(prompt_with_template)
-    
+
+    pipe = get_pipe()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    gen = torch.Generator(device=device).manual_seed(seed)
+
     image = pipe(
-        image=input_image, 
+        image=input_image,
         prompt=prompt_with_template,
         guidance_scale=guidance_scale,
         width=input_image.size[0],
         height=input_image.size[1],
-        generator=torch.Generator().manual_seed(seed),
+        generator=gen,
     ).images[0]
     return [input_image, image], seed, prompt_with_template