File size: 3,104 Bytes
17a4159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a07bcfd
17a4159
dd8f5e4
 
 
 
 
17a4159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dd8f5e4
17a4159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b5c8a75
17a4159
 
 
dd8f5e4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import os
import random
import time
from typing import Any

import gradio as gr
import pillow_avif  # noqa: F401
import pillow_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)

pillow_heif.register_avif_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/weights-light-switcher-space"
model = light_switcher_lite.Model.from_pretrained(path, device=DEVICE_CPU, dtype=DTYPE)

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>
    For premium-quality results,
    <a href="https://chat.finegrain.ai?utm_source=hf&utm_campaign=light-switcher">try Finegrain Chat</a>
    — it's free to test!
</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, show_fullscreen_button=False)
            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, run_button],
        fn=process,
        cache_examples=True,
        cache_mode="eager",
        run_on_click=False,
    )

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