gradio_demo.py

# -------------------------------------------
#
# File Name: gradio_demo.py
# Author: WANG Yiyang
# Created: October.8, 2025
# Description: Gradio demo for DiffCamera.
# -------------------------------------------
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
from transformers import pipeline
import argparse
import copy
import cv2
import os
from peft import LoraConfig, set_peft_model_state_dict
from peft.utils import get_peft_model_state_dict
from diffusers import (
    AutoencoderKL,
    FlowMatchEulerDiscreteScheduler,
)
from diffusers.utils.torch_utils import is_compiled_module
from src.flux.pipeline.pipeline_flux_new import FluxPipeline, flux_pipeline_call
from src.flux.model.transformer_flux_new import FluxTransformer2DModel
from accelerate import Accelerator
import torch
import torch.nn as nn
from utils import parse_args
import safetensors
from diffusers.utils import (
    check_min_version,
    convert_unet_state_dict_to_peft,
    is_wandb_available,
)
from PIL import Image
from torchvision import transforms
import numpy as np


def load_diffcamera_model(args):
    # Load scheduler and models
    accelerator = Accelerator(
        mixed_precision=args.mixed_precision,
    )
    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="scheduler"
    )
    noise_scheduler_copy = copy.deepcopy(noise_scheduler)
    device = args.device
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="vae",
        revision=args.revision,
        variant=args.variant,
    )
    transformer = FluxTransformer2DModel.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant
    )
    transformer.set_frame_embedder()

    # We only train the additional adapter LoRA layers
    transformer.requires_grad_(False)
    vae.requires_grad_(False)

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
    # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )

    vae.to(device, dtype=weight_dtype)
    transformer.to(device, dtype=weight_dtype)

    in_channel = 4
    encode_feature_dim = 4096

    # from (k,x,y,d) to T5
    camera_emb_token_num = 1
    camera_projector = torch.nn.Linear(in_channel, encode_feature_dim * camera_emb_token_num) # from (K, x, y) to a positional embedding
    camera_projector = camera_projector.to(device=device, dtype=weight_dtype)
    camera_projector.requires_grad_(True)

    camera_projector_768 = torch.nn.Linear(in_channel, 768).to(device=device, dtype=weight_dtype)

    transformer = accelerator.prepare(transformer)
    camera_projector = accelerator.prepare(camera_projector)
    camera_projector_768 = accelerator.prepare(camera_projector_768)

    def get_linear_modules(model):
        linear_modules = []
        
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                linear_modules.append(name)
        
        return linear_modules

    target_modules = get_linear_modules(transformer)

    lora_rank = 64
    transformer_lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_rank,
        init_lora_weights="gaussian",
        target_modules=target_modules,
    )
    transformer.add_adapter(transformer_lora_config)

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    def load_model_hook(models, input_dir):
        transformer_ = None
        camera_projector_ = None
        projector_ = None
        camera_projector_768_ = None

        while len(models) > 0:
            model = models.pop()

            if isinstance(model, type(unwrap_model(transformer))):
                transformer_ = model
            elif model is unwrap_model(camera_projector):
                camera_projector_ = model

            elif model is unwrap_model(camera_projector_768):
                camera_projector_768_ = model
            else:
                raise ValueError(f"unexpected save model: {model.__class__}")

        # load projector weights
        if projector_ is not None:
            projector_state_dict = safetensors.torch.load_file(
                os.path.join(input_dir, "projector.safetensors"),
            )
            projector_.load_state_dict(projector_state_dict)
            print(f"Projector loaded from {os.path.join(input_dir, 'projector.safetensors')}")
        if camera_projector_ is not None:
            camera_projector_state_dict = safetensors.torch.load_file(
                os.path.join(input_dir, "camera_projector.safetensors"),
            )
            camera_projector_.load_state_dict(camera_projector_state_dict)
            print(f"Camera projector loaded from {os.path.join(input_dir, 'camera_projector.safetensors')}")
        if camera_projector_768_ is not None:
            camera_projector_768_state_dict = safetensors.torch.load_file(
                os.path.join(input_dir, "camera_projector_768.safetensors"),
            )
            camera_projector_768_.load_state_dict(camera_projector_768_state_dict)
            print(f"Camera projector 768 loaded from {os.path.join(input_dir, 'camera_projector_768.safetensors')}")
        
        # load lora weights
        lora_state_dict = FluxPipeline.lora_state_dict(input_dir)

        transformer_state_dict = {
            f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")
        }
        transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
        incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
        if incompatible_keys is not None:
            # check only for unexpected keys
            unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
            if unexpected_keys:
                logger.warning(
                    f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
                    f" {unexpected_keys}. "
                )

        # Make sure the trainable params are in float32. This is again needed since the base models
        # are in `weight_dtype`. More details:
        # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
        if args.mixed_precision == "fp16":
            models = [transformer_]
            if camera_projector_ is not None:
                models.append(camera_projector_)
            if projector_ is not None:
                models.append(projector_)
            # only upcast trainable parameters (LoRA) into fp32
            cast_training_params(models)

    accelerator.register_load_state_pre_hook(load_model_hook)
    accelerator.print(f"Resuming from checkpoint {args.resume_from_checkpoint}")
    accelerator.load_state(args.resume_from_checkpoint)

    transformer = unwrap_model(transformer)
    true_cfg_scale = args.cfg_scale
    pipeline_args = {
        "vae": vae,
        "transformer": unwrap_model(transformer),
        "scheduler": copy.deepcopy(noise_scheduler),
        "image_lst": None,
        'prompt_embeds': None,
        'pooled_prompt_embeds': None,
        'text_ids': None,
        "height":args.image_size, "width": args.image_size, 
        "num_inference_steps":20, 
        "true_cfg_scale": true_cfg_scale,
    }

    return pipeline_args, weight_dtype, camera_projector, camera_projector_768

def diffcamera_inference(image, depth, K_scalar, x, y, pipeline_args, args, weight_dtype, camera_projector, camera_projector_768):
    # conver image, depth to tensors
    device = args.device
    vae_dtype = pipeline_args['vae'].dtype

    encode_feature_dim = 4096

    # center crop
    image = image.resize((args.image_size, args.image_size))
    depth = depth.resize((args.image_size, args.image_size))

    image = transforms.ToTensor()(image).to(device, dtype=vae_dtype)
    depth = transforms.ToTensor()(depth).to(device, dtype=vae_dtype)

    K_scalar = torch.tensor([[K_scalar*1.0]]).to(args.device, dtype=weight_dtype)
    focus_coordinates = torch.tensor([[x, y]]).to(args.device, dtype=weight_dtype)
    K_scalar = K_scalar / 31.0
    projector_input = torch.cat([K_scalar, focus_coordinates], dim=1).to(device, dtype=weight_dtype)

    x, y = int(focus_coordinates[:, 0] * args.image_size), int(focus_coordinates[:, 1] * args.image_size)

    d = depth[0, x, y].unsqueeze(0).unsqueeze(0)
    d = d.to(device, dtype=weight_dtype)
    if args.zero_depth:
        d = torch.zeros_like(d).to(device, dtype=weight_dtype)
    projector_input = torch.cat([projector_input, d], dim=1).to(device, dtype=weight_dtype)

    if args.include_depth and args.zero_depth:
        depth = torch.zeros_like(depth).to(device, dtype=weight_dtype)

    with torch.no_grad():
        camera_emb = camera_projector(projector_input).to(dtype=weight_dtype)
        camera_emb = camera_emb.view(-1, camera_emb_token_num, encode_feature_dim)
        camera_emb_768 = camera_projector_768(projector_input).to(dtype=weight_dtype)
    text_ids = torch.zeros(camera_emb.shape[1], 3).to(device=device, dtype=weight_dtype)
    if args.include_depth:
        image_lst = [image, depth]
    else:
        image_lst = [image]
    pipeline_args['image_lst'] = image_lst
    pipeline_args['prompt_embeds'] = camera_emb
    pipeline_args['pooled_prompt_embeds'] = camera_emb_768
    pipeline_args['text_ids'] = text_ids

    generator = None
    pred_image = flux_pipeline_call(**pipeline_args, generator=generator).images[0]

    return pred_image

args = parse_args()
pipeline_args, weight_dtype, camera_projector, camera_projector_768 = load_diffcamera_model(args)
pipe = pipeline(task="depth-estimation", model=args.depth_model_path, device=args.device)

def draw_focus_box_on_image(image, norm_x, norm_y):
    image = image.copy()
    draw = ImageDraw.Draw(image)
    w, h = image.size
    x = int(norm_x * w)
    y = int(norm_y * h)

    # Draw red center dot
    dot_radius = max(3, int(0.01 * min(w, h)))
    draw.ellipse(
        (x - dot_radius, y - dot_radius, x + dot_radius, y + dot_radius),
        fill=(0, 255, 0)
    )

    # 16:9 aspect ratio focus box
    box_half_width = int(0.09 * w)    # horizontal radius
    box_half_height = int(box_half_width * 12 / 16)  # 16:9 aspect

    left = x - box_half_width
    right = x + box_half_width
    top = y - box_half_height
    bottom = y + box_half_height

    # Gaps in middle of each edge
    gap_x = int(box_half_width * 0.3)
    gap_y = int(box_half_height * 0.3)

    # Top edge (left and right segments)
    draw.line([(left, top), (x - gap_x, top)], fill=(0, 255, 0), width=2)
    draw.line([(x + gap_x, top), (right, top)], fill=(0, 255, 0), width=2)

    # Bottom edge
    draw.line([(left, bottom), (x - gap_x, bottom)], fill=(0, 255, 0), width=2)
    draw.line([(x + gap_x, bottom), (right, bottom)], fill=(0, 255, 0), width=2)

    # Left edge
    draw.line([(left, top), (left, y - gap_y)], fill=(0, 255, 0), width=2)
    draw.line([(left, y + gap_y), (left, bottom)], fill=(0, 255, 0), width=2)

    # Right edge
    draw.line([(right, top), (right, y - gap_y)], fill=(0, 255, 0), width=2)
    draw.line([(right, y + gap_y), (right, bottom)], fill=(0, 255, 0), width=2)

    return image



# Handle clicks: update image with point + show coordinates
def on_image_click(evt: gr.SelectData, image: Image.Image):
    image = image['composite']
    x, y = evt.index
    w, h = image.size
    norm_x = x / w
    norm_y = y / h
    image_with_box = draw_focus_box_on_image(image, norm_x, norm_y)
    coord_str = f"Normalized click: ({norm_y:.3f}, {norm_x:.3f})"
    return image_with_box, coord_str, (norm_y, norm_x)

def generate_depth(image):
    image = image['composite'] 
    if image is None:
        return None
    depth = pipe(image)["depth"]
    return depth

def f(rgb_image: Image.Image, depth_image: Image.Image, K_scalar, x, y) -> Image.Image:
    # rgb_np = np.array(rgb_image)
    # depth_resized = depth_image.resize(rgb_image.size)
    # depth_np = np.array(depth_resized).astype(np.uint8)
    # if len(depth_np.shape) == 3:
    #     depth_np = depth_np[:, :, 0]
    # rgb_np[:, :, 2] = depth_np
    # return Image.fromarray(rgb_np)
    # return rgb_image
    result = diffcamera_inference(rgb_image, depth_image, K_scalar, x, y, pipeline_args, args, weight_dtype,  camera_projector, camera_projector_768)
    print(type(result))
    return result

def update_image(img):
    print('loading image')
    image = img['composite']
    if image.size[0] != image.size[1]:
        image = image.resize((image.size[0], image.size[0]))
    image = draw_focus_box_on_image(image, 0.5, 0.5)
    return image, f"Normalized click: ({0.5:.3f}, {0.5:.3f})", (0.5, 0.5), None  # clear depth

def process(image, depth, coords, strength):
    image = image['composite']
    # return image, depth
    generated = False
    if depth is None:
        print('Calculating depth...')
        depth = pipe(image)["depth"]
        generated = True
    print(f"Processing with coords={coords}, strength={strength}")
    output = f(image, depth, strength, coords[0], coords[1])
    # return output, depth if generated else gr.update()
    return output, depth

with gr.Blocks() as demo:
    with gr.Row():
        image_editor = gr.ImageEditor(
            # height=500,
            label="Editable Input Image",
            type="pil",
            image_mode = "RGB",
            # canvas_size=(1024,1024),
            crop_size="1:1",
        )
        input_image = gr.Image(
            label="Input Image (Click to set the focus point)",
            type="pil",
            interactive=False,
        )
        depth_image = gr.Image(label="Depth Map (optional or auto-generated)", type="pil")

    stored_coords = gr.State((0.5, 0.5))

    click_coords = gr.Textbox(label="Focus Point Coordinates", interactive=False)

    strength_slider = gr.Slider(minimum=0, maximum=20, step=1, label="Bokeh Level", value=10)

    image_editor.change(
        fn=update_image,
        inputs=image_editor,
        outputs=[input_image, click_coords, stored_coords, depth_image]
    )

    input_image.select(
        fn=on_image_click,
        inputs=image_editor,
        outputs=[input_image, click_coords, stored_coords]
    )

    with gr.Row():
        generate_button = gr.Button("Generate Depth from Input Image")
        run_button = gr.Button("Refocus")

    output_image = gr.Image(label="Output Image")

    generate_button.click(fn=generate_depth, inputs=image_editor, outputs=depth_image)


    run_button.click(
        fn=process,
        inputs=[image_editor, depth_image, stored_coords, strength_slider],
        outputs=[output_image, depth_image]
    )


demo.launch(server_port=8888)