import spaces import gradio as gr import torch import torch.nn.functional as F import numpy as np from PIL import Image import cv2 import os from diffusers.utils import load_image, check_min_version from controlnet_flux import FluxControlNetModel from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline from diffusers.models.attention_processor import Attention from transformers import AutoProcessor, AutoModelForMaskGeneration, pipeline from dataclasses import dataclass from typing import Any, List, Dict, Optional, Union, Tuple from huggingface_hub import hf_hub_download import random device = "cuda" if torch.cuda.is_available() else "cpu" # --- Helper Dataclasses (Identical to previous version) --- @dataclass class BoundingBox: xmin: int ymin: int xmax: int ymax: int @property def xyxy(self) -> List[float]: return [self.xmin, self.ymin, self.xmax, self.ymax] @dataclass class DetectionResult: score: float label: str box: BoundingBox mask: Optional[np.array] = None @classmethod def from_dict(cls, detection_dict: Dict) -> 'DetectionResult': return cls(score=detection_dict['score'], label=detection_dict['label'], box=BoundingBox(xmin=detection_dict['box']['xmin'], ymin=detection_dict['box']['ymin'], xmax=detection_dict['box']['xmax'], ymax=detection_dict['box']['ymax'])) # --- Helper Functions (Identical to previous version) --- def mask_to_polygon(mask: np.ndarray) -> List[List[int]]: contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return [] largest_contour = max(contours, key=cv2.contourArea) return largest_contour.reshape(-1, 2).tolist() def polygon_to_mask(polygon: List[Tuple[int, int]], image_shape: Tuple[int, int]) -> np.ndarray: mask = np.zeros(image_shape, dtype=np.uint8) if not polygon: return mask pts = np.array(polygon, dtype=np.int32) cv2.fillPoly(mask, [pts], color=(255,)) return mask def get_boxes(results: List[DetectionResult]) -> List[List[List[float]]]: boxes = [result.box.xyxy for result in results] return [boxes] def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]: masks = masks.cpu().float().permute(0, 2, 3, 1).mean(axis=-1) masks = (masks > 0).int().numpy().astype(np.uint8) masks = list(masks) if polygon_refinement: for idx, mask in enumerate(masks): shape = mask.shape polygon = mask_to_polygon(mask) refined_mask = polygon_to_mask(polygon, shape) masks[idx] = refined_mask return masks def detect( object_detector, image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None ) -> List[DetectionResult]: labels = [label if label.endswith(".") else label + "." for label in labels] results = object_detector(image, candidate_labels=labels, threshold=threshold) return [DetectionResult.from_dict(result) for result in results] def segment( segmentator, processor, image: Image.Image, detection_results: List[DetectionResult], polygon_refinement: bool = False ) -> List[DetectionResult]: if not detection_results: return [] boxes = get_boxes(detection_results) inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to(device) with torch.no_grad(): outputs = segmentator(**inputs) masks = processor.post_process_masks( masks=outputs.pred_masks, original_sizes=inputs.original_sizes, reshaped_input_sizes=inputs.reshaped_input_sizes )[0] masks = refine_masks(masks, polygon_refinement) for detection_result, mask in zip(detection_results, masks): detection_result.mask = mask return detection_results def grounded_segmentation( detect_pipeline, segmentator, segment_processor, image: Image.Image, labels: List[str], ) -> Tuple[np.ndarray, List[DetectionResult]]: detections = detect(detect_pipeline, image, labels, threshold=0.3) detections = segment(segmentator, segment_processor, image, detections, polygon_refinement=True) return np.array(image), detections def segment_image(image, object_name, detector, segmentator, seg_processor): """ Segments a specific object from an image and returns the segmented object on a white background. Args: image (PIL.Image.Image): The input image. object_name (str): The name of the object to segment. detector: The object detection pipeline. segmentator: The mask generation model. seg_processor: The processor for the mask generation model. Returns: PIL.Image.Image: The image with the segmented object on a white background. Raises: gr.Error: If the object cannot be segmented. """ image_array, detections = grounded_segmentation(detector, segmentator, seg_processor, image, [object_name]) if not detections or detections[0].mask is None: raise gr.Error(f"Could not segment the subject '{object_name}' in the image. Please try a clearer image or a more specific subject name.") mask_expanded = np.expand_dims(detections[0].mask / 255, axis=-1) segment_result = image_array * mask_expanded + np.ones_like(image_array) * (1 - mask_expanded) * 255 return Image.fromarray(segment_result.astype(np.uint8)) def make_diptych(image): """ Creates a diptych image by concatenating the input image with a black image of the same size. Args: image (PIL.Image.Image): The input image. Returns: PIL.Image.Image: The diptych image. """ ref_image_np = np.array(image) diptych_np = np.concatenate([ref_image_np, np.zeros_like(ref_image_np)], axis=1) return Image.fromarray(diptych_np) # --- Custom Attention Processor (Identical to previous version) --- class CustomFluxAttnProcessor2_0: def __init__(self, height=44, width=88, attn_enforce=1.0): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") self.height = height self.width = width self.num_pixels = height * width self.step = 0 self.attn_enforce = attn_enforce def __call__( self, attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, attention_mask: Optional[torch.FloatTensor] = None, image_rotary_emb: Optional[torch.Tensor] = None, ) -> torch.FloatTensor: self.step += 1 batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) inner_dim, head_dim = key.shape[-1], key.shape[-1] // attn.heads query, key, value = [x.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) for x in [query, key, value]] if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) if encoder_hidden_states is not None: encoder_q = attn.add_q_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) encoder_k = attn.add_k_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) encoder_v = attn.add_v_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) if attn.norm_added_q is not None: encoder_q = attn.norm_added_q(encoder_q) if attn.norm_added_k is not None: encoder_k = attn.norm_added_k(encoder_k) query, key, value = [torch.cat([e, x], dim=2) for e, x in zip([encoder_q, encoder_k, encoder_v], [query, key, value])] if image_rotary_emb is not None: from diffusers.models.embeddings import apply_rotary_emb query = apply_rotary_emb(query, image_rotary_emb) key = apply_rotary_emb(key, image_rotary_emb) if self.attn_enforce != 1.0: attn_probs = (torch.einsum('bhqd,bhkd->bhqk', query, key) * attn.scale).softmax(dim=-1) img_attn_probs = attn_probs[:, :, -self.num_pixels:, -self.num_pixels:].reshape((batch_size, attn.heads, self.height, self.width, self.height, self.width)) img_attn_probs[:, :, :, self.width//2:, :, :self.width//2] *= self.attn_enforce img_attn_probs = img_attn_probs.reshape((batch_size, attn.heads, self.num_pixels, self.num_pixels)) attn_probs[:, :, -self.num_pixels:, -self.num_pixels:] = img_attn_probs hidden_states = torch.einsum('bhqk,bhkd->bhqd', attn_probs, value) else: hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim).to(query.dtype) if encoder_hidden_states is not None: encoder_hs, hs = hidden_states[:, : encoder_hidden_states.shape[1]], hidden_states[:, encoder_hidden_states.shape[1] :] hs = attn.to_out[0](hs) hs = attn.to_out[1](hs) encoder_hs = attn.to_add_out(encoder_hs) return hs, encoder_hs else: return hidden_states # --- Model Loading (executed once at startup) --- print("--- Loading Models: This may take a few minutes and requires >40GB VRAM ---") controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16) pipe = FluxControlNetInpaintingPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", controlnet=controlnet, torch_dtype=torch.bfloat16).to(device) pipe.transformer.to(torch.bfloat16) pipe.controlnet.to(torch.bfloat16) base_attn_procs = pipe.transformer.attn_processors.copy() print("Loading segmentation models...") detector_id, segmenter_id = "IDEA-Research/grounding-dino-tiny", "facebook/sam-vit-base" segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(device) segment_processor = AutoProcessor.from_pretrained(segmenter_id) object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=device) print("--- All models loaded successfully! ---") def get_duration( input_image: Image.Image, subject_name: str, do_segmentation: bool, full_prompt: str, attn_enforce: float, ctrl_scale: float, width: int, height: int, pixel_offset: int, num_steps: int, guidance: float, real_guidance: float, seed: int, randomize_seed: bool, progress=gr.Progress(track_tqdm=True) ): """ Calculates the estimated duration for the Spaces GPU based on image dimensions. Args: input_image (PIL.Image.Image): The input reference image. subject_name (str): Name of the subject for segmentation. do_segmentation (bool): Whether to perform segmentation. full_prompt (str): The full text prompt. attn_enforce (float): Attention enforcement value. ctrl_scale (float): ControlNet conditioning scale. width (int): Target width of the generated image. height (int): Target height of the generated image. pixel_offset (int): Padding offset in pixels. num_steps (int): Number of inference steps. guidance (float): Distilled guidance scale. real_guidance (float): Real guidance scale. seed (int): Random seed. randomize_seed (bool): Whether to randomize the seed. progress (gr.Progress): Gradio progress tracker. Returns: int: Estimated duration in seconds. """ if width > 768 or height > 768: return 210 else: return 120 @spaces.GPU(duration=get_duration) def run_diptych_prompting( input_image: Image.Image, subject_name: str, do_segmentation: bool, full_prompt: str, attn_enforce: float, ctrl_scale: float, width: int, height: int, pixel_offset: int, num_steps: int, guidance: float, real_guidance: float, seed: int, randomize_seed: bool, progress=gr.Progress(track_tqdm=True) ): """ Runs the diptych prompting image generation process. Args: input_image (PIL.Image.Image): The input reference image. subject_name (str): The name of the subject for segmentation (if `do_segmentation` is True). do_segmentation (bool): If True, the subject will be segmented from the reference image. full_prompt (str): The complete text prompt used for image generation. attn_enforce (float): Controls the attention enforcement in the custom attention processor. ctrl_scale (float): The conditioning scale for ControlNet. width (int): The desired width of the final generated image. height (int): The desired height of the final generated image. pixel_offset (int): Padding added around the image during diptych creation. num_steps (int): The number of inference steps for the diffusion process. guidance (float): The distilled guidance scale for the diffusion process. real_guidance (float): The real guidance scale for the diffusion process. seed (int): The random seed for reproducibility. randomize_seed (bool): If True, a random seed will be used instead of the provided `seed`. progress (gr.Progress): Gradio progress tracker to update UI during execution. Returns: tuple: A tuple containing: - PIL.Image.Image: The final generated image. - PIL.Image.Image: The processed reference image (left panel of the diptych). - PIL.Image.Image: The full diptych image generated by the pipeline. - str: The final prompt used. - int: The actual seed used for generation. Raises: gr.Error: If a reference image is not uploaded, prompts are empty, or segmentation fails. """ if randomize_seed: actual_seed = random.randint(0, 9223372036854775807) else: actual_seed = seed if input_image is None: raise gr.Error("Please upload a reference image.") if not full_prompt: raise gr.Error("Full Prompt is empty. Please fill out the prompt fields.") # 1. Prepare dimensions and reference image padded_width = width + pixel_offset * 2 padded_height = height + pixel_offset * 2 diptych_size = (padded_width * 2, padded_height) reference_image = input_image.resize((padded_width, padded_height)).convert("RGB") # 2. Process reference image based on segmentation flag progress(0, desc="Preparing reference image...") if do_segmentation: if not subject_name: raise gr.Error("Subject Name is required when 'Do Segmentation' is checked.") progress(0.05, desc="Segmenting reference image...") processed_image = segment_image(reference_image, subject_name, object_detector, segmentator, segment_processor) else: processed_image = reference_image # 3. Create diptych and mask progress(0.2, desc="Creating diptych and mask...") mask_image = np.concatenate([np.zeros((padded_height, padded_width, 3)), np.ones((padded_height, padded_width, 3)) * 255], axis=1) mask_image = Image.fromarray(mask_image.astype(np.uint8)) diptych_image_prompt = make_diptych(processed_image) # 4. Setup Attention Processor progress(0.3, desc="Setting up attention processors...") new_attn_procs = base_attn_procs.copy() for k in new_attn_procs: new_attn_procs[k] = CustomFluxAttnProcessor2_0(height=padded_height // 16, width=padded_width * 2 // 16, attn_enforce=attn_enforce) pipe.transformer.set_attn_processor(new_attn_procs) # 5. Run Inference progress(0.4, desc="Running diffusion process...") generator = torch.Generator(device="cuda").manual_seed(actual_seed) full_diptych_result = pipe( prompt=full_prompt, height=diptych_size[1], width=diptych_size[0], control_image=diptych_image_prompt, control_mask=mask_image, num_inference_steps=num_steps, generator=generator, controlnet_conditioning_scale=ctrl_scale, guidance_scale=guidance, negative_prompt="", true_guidance_scale=real_guidance ).images[0] # 6. Final cropping progress(0.95, desc="Finalizing image...") final_image = full_diptych_result.crop((padded_width, 0, padded_width * 2, padded_height)) final_image = final_image.crop((pixel_offset, pixel_offset, padded_width - pixel_offset, padded_height - pixel_offset)) # 7. Return all outputs return final_image, processed_image, full_diptych_result, full_prompt, actual_seed # --- Gradio UI Definition --- css = ''' .gradio-container{max-width: 960px;margin: 0 auto} ''' with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo: gr.Markdown( """ # Diptych Prompting: Zero-Shot Subject-Driven & Style-Driven Image Generation ### Demo for the paper "[Large-Scale Text-to-Image Model with Inpainting is a Zero-Shot Subject-Driven Image Generator](https://diptychprompting.github.io/)" """ ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(type="pil", label="Reference Image") with gr.Group() as subject_driven_group: subject_name = gr.Textbox(label="Subject Name", placeholder="e.g., a plush bear") target_prompt = gr.Textbox(label="Target Prompt", placeholder="e.g., riding a skateboard on the moon") run_button = gr.Button("Generate Image", variant="primary") with gr.Accordion("Advanced Settings", open=False): mode = gr.Radio(["Subject-Driven", "Style-Driven (unstable)"], label="Generation Mode", value="Subject-Driven") with gr.Group(visible=False) as style_driven_group: original_style_description = gr.Textbox(label="Original Image Description", placeholder="e.g., in watercolor painting style") do_segmentation = gr.Checkbox(label="Do Segmentation", value=True) attn_enforce = gr.Slider(minimum=1.0, maximum=2.0, value=1.3, step=0.05, label="Attention Enforcement") full_prompt = gr.Textbox(label="Full Prompt (Auto-generated, editable)", lines=3) ctrl_scale = gr.Slider(minimum=0.5, maximum=1.0, value=0.95, step=0.01, label="ControlNet Scale") num_steps = gr.Slider(minimum=20, maximum=50, value=28, step=1, label="Inference Steps") guidance = gr.Slider(minimum=1.0, maximum=10.0, value=3.5, step=0.1, label="Distilled Guidance Scale") real_guidance = gr.Slider(minimum=1.0, maximum=10.0, value=4.5, step=0.1, label="Real Guidance Scale") width = gr.Slider(minimum=512, maximum=1024, value=768, step=64, label="Image Width") height = gr.Slider(minimum=512, maximum=1024, value=768, step=64, label="Image Height") pixel_offset = gr.Slider(minimum=0, maximum=32, value=8, step=1, label="Padding (Pixel Offset)") seed = gr.Slider(minimum=0, maximum=9223372036854775807, value=42, step=1, label="Seed") randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) with gr.Column(scale=1): output_image = gr.Image(type="pil", label="Generated Image") with gr.Accordion("Other Outputs", open=False) as other_outputs_accordion: processed_ref_image = gr.Image(label="Processed Reference (Left Panel)") full_diptych_image = gr.Image(label="Full Diptych Output") final_prompt_used = gr.Textbox(label="Final Prompt Used") # --- UI Event Handlers --- def toggle_mode_visibility(mode_choice): """ Hides/shows the relevant input textboxes based on the selected mode. Args: mode_choice (str): The selected generation mode ("Subject-Driven" or "Style-Driven"). Returns: tuple: Gradio update objects for `subject_driven_group` and `style_driven_group` visibility. """ if mode_choice == "Subject-Driven": return gr.update(visible=True), gr.update(visible=False) else: return gr.update(visible=False), gr.update(visible=True) def update_derived_fields(mode_choice, subject, style_desc, target): """ Updates the full prompt and segmentation checkbox based on other inputs. Args: mode_choice (str): The selected generation mode ("Subject-Driven" or "Style-Driven"). subject (str): The subject name (relevant for "Subject-Driven" mode). style_desc (str): The original style description (relevant for "Style-Driven" mode). target (str): The target prompt. Returns: tuple: Gradio update objects for `full_prompt` value and `do_segmentation` checkbox value. """ if mode_choice == "Subject-Driven": prompt = f"A diptych with two side-by-side images of same {subject}. On the left, a photo of {subject}. On the right, replicate this {subject} exactly but as {target}" return gr.update(value=prompt), gr.update(value=True) else: # Style-Driven prompt = f"A diptych with two side-by-side images of same style. On the left, {style_desc}. On the right, replicate this style exactly but as {target}" return gr.update(value=prompt), gr.update(value=False) # --- UI Connections --- # When mode changes, toggle visibility of the specific prompt fields mode.change( fn=toggle_mode_visibility, inputs=mode, outputs=[subject_driven_group, style_driven_group], queue=False ) # A list of all inputs that affect the full prompt or segmentation checkbox prompt_component_inputs = [mode, subject_name, original_style_description, target_prompt] # A list of the UI elements that are derived from the above inputs derived_outputs = [full_prompt, do_segmentation] # When any prompt component changes, update the derived fields for component in prompt_component_inputs: component.change(update_derived_fields, inputs=prompt_component_inputs, outputs=derived_outputs, queue=False, show_progress="hidden") run_button.click( fn=run_diptych_prompting, inputs=[ input_image, subject_name, do_segmentation, full_prompt, attn_enforce, ctrl_scale, width, height, pixel_offset, num_steps, guidance, real_guidance, seed, randomize_seed ], outputs=[output_image, processed_ref_image, full_diptych_image, final_prompt_used, seed] ) def run_subject_driven_example(input_image, subject_name, target_prompt): """ Helper function to run an example for the subject-driven mode. Args: input_image (PIL.Image.Image): The input reference image for the example. subject_name (str): The subject name for the example. target_prompt (str): The target prompt for the example. Returns: tuple: The outputs from `run_diptych_prompting`. """ # Construct the full prompt for subject-driven mode full_prompt = f"A diptych with two side-by-side images of same {subject_name}. On the left, a photo of {subject_name}. On the right, replicate this {subject_name} exactly but as {target_prompt}" # Call the main function with all arguments, using defaults for subject-driven mode return run_diptych_prompting( input_image=input_image, subject_name=subject_name, do_segmentation=True, full_prompt=full_prompt, attn_enforce=1.3, ctrl_scale=0.95, width=768, height=768, pixel_offset=8, num_steps=28, guidance=3.5, real_guidance=4.5, seed=42, randomize_seed=False, ) gr.Examples( examples=[ ["./assets/cat_squished.png", "a cat toy", "a cat toy riding a skate"], ["./assets/hf.png", "hugging face logo", "a hugging face logo on a hat"], ["./assets/bear_plushie.jpg", "a bear plushie", "a bear plushie drinking bubble tea"] ], inputs=[input_image, subject_name, target_prompt], outputs=[output_image, processed_ref_image, full_diptych_image, final_prompt_used, seed], fn=run_subject_driven_example, cache_examples="lazy" ) if __name__ == "__main__": demo.launch(share=True, debug=True, mcp_server=True)