Update app.py

app.py

@@ -1,116 +1,42 @@
 import gradio as gr
 import numpy as np
+
 import spaces
 import torch
 import random
 from PIL import Image
-
 from diffusers import FluxKontextPipeline
-from huggingface_hub import hf_hub_download
-
-# --------------------
-# Globals (CPU-only; DO NOT touch CUDA here)
-# --------------------
-MAX_SEED = np.iinfo(np.int32).max
-
-DEVICE = "cpu"   # set in gpu_startup()
-DTYPE = torch.float32  # set in gpu_startup()
+from diffusers import FluxTransformer2DModel
+from diffusers.utils import load_image
 
-PIPE = None      # cached pipeline (loaded on first GPU run)
-LORA_LOADED = False
-
-# --------------------
-# ZeroGPU: allocate GPU at startup & decide dtype
-# --------------------
-@spaces.GPU
-def gpu_startup():
-    """
-    Runs when the Space boots on ZeroGPU. Safe place to detect CUDA and dtype.
-    """
-    global DEVICE, DTYPE
-    has_cuda = torch.cuda.is_available()
-    DEVICE = "cuda" if has_cuda else "cpu"
-
-    # Prefer bfloat16 on CUDA if supported; otherwise fp16; CPU uses fp32
-    if DEVICE == "cuda":
-        if torch.cuda.is_bf16_supported():
-            DTYPE = torch.bfloat16
-        else:
-            DTYPE = torch.float16
-    else:
-        DTYPE = torch.float32
-
-    print(f"[startup] device={DEVICE}, dtype={DTYPE}")
-
-# --------------------
-# Lazy model loader (runs inside GPU context)
-# --------------------
-def get_pipeline():
-    """
-    Create/cache the FluxKontextPipeline and load LoRA once.
-    Must be called from within a @spaces.GPU function.
-    """
-    global PIPE, LORA_LOADED
-
-    if PIPE is None:
-        # Load base pipeline with chosen dtype, then move to device
-        PIPE = FluxKontextPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-Kontext-dev",
-            torch_dtype=DTYPE
-        ).to(DEVICE)
+from huggingface_hub import hf_hub_download
 
-    if not LORA_LOADED:
-        # Load LoRA weights once and set adapter
-        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=[1.0])
-        LORA_LOADED = True
+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=[1.0])
 
-    return PIPE
+MAX_SEED = np.iinfo(np.int32).max
 
-# --------------------
-# Inference (GPU entrypoint)
-# --------------------
 @spaces.GPU
 def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
-    pipe = get_pipeline()
-
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
-    if input_image is None:
-        return None, seed
-
+        
     input_image = input_image.convert("RGB")
-
-    prompt_with_template = (
-        f"Change the lighting conditions in this image and add {prompt}. "
-        "Change the background details but maintain the foreground. "
-        "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 under different lighting conditions."
-    )
-
-    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
-
+    prompt_with_template = f"Change 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."
+    
     image = pipe(
-        image=input_image,
+        image=input_image, 
         prompt=prompt_with_template,
-        guidance_scale=float(guidance_scale),
+        guidance_scale=guidance_scale,
         width=input_image.size[0],
         height=input_image.size[1],
-        generator=generator,
+        generator=torch.Generator().manual_seed(seed),
     ).images[0]
-
     return image, seed
 
-# --------------------
-# UI
-# --------------------
-css = """
+css="""
 #col-container {
     margin: 0 auto;
     max-width: 960px;
@@ -118,9 +44,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
+    
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# FLUX.1 Kontext [dev] Relight 💡")
-        gr.Markdown("Kontext[dev] used for object relighting ✨")
+        gr.Markdown(f"""# FLUX.1 Kontext [dev] Relight 💡
+        """)
+        gr.Markdown(f"""Kontext[dev] used for object relighting ✨
+         """)
 
         with gr.Row():
             with gr.Column():
@@ -134,36 +63,48 @@ with gr.Blocks(css=css) as demo:
                         container=False,
                     )
                     run_button = gr.Button("Run", scale=0)
-
                 with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+                    
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10, step=0.1, value=2.5)
-
+                    
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1,
+                        maximum=10,
+                        step=0.1,
+                        value=2.5,
+                    )       
+                    
             with gr.Column():
                 result = gr.Image(label="Result", show_label=False, interactive=False)
-
+        
         gr.Examples(
-            examples=[
-                ["./assets/5_before.png", "sunset over sea lighting coming from the top right part of the photo", 0, True, 2.5],
-                ["./assets/3_before.png", "sci-fi RGB glowing, studio lighting", 0, True, 2.5],
-                ["./assets/2_before.png", "neon light, city", 0, True, 2.5],
-                ["./assets/before_6.png", "bright sunlight, warm, luminous", 0, True, 2.5],
-            ],
-            inputs=[input_image, prompt, seed, randomize_seed, guidance_scale],
-            outputs=[result, seed],
-            fn=infer,
-            cache_examples="lazy",
-        )
-
-        gr.on(
-            triggers=[run_button.click, prompt.submit],
-            fn=infer,
-            inputs=[input_image, prompt, seed, randomize_seed, guidance_scale],
-            outputs=[result, seed],
-        )
+        examples=[
+            ["./assets/5_before.png", "sunset over sea lighting coming from the top right part of the photo", 0, True, 2.5],
+            ["./assets/3_before.png", "sci-fi RGB glowing, studio lighting",0, True,2.5],
+            ["./assets/2_before.png", "neon light, city",0, True, 2.5],
+            ["./assets/before_6.png", "bright sunlight, warm, luminous", 0, True, 2.5]
+        ],
+        inputs=[input_image, prompt, seed, randomize_seed, guidance_scale],
+        outputs=[result, seed],
+        fn=infer,
+        cache_examples="lazy"
+    )
+    
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = [input_image, prompt, seed, randomize_seed, guidance_scale],
+        outputs = [result, seed]
+    )
+    
 
-if __name__ == "__main__":
-    # Ensure ZeroGPU allocates a GPU at boot and dtype/device are decided early
-    gpu_startup()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()