File size: 3,284 Bytes
17a4159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a07bcfd
17a4159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
113
114
115
116
117
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 gradio_imageslider import ImageSlider
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>
  Given an image with a lamp switched off, the model should turn it on.
</p>
<p>
  πŸ”Œ For higher resolution results with control over lighting intensity and warmth,&nbsp;
  <a href="https://finegrain.ai">head to the Finegrain API</a> πŸ”Œ
</p>
<p>
  <a href="https://discord.gg/zFKg5TjXub" target="_blank">[Discord]</a>
  <a href="https://github.com/finegrain-ai" target="_blank">[GitHub]</a>
  <a href="https://finegrain.ai">[Finegrain API]</a>
</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 = ImageSlider(label="Before / After")
            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,
        run_on_click=False,
    )

demo.launch(share=False)