flx-upscale / app.py
fantaxy's picture
Update app.py
833ef3a verified
raw
history blame
5.87 kB
import logging
import random
import warnings
import os
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from gradio_imageslider import ImageSlider
from PIL import Image
from huggingface_hub import snapshot_download
css = """
#col-container {
margin: 0 auto;
max-width: 512px;
}
"""
if torch.cuda.is_available():
power_device = "GPU"
device = "cuda"
else:
power_device = "CPU"
device = "cpu"
huggingface_token = os.getenv("HUGGINFACE_TOKEN")
model_path = snapshot_download(
repo_id="black-forest-labs/FLUX.1-dev",
repo_type="model",
ignore_patterns=["*.md", "*..gitattributes"],
local_dir="FLUX.1-dev",
token=huggingface_token, # type a new token-id.
)
# Load pipeline
controlnet = FluxControlNetModel.from_pretrained(
"jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
).to(device)
pipe = FluxControlNetPipeline.from_pretrained(
model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
)
pipe.to(device)
MAX_SEED = 1000000
MAX_PIXEL_BUDGET = 1024 * 1024
def process_input(input_image, upscale_factor, **kwargs):
w, h = input_image.size
w_original, h_original = w, h
aspect_ratio = w / h
was_resized = False
if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
warnings.warn(
f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels."
)
gr.Info(
f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing input to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels budget."
)
input_image = input_image.resize(
(
int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
)
)
was_resized = True
# resize to multiple of 8
w, h = input_image.size
w = w - w % 8
h = h - h % 8
return input_image.resize((w, h)), w_original, h_original, was_resized
@spaces.GPU#(duration=42)
def infer(
seed,
randomize_seed,
input_image,
num_inference_steps,
upscale_factor,
controlnet_conditioning_scale,
progress=gr.Progress(track_tqdm=True),
):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
true_input_image = input_image
input_image, w_original, h_original, was_resized = process_input(
input_image, upscale_factor
)
# rescale with upscale factor
w, h = input_image.size
control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
generator = torch.Generator().manual_seed(seed)
gr.Info("Upscaling image...")
image = pipe(
prompt="",
control_image=control_image,
controlnet_conditioning_scale=controlnet_conditioning_scale,
num_inference_steps=num_inference_steps,
guidance_scale=3.5,
height=control_image.size[1],
width=control_image.size[0],
generator=generator,
).images[0]
if was_resized:
gr.Info(
f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
)
# resize to target desired size
image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
image.save("output.jpg")
# convert to numpy
return [true_input_image, image, seed]
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
with gr.Row():
run_button = gr.Button(value="Run")
with gr.Row():
with gr.Column(scale=4):
input_im = gr.Image(label="Input Image", type="pil")
with gr.Column(scale=1):
num_inference_steps = gr.Slider(
label="Number of Inference Steps",
minimum=8,
maximum=50,
step=1,
value=28,
)
upscale_factor = gr.Slider(
label="Upscale Factor",
minimum=1,
maximum=4,
step=1,
value=4,
)
controlnet_conditioning_scale = gr.Slider(
label="Controlnet Conditioning Scale",
minimum=0.1,
maximum=1.5,
step=0.1,
value=0.6,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=42,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
result = ImageSlider(label="Input / Output", type="pil", interactive=True)
examples = gr.Examples(
examples=[
[42, False, "z1.webp", 28, 4, 0.6],
[42, False, "z2.webp", 28, 4, 0.6],
],
inputs=[
seed,
randomize_seed,
input_im,
num_inference_steps,
upscale_factor,
controlnet_conditioning_scale,
],
fn=infer,
outputs=result,
cache_examples="lazy",
)
gr.on(
[run_button.click],
fn=infer,
inputs=[
seed,
randomize_seed,
input_im,
num_inference_steps,
upscale_factor,
controlnet_conditioning_scale,
],
outputs=result,
show_api=False,
# show_progress="minimal",
)
demo.queue().launch(share=False)