Create app.py

app.py

@@ -0,0 +1,340 @@
+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
+
+# --- Constants and Setup ---
+# Ensure all required modules are available
+check_min_version("0.29.0.dev0")
+
+# Set a seed for reproducibility. The original script uses a fixed seed.
+generator = torch.Generator(device="cuda").manual_seed(42)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# --- Helper Dataclasses (Identical to diptych_prompting_inference.py) ---
+@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 diptych_prompting_inference.py) ---
+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):
+    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):
+    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 (EXACTLY as in diptych_prompting_inference.py) ---
+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! ---")
+
+
+# --- Main Inference Function for Gradio ---
+def run_diptych_prompting(
+    input_image: Image.Image,
+    subject_name: str,
+    target_prompt: str,
+    attn_enforce: float,
+    ctrl_scale: float,
+    width: int,
+    height: int,
+    pixel_offset: int,
+    num_steps: int,
+    guidance: float,
+    progress=gr.Progress(track_tqdm=True)
+):
+    if input_image is None: raise gr.Error("Please upload a reference image.")
+    if not subject_name: raise gr.Error("Please provide the subject's name (e.g., 'a red car').")
+    if not target_prompt: raise gr.Error("Please provide a target prompt.")
+
+    # 1. Prepare dimensions (logic from original script's main block)
+    padded_width = width + pixel_offset * 2
+    padded_height = height + pixel_offset * 2
+    diptych_size = (padded_width * 2, padded_height)
+
+    # 2. Prepare prompts and images
+    progress(0, desc="Resizing and segmenting reference image...")
+    base_prompt = f"a photo of {subject_name}"
+    diptych_text_prompt = f"A diptych with two side-by-side images of same {subject_name}. On the left, {base_prompt}. On the right, replicate this {subject_name} exactly but as {target_prompt}"
+    
+    reference_image = input_image.resize((padded_width, padded_height)).convert("RGB")
+    segmented_image = segment_image(reference_image, subject_name, object_detector, segmentator, segment_processor)
+
+    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(segmented_image)
+
+    # 3. Setup Attention Processor (logic from original script's main block)
+    progress(0.3, desc="Setting up attention processors...")
+    new_attn_procs = base_attn_procs.copy()
+    for k in new_attn_procs:
+        # Use full diptych dimensions for the attention processor
+        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)
+
+    # 4. Run Inference (using parameters identical to the original script)
+    progress(0.4, desc="Running diffusion process...")
+    result = pipe(
+        prompt=diptych_text_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,          # This is used for guidance embeds if enabled
+        negative_prompt="",
+        true_guidance_scale=guidance      # **CRITICAL FIX**: This matches the original script's CFG scale
+    ).images[0]
+
+    # 5. Final cropping (logic from original script's main block)
+    progress(0.95, desc="Finalizing image...")
+    # Crop the right panel
+    result = result.crop((padded_width, 0, padded_width * 2, padded_height))
+    # Crop the pixel offset padding
+    result = result.crop((pixel_offset, pixel_offset, padded_width - pixel_offset, padded_height - pixel_offset))
+
+    return result
+
+
+# --- Gradio UI Definition ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Diptych Prompting: Zero-Shot Subject-Driven Image Generation
+        ### Official Gradio Demo for the paper "[Large-Scale Text-to-Image Model with Inpainting is a Zero-Shot Subject-Driven Image Generator](https://diptychprompting.github.io/)"
+        **Instructions:**
+        1. Upload a clear image of a subject.
+        2. Describe the subject in the "Subject Name" box (e.g., 'a plush bear', 'a red sports car').
+        3. Describe the new scene for your subject in the "Target Prompt" box.
+        
+        > **Note:** This demo runs on a powerful GPU and requires over **40GB of VRAM**. Inference may take 1-2 minutes.
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="1. Reference Image")
+            subject_name = gr.Textbox(label="2. Subject Name", placeholder="e.g., a plush bear")
+            target_prompt = gr.Textbox(label="3. 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):
+                attn_enforce = gr.Slider(minimum=1.0, maximum=2.0, value=1.3, step=0.05, label="Attention Enforcement")
+                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=30, step=1, label="Inference Steps")
+                guidance = gr.Slider(minimum=1.0, maximum=10.0, value=3.5, step=0.1, label="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)")
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="pil", label="Generated Image")
+
+    gr.Examples(
+        examples=[
+            ["./assets/bear_plushie.jpg", "a bear plushie", "a bear plushie riding a skateboard"],
+            ["./assets/corgi.jpg", "a corgi dog", "a corgi dog wearing a superhero cape and flying"],
+            ["./assets/teapot.png", "a blue and white teapot", "a blue and white teapot in a field of flowers"],
+        ],
+        inputs=[input_image, subject_name, target_prompt],
+        outputs=output_image,
+        fn=run_diptych_prompting,
+        cache_examples=False,
+    )
+
+    run_button.click(
+        fn=run_diptych_prompting,
+        inputs=[input_image, subject_name, target_prompt, attn_enforce, ctrl_scale, width, height, pixel_offset, num_steps, guidance],
+        outputs=output_image
+    )
+
+if __name__ == "__main__":
+    if not os.path.exists("./assets"):
+        os.makedirs("./assets")
+        print("Created './assets' directory. Please add example images like 'bear_plushie.jpg' there for the examples to work.")
+        
+    demo.launch(share=True)