Spaces:
Running
on
Zero
Running
on
Zero
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,
<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)
|