app.py

import os
import random
import time
from typing import Any

import gradio as gr
import pi_heif
import spaces
import torch
from PIL import Image


class ModuleInterface:
    def __init__(self, d: dict[str, Any]):
        self.d = d

    @property
    def Model(self) -> Any:
        return self.d["Model"]

    def process(self, input_image: Image.Image, model: Any, seed: int) -> Image.Image:
        return self.d["process"](input_image, model, seed)


assert (src := os.getenv("LIGHT_SWITCHER_LITE")), "LIGHT_SWITCHER_LITE not set"
exec_globals: dict[str, Any] = {}
exec(src, exec_globals)
light_switcher_lite = ModuleInterface(exec_globals)

pi_heif.register_heif_opener()


DEVICE_CPU = torch.device("cpu")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32

# CPU -> GPU dance because of ZeroGPU

path = "finegrain/merged-turbo-light-switcher"
model = light_switcher_lite.Model.from_pretrained(path, device=DEVICE_CPU, dtype=DTYPE)

model = model.to(DEVICE)


@spaces.GPU
def process(input_image: Image.Image, seed: int = 42) -> tuple[tuple[Image.Image, Image.Image], dict[str, Any]]:
    output_image = light_switcher_lite.process(input_image, model, seed)
    resized_input_image = input_image.resize(output_image.size)
    return ((resized_input_image, output_image), gr.update(value=random.choice(BUTTON_LABELS)))


TITLE = """
<h1>Finegrain Light Switcher (Lite Version)</h1>
<p>
    Instantly turn lamps on in your images.
</p><p>
    This Space showcases one of Finegrain's earlier generative AI models. Today, we're building the
    Agentic Camera: an on-device system that produces the intended photo at capture time, eliminating
    image editing in most cases.
</p><p>
    ✨ Learn more at <a href="https://finegrain.ai">finegrain.ai</a>
    or follow <a href="https://huggingface.co/finegrain">Finegrain</a> for updates.
</p>
"""

BUTTON_LABELS = [
    "💡",
    "Let there be light!",
    "Light it up like a Christmas tree!",
    "Turn it on!",
    "Aziz, Light!",
    "Make it shine ✨",
    "Flip the magic switch.",
]

random.seed(time.time())

with gr.Blocks() as demo:
    gr.HTML(TITLE)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.ClearButton(components=None, value=random.choice(BUTTON_LABELS))
        with gr.Column():
            output_slider = gr.ImageSlider(label="Before / After", max_height=1500, buttons=["download"])
            run_button.add(output_slider)

    with gr.Accordion("Advanced Options", open=False):
        seed = gr.Slider(minimum=0, maximum=999, value=42, step=1, label="Seed")

    run_button.click(
        fn=process,
        inputs=[input_image, seed],
        outputs=[output_slider, run_button],
    )

    gr.Examples(
        examples=[
            "examples/01.webp",
            "examples/02.webp",
            "examples/03.webp",
            "examples/04.webp",
            "examples/05.webp",
        ],
        inputs=[input_image],
        outputs=[output_slider],
        fn=lambda image: process(image)[0],  # skip the ClearButton update (breaks example caching)
        cache_examples=True,
        cache_mode="eager",
        run_on_click=False,
    )

demo.launch(share=False, ssr_mode=False)