Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| import random | |
| from PIL import Image | |
| from diffusers import FluxKontextPipeline | |
| from diffusers import FluxTransformer2DModel | |
| 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-v2", weight_name="relighting-kontext-dev-lora.safetensors", adapter_name="lora") | |
| pipe.set_adapters(["lora"], adapter_weights=[1.0]) | |
| MAX_SEED = np.iinfo(np.int32).max | |
| # Illumination options mapping | |
| ILLUMINATION_OPTIONS = { | |
| "sunset over sea": "add warm sunset light reflecting off water, creating a peaceful atmosphere with soft, diffused lighting and gentle color gradients", | |
| "sunshine from window": "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", | |
| "golden time": "Set golden color temperature. Increase contrast moderately. Add rim lighting on subjects. Enhance warm colors. Create soft directional shadows. Reduce blue tones.", | |
| "sci-fi RGB glowing, cyberpunk": "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects", | |
| "neon night, city":"add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts", | |
| "natural lighting": "Set neutral white color temperature. Maintain balanced exposure. Create soft shadows with gradual falloff. Preserve original color saturation. Add subtle ambient occlusion.", | |
| "warm atmosphere, at home, bedroom": "Set very warm color temperature. Reduce brightness moderately. Add soft diffused lighting. Increase red and yellow tones. Create minimal shadows.", | |
| "magic lit": "Add subtle colored rim lighting in purple and blue tones. Increase brightness selectively on edges. Add soft glow effects. Create ethereal light sources with no clear origin.", | |
| "evil, gothic, in a cave": "Reduce overall brightness significantly. Set cool color temperature. Increase contrast substantially. Create deep black shadows. Add single point light source from above or side.", | |
| "light and shadow": "Maximize contrast. Create equal areas of deep black and bright white. Sharp shadow boundaries. Maintain mid-tone details only in transition zones.", | |
| "shadow from window": "Add window frame shadow patterns. Create geometric shadow shapes. Maintain directional sunlight. Increase contrast in shadowed areas.", | |
| "soft studio lighting": "Add multiple diffused light sources. Reduce harsh shadows. Set neutral white color temperature. Create even illumination across subject. Minimize contrast.", | |
| "home atmosphere, cozy bedroom illumination": "Set very warm color temperature. Add table lamp or bedside lighting. Create pools of warm light. Reduce overall brightness.", | |
| "neon, Wong Kar-wai, warm": "Add warm neon lighting in amber and red tones. Create moody atmosphere with selective lighting. Increase film grain. Add lens flare effects.", | |
| "moonlight through curtains": "Set cool blue color temperature. Reduce brightness significantly. Add curtain shadow patterns. Create soft diffused lighting through fabric.", | |
| "stormy sky lighting": "Set dramatic lighting with rapid brightness changes. Add cool color temperature. Create low-key lighting with high contrast. Add rim lighting effects.", | |
| "underwater glow, deep sea": "Set blue-green color temperature. Add volumetric lighting effects. Reduce red channel significantly. Create particle effects in light beams. Add caustic light patterns.", | |
| "foggy forest at dawn": "Add volumetric fog effects. Set cool color temperature. Create god rays through trees. Reduce contrast moderately. Add atmospheric perspective.", | |
| "golden hour in a meadow": "Set golden color temperature. Add backlighting from low sun position. Create lens flare. Increase warm color saturation. Add rim lighting on grass.", | |
| "rainbow reflections, neon": "Add chromatic dispersion effects. Create rainbow light patterns. Maximize saturation. Add prismatic reflections. Set multiple colored light sources.", | |
| "cozy candlelight": "Set very warm orange color temperature. Add flickering light animation. Create localized warm pools of light. Increase shadows dramatically.", | |
| "apocalyptic, smoky atmosphere": "Reduce overall brightness moderately. Add orange-red tint from fires. Create volumetric smoke effects. Set warm color temperature. Add particle effects.", | |
| "red glow, emergency lights": "Replace lighting with red sources. Add flashing strobing effects. Increase contrast. Create harsh shadows. Set monochromatic red color scheme.", | |
| "mystical glow, enchanted forest": "Add supernatural green and blue glowing effects. Create floating light particles. Add bioluminescent elements. Set ethereal rim lighting.", | |
| "campfire light": "Set warm orange color temperature. Add flickering warm light source from below. Create dancing shadows. Increase warm color channels. Add spark particle effects.", | |
| "harsh, industrial lighting": "Add multiple bright fluorescent sources. Set cool white color temperature. Create hard shadows. Increase contrast significantly. Add industrial light fixtures.", | |
| "sunrise in the mountains": "Set neutral to warm color temperature. Add horizontal backlighting. Create atmospheric haze. Add lens flare from sun. Increase contrast on mountain silhouettes.", | |
| "evening glow in the desert": "Set warm color temperature. Add strong directional lighting from low angle. Create long shadows. Increase red and orange channels.", | |
| "moonlight in a dark alley": "Set cool blue color temperature. Reduce brightness substantially. Add single overhead light source. Create deep shadows with urban light pollution.", | |
| "golden glow at a fairground": "Add multiple warm colored light sources. Set golden color temperature. Create bokeh light effects. Add carnival lighting with bulbs.", | |
| "midnight in the forest": "Reduce brightness dramatically. Set cool color temperature. Add minimal ambient lighting. Create near-black shadows with subtle detail retention.", | |
| "purple and pink hues at twilight": "Set cool color temperature with magenta tint. Add purple and pink color grading. Create gradient sky lighting.", | |
| "foggy morning, muted light": "Reduce contrast moderately. Add volumetric fog effects. Set neutral color temperature. Create soft diffused lighting throughout.", | |
| "candle-lit room, rustic vibe": "Set very warm orange color temperature. Add multiple small warm light sources. Create atmospheric shadows. Add warm ambient lighting.", | |
| "fluorescent office lighting": "Set cool white color temperature. Add even overhead lighting. Create minimal shadows. Add slight green tint. Increase brightness uniformly.", | |
| "lightning flash in storm": "Add brief intense white light burst. Create stark shadows momentarily. Set very cool color temperature. Add dramatic contrast.", | |
| "night, cozy warm light from fireplace": "Set warm orange-red color temperature. Add warm light source from one side. Create flickering effects. Add warm ambient glow.", | |
| "ethereal glow, magical forest": "Add supernatural lighting with no clear source. Create floating light particles. Add rim lighting in blue and green tones. Set mystical atmosphere.", | |
| "dusky evening on a beach": "Set cool color temperature. Add horizon backlighting. Create silhouette effects. Add water reflections of sky colors.", | |
| "afternoon light filtering through trees": "Add dappled lighting patterns. Create leaf shadow patterns. Set neutral color temperature. Add green color cast from foliage.", | |
| "blue neon light, urban street": "Replace ambient lighting with blue neon. Add neon glow effects on surfaces. Create urban light pollution. Set cool color temperature.", | |
| "red and blue police lights in rain": "Add alternating red and blue light sources. Create reflections on wet surfaces. Add emergency lighting strobing effects.", | |
| "aurora borealis glow, arctic landscape": "Add green and purple sky lighting. Create dancing light effects. Set cool color temperature.", | |
| "sunrise through foggy mountains": "Add horizontal warm lighting through mist. Create volumetric fog effects. Set warm color temperature. Add atmospheric perspective.", | |
| "golden hour on a city skyline": "Set golden color temperature. Add backlighting on buildings. Create silhouette effects. Add lens flare from sun position.", | |
| "mysterious twilight, heavy mist": "Reduce visibility substantially. Add volumetric fog throughout. Set cool color temperature. Create atmospheric depth.", | |
| "early morning rays, forest clearing": "Add god rays through trees. Set warm color temperature. Create dappled sunlight patterns. Add atmospheric haze.", | |
| "colorful lantern light at festival": "Add multiple colored light sources from lanterns. Create warm festival lighting. Set varied color temperatures. Add bokeh effects.", | |
| "soft glow through stained glass": "Add colored light filtering effects. Create rainbow light patterns on surfaces. Set warm color temperature with color variations.", | |
| "harsh spotlight in dark room": "Add single intense directional light source. Create dramatic shadows. Set high contrast. Reduce ambient lighting to minimum.", | |
| "mellow evening glow on a lake": "Set warm color temperature. Add water reflections. Create calm atmospheric lighting. Add subtle lens flare.", | |
| "crystal reflections in a cave": "Add sparkle and reflection effects on crystal surfaces. Create prismatic light dispersion. Set cool color temperature. Add caustic lighting patterns.", | |
| "vibrant autumn lighting in a forest": "Enhance warm colors significantly. Set warm color temperature. Add dappled sunlight through autumn leaves.", | |
| "gentle snowfall at dusk": "Add cool blue color temperature. Create soft diffused lighting. Add snowflake particle effects. Reduce contrast.", | |
| "hazy light of a winter morning": "Set neutral color temperature. Add atmospheric haze. Create soft lighting conditions. Reduce saturation slightly.", | |
| "soft, diffused foggy glow": "Add volumetric fog throughout scene. Reduce contrast substantially. Create soft lighting with no harsh shadows. Set neutral color temperature.", | |
| "underwater luminescence": "Set blue-green color temperature. Add bioluminescent glowing effects. Create caustic light patterns. Add particle effects in water.", | |
| "rain-soaked reflections in city lights": "Add reflections on wet surfaces. Create light streaking effects. Set urban lighting with multiple colored sources. Add rain particle effects.", | |
| "golden sunlight streaming through trees": "Add god rays effect. Set golden color temperature. Create dappled light patterns. Add atmospheric particles in light beams.", | |
| "fireflies lighting up a summer night": "Add small glowing points of light. Set warm color temperature. Create subtle animated glow effects. Maintain dark ambient lighting.", | |
| "glowing embers from a forge": "Add orange-red glowing particle effects. Set very warm color temperature. Create intense heat glow. Add sparking effects.", | |
| "dim candlelight in a gothic castle": "Set very warm orange color temperature. Add multiple small flickering light sources. Create dramatic shadows and enhance stone textures.", | |
| "midnight sky with bright starlight": "Set very cool blue color temperature. Add point light sources for stars. Create minimal ambient lighting. Add astronomical lighting effects.", | |
| "warm sunset in a rural village": "Set warm color temperature. Add golden hour lighting. Create peaceful ambient lighting. Add warm color enhancement.", | |
| "flickering light in a haunted house": "Add animated flickering light effects. Create unstable lighting conditions. Set cool color temperature with warm light sources.", | |
| "desert sunset with mirage-like glow": "Set warm orange color temperature. Add heat distortion effects. Create atmospheric haze. Add intense warm lighting.", | |
| "golden beams piercing through storm clouds": "Add dramatic god rays through clouds. Set golden color temperature. Create high contrast between light and shadow areas." | |
| } | |
| # Lighting direction options | |
| DIRECTION_OPTIONS = { | |
| "none": "", | |
| "left side": "Position the light source from the left side of the frame, creating shadows falling to the right.", | |
| "right side": "Position the light source from the right side of the frame, creating shadows falling to the left.", | |
| "top": "Position the light source from directly above, creating downward shadows.", | |
| "top left": "Position the light source from the top left corner, creating diagonal shadows falling down and to the right.", | |
| "top right": "Position the light source from the top right corner, creating diagonal shadows falling down and to the left.", | |
| "bottom": "Position the light source from below, creating upward shadows and dramatic under-lighting.", | |
| "front": "Position the light source from the front, minimizing shadows and creating even illumination.", | |
| "back": "Position the light source from behind the subject, creating silhouette effects and rim lighting." | |
| } | |
| def infer(input_image, prompt, illumination_dropdown, direction_dropdown, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| input_image = input_image.convert("RGB") | |
| # Combine illumination and direction if direction is not "Auto" | |
| if direction_dropdown != "none": | |
| prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown}" | |
| else: | |
| prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown}" | |
| prompt_with_template = f"Relight the image, {prompt_prefix}. {prompt} Maintain the identity of the foreground subjects." | |
| image = pipe( | |
| 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), | |
| ).images[0] | |
| return image, seed | |
| def update_prompt_from_dropdown(illumination_option): | |
| """Update the prompt textbox based on dropdown selection""" | |
| if illumination_option == "custom": | |
| return "" # Clear the prompt for custom input | |
| else: | |
| return ILLUMINATION_OPTIONS[illumination_option] | |
| css=""" | |
| #col-container { | |
| margin: 0 auto; | |
| max-width: 960px; | |
| } | |
| """ | |
| 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"""Kontext[dev] finetuned for scene relighting ✨ | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="Upload the image for relighting", type="pil") | |
| with gr.Row(): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| show_label=True, | |
| max_lines=3, | |
| placeholder="select an illumination style above or enter your custom description...", | |
| value="Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", | |
| container=True | |
| ) | |
| illumination_dropdown = gr.Dropdown( | |
| choices=list(ILLUMINATION_OPTIONS.keys()) + ["custom"], | |
| value="sunshine from window", | |
| label="Choose Illumination Style", | |
| info="Select a predefined lighting style or choose 'Custom' to use your own prompt", | |
| scale=2 | |
| ) | |
| direction_dropdown = gr.Dropdown( | |
| choices=list(DIRECTION_OPTIONS.keys()), | |
| value="auto", | |
| label="Light Direction", | |
| info="Choose the direction of the light source", | |
| scale=1 | |
| ) | |
| run_button = gr.Button("Run", scale=0, variant="primary") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| 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, | |
| ) | |
| with gr.Column(): | |
| result = gr.Image(label="Result", show_label=False, interactive=False) | |
| # update prompt when dropdown changes | |
| illumination_dropdown.change( | |
| fn=update_prompt_from_dropdown, | |
| inputs=[illumination_dropdown], | |
| outputs=[prompt] | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["./assets/pexels-creationhill-1681010.jpg", "Add multiple colored light sources from lanterns. Create warm festival lighting. Set varied color temperatures. Add bokeh effects.", "colorful lantern light at festival", "auto", 0, True, 2.5], | |
| ["./assets/pexels-creationhill-1681010.jpg", "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects", "sci-fi RGB glowing, cyberpunk", "left side", 0, True, 2.5], | |
| ["./assets/pexels-moose-photos-170195-1587009.jpg", "Set blue-green color temperature. Add volumetric lighting effects. Reduce red channel significantly. Create particle effects in light beams. Add caustic light patterns.", "underwater glow, deep sea", "top", 0, True, 2.5], | |
| ["./assets/pexels-moose-photos-170195-1587009.jpg", "Replace lighting with red sources. Add flashing strobing effects. Increase contrast. Create harsh shadows. Set monochromatic red color scheme.", "red glow, emergency lights", "right side", 0, True, 2.5], | |
| ["./assets/pexels-simon-robben-55958-614810.jpg", "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", "sunshine from window", "top right", 0, True, 2.5], | |
| ["./assets/pexels-simon-robben-55958-614810.jpg", "add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts", "neon night, city", "top left", 0, True, 2.5] | |
| ], | |
| inputs=[input_image, prompt, illumination_dropdown, direction_dropdown, 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, illumination_dropdown, direction_dropdown, seed, randomize_seed, guidance_scale], | |
| outputs = [result, seed] | |
| ) | |
| demo.launch() |