Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,564 Bytes
7acb2a5 4686696 7acb2a5 4686696 7acb2a5 4686696 7acb2a5 4686696 7acb2a5 4686696 7acb2a5 15bf55b 7acb2a5 29dbe4a |
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 |
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="InstantX/InstantIR", filename="models/adapter.pt", local_dir=".")
hf_hub_download(repo_id="InstantX/InstantIR", filename="models/aggregator.pt", local_dir=".")
hf_hub_download(repo_id="InstantX/InstantIR", filename="models/previewer_lora_weights.bin", local_dir=".")
import torch
from PIL import Image
from diffusers import DDPMScheduler
from schedulers.lcm_single_step_scheduler import LCMSingleStepScheduler
from module.ip_adapter.utils import load_adapter_to_pipe
from pipelines.sdxl_instantir import InstantIRPipeline
# prepare models under ./models
instantir_path = f'./models'
# load pretrained models
pipe = InstantIRPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=torch.float16)
# load adapter
load_adapter_to_pipe(
pipe,
f"{instantir_path}/adapter.pt",
image_encoder_or_path = 'facebook/dinov2-large',
)
# load previewer lora
pipe.prepare_previewers(instantir_path)
pipe.scheduler = DDPMScheduler.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', subfolder="scheduler")
lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)
# load aggregator weights
pretrained_state_dict = torch.load(f"{instantir_path}/aggregator.pt")
pipe.aggregator.load_state_dict(pretrained_state_dict)
# send to GPU and fp16
pipe.to(device='cuda', dtype=torch.float16)
pipe.aggregator.to(device='cuda', dtype=torch.float16)
PROMPT = "Photorealistic, highly detailed, hyper detailed photo - realistic maximum detail, 32k, \
ultra HD, extreme meticulous detailing, skin pore detailing, \
hyper sharpness, perfect without deformations, \
taken using a Canon EOS R camera, Cinematic, High Contrast, Color Grading. "
NEG_PROMPT = "blurry, out of focus, unclear, depth of field, over-smooth, \
sketch, oil painting, cartoon, CG Style, 3D render, unreal engine, \
dirty, messy, worst quality, low quality, frames, painting, illustration, drawing, art, \
watermark, signature, jpeg artifacts, deformed, lowres"
def infer(prompt, input_image, steps=30, cfg_scale=7.0, guidance_end=1.0,
creative_restoration=False, seed=3407, height=1024, width=1024):
# load a broken image
low_quality_image = Image.open(input_image).convert("RGB")
lq = [resize_img(low_quality_image, size=(width, height))]
generator = torch.Generator(device=device).manual_seed(seed)
timesteps = [
i * (1000//steps) + pipe.scheduler.config.steps_offset for i in range(0, steps)
]
timesteps = timesteps[::-1]
prompt = PROMPT if len(prompt)==0 else prompt
neg_prompt = NEG_PROMPT
# InstantIR restoration
image = pipe(
prompt=[prompt]*len(lq),
image=lq,
num_inference_steps=steps,
generator=generator,
timesteps=timesteps,
negative_prompt=[neg_prompt]*len(lq),
guidance_scale=cfg_scale,
previewer_scheduler=lcm_scheduler,
).images[0]
return image
import gradio as gr
with gr.Blocks() as demo:
with gr.Column():
with gr.Row():
with gr.Column():
lq_img = gr.Image(label="Low-quality image", type="filepath")
with gr.Group():
prompt = gr.Textbox(label="Prompt", value="")
submit_btn = gr.Button("InstantIR magic!")
output_img = gr.Image(label="InstantIR restored")
submit_btn.click(
fn=infer,
inputs=[prompt, lq_img],
outputs=[output_img]
)
demo.launch(show_error=True) |