|
|
import torch |
|
|
import gradio as gr |
|
|
from pipeline_controlnet_sd_xl_raw import StableDiffusionXLControlNetRAWPipeline |
|
|
from diffusers import ControlNetModel, UniPCMultistepScheduler |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
import traceback |
|
|
|
|
|
|
|
|
pipe = StableDiffusionXLControlNetRAWPipeline.from_pretrained( |
|
|
"wencheng256/DiffusionRAW", |
|
|
torch_dtype=torch.float16 |
|
|
) |
|
|
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) |
|
|
pipe.enable_model_cpu_offload() |
|
|
|
|
|
|
|
|
def tensor_to_pil(img_tensor: torch.Tensor) -> Image.Image: |
|
|
if img_tensor.is_cuda: |
|
|
img_tensor = img_tensor.cpu() |
|
|
if img_tensor.dtype != torch.float32: |
|
|
img_tensor = img_tensor.float() |
|
|
img_tensor = img_tensor.clamp(0, 1) |
|
|
return transforms.ToPILImage()(img_tensor) |
|
|
|
|
|
|
|
|
def load_pth_data(pth_path): |
|
|
data = torch.load(pth_path) |
|
|
rgb_tensor = data["rgb"] |
|
|
raw_tensor = data["raw"] |
|
|
mask_tensor = data["mask"] |
|
|
cond_tensor = data["condition"] |
|
|
|
|
|
raw_image_pil = tensor_to_pil(raw_tensor[0][:, :448]) |
|
|
rgb_tensor_pil = tensor_to_pil(torch.flip(rgb_tensor[0], dims=[0])[:, :448]) |
|
|
mask_image_pil = tensor_to_pil(1 - mask_tensor[0]) |
|
|
|
|
|
return rgb_tensor_pil, raw_image_pil, mask_image_pil, raw_tensor, mask_tensor, cond_tensor |
|
|
|
|
|
|
|
|
def infer_fn(prompt, mask_edited, raw_tensor_state, mask_tensor_state, cond_tensor_state): |
|
|
try: |
|
|
if isinstance(mask_edited, dict): |
|
|
mask_edited = mask_edited["composite"] |
|
|
|
|
|
mask_edited_tensor = transforms.ToTensor()(mask_edited) |
|
|
mask_edited_tensor = 1-mask_edited_tensor[:1].unsqueeze(0).half() |
|
|
|
|
|
raw_t = raw_tensor_state.half() |
|
|
cond_t = cond_tensor_state.half() |
|
|
|
|
|
generator = torch.manual_seed(0) |
|
|
|
|
|
result = pipe( |
|
|
prompt=prompt, |
|
|
num_inference_steps=20, |
|
|
generator=generator, |
|
|
image=raw_t, |
|
|
mask_image=mask_edited_tensor, |
|
|
control_image=cond_t |
|
|
).images[0] |
|
|
|
|
|
return tensor_to_pil(result) |
|
|
|
|
|
except Exception as e: |
|
|
traceback.print_exc() |
|
|
return "Error occurred during inference. Please check the terminal logs!" |
|
|
|
|
|
|
|
|
def build_demo(): |
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown("# DiffusionRAW") |
|
|
|
|
|
pth_options = ["./data1.pth", "./data2.pth", "./data3.pth"] |
|
|
pth_selector = gr.Dropdown(choices=pth_options, value=pth_options[0], label="Select a PTH file") |
|
|
load_button = gr.Button("Load") |
|
|
|
|
|
with gr.Row(): |
|
|
raw_display = gr.Image(label="Raw Image", interactive=False) |
|
|
rgb_display = gr.Image(label="sRGB Image", interactive=False) |
|
|
mask_editor = gr.Sketchpad( |
|
|
label="Mask (Sketch)", |
|
|
interactive=True, |
|
|
width=512, |
|
|
height=512 |
|
|
) |
|
|
|
|
|
raw_tensor_state = gr.State() |
|
|
mask_tensor_state = gr.State() |
|
|
cond_tensor_state = gr.State() |
|
|
|
|
|
load_button.click( |
|
|
fn=load_pth_data, |
|
|
inputs=[pth_selector], |
|
|
outputs=[ |
|
|
rgb_display, |
|
|
raw_display, |
|
|
mask_editor, |
|
|
raw_tensor_state, |
|
|
mask_tensor_state, |
|
|
cond_tensor_state |
|
|
] |
|
|
) |
|
|
|
|
|
prompt_input = gr.Textbox(label="Prompt", value="An RAW Image.", lines=1) |
|
|
generate_button = gr.Button("Generate") |
|
|
output_image = gr.Image(label="Output") |
|
|
|
|
|
generate_button.click( |
|
|
fn=infer_fn, |
|
|
inputs=[ |
|
|
prompt_input, |
|
|
mask_editor, |
|
|
raw_tensor_state, |
|
|
mask_tensor_state, |
|
|
cond_tensor_state |
|
|
], |
|
|
outputs=[output_image] |
|
|
) |
|
|
|
|
|
return demo |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo = build_demo() |
|
|
demo.launch() |
|
|
|
