import spaces import os import gradio as gr import torch import numpy as np import cv2 import safetensors from PIL import Image, ImageDraw from diffusers import AutoencoderKL from diffusers.utils import load_image, check_min_version from controlnet_flux import FluxControlNetModel from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline from transformers import AutoProcessor, pipeline, AutoModelForMaskGeneration from diffusers.models.attention_processor import Attention from dataclasses import dataclass from typing import Any, List, Dict, Optional, Union, Tuple from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, FluxTransformer2DModel, FluxPipeline from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel # Ensure that the minimal version of diffusers is installed check_min_version("0.30.2") HF_TOKEN = os.getenv("HF_TOKEN") os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1' dtype = torch.bfloat16 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype, use_safetensors=True, token=HF_TOKEN ).to("cuda") # quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True) # transformer_8bit = FluxTransformer2DModel.from_pretrained( # "black-forest-labs/FLUX.1-dev", # subfolder="transformer", # quantization_config=quant_config, # torch_dtype=dtype, # token=HF_TOKEN # ) # Quantize the text encoder to 8-bit precision quant_config = BitsAndBytesConfig(load_in_8bit=True) text_encoder_8bit = T5EncoderModel.from_pretrained( "black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2", quantization_config=quant_config, torch_dtype=torch.float16, token=HF_TOKEN ) # # Load necessary models and processors # controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16) # pipe = FluxControlNetInpaintingPipeline.from_pretrained( # "LPX55/FLUX.1-merged_uncensored", # vae=good_vae, # # transformer=transformer_8bit, # controlnet=controlnet, # torch_dtype=dtype, # use_safetensors=True, # token=HF_TOKEN # ).to("cuda") 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("cuda") pipe.transformer.to(torch.bfloat16) pipe.controlnet.to(torch.bfloat16) pipe.text_encoder_2 = text_encoder_8bit base_attn_procs = pipe.transformer.attn_processors.copy() detector_id = "IDEA-Research/grounding-dino-tiny" segmenter_id = "facebook/sam-vit-base" segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).cuda() segment_processor = AutoProcessor.from_pretrained(segmenter_id) object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=torch.device("cuda")) @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'])) 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) polygon = largest_contour.reshape(-1, 2).tolist() return polygon def polygon_to_mask(polygon: List[Tuple[int, int]], image_shape: Tuple[int, int]) -> np.ndarray: mask = np.zeros(image_shape, dtype=np.uint8) 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 = [] for result in results: xyxy = result.box.xyxy boxes.append(xyxy) return [boxes] def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]: masks = masks.cpu().float() masks = masks.permute(0, 2, 3, 1) masks = masks.mean(axis=-1) masks = (masks > 0).int() masks = masks.numpy().astype(np.uint8) masks = list(masks) if polygon_refinement: for idx, mask in enumerate(masks): shape = mask.shape polygon = mask_to_polygon(mask) mask = polygon_to_mask(polygon, shape) masks[idx] = mask return masks def detect( object_detector, image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None ) -> List[Dict[str, Any]]: device = "cuda" if torch.cuda.is_available() else "cpu" detector_id = detector_id if detector_id is not None else detector_id labels = [label if label.endswith(".") else label+"." for label in labels] results = object_detector(image, candidate_labels=labels, threshold=threshold) results = [DetectionResult.from_dict(result) for result in results] return results def segment( segmentator, processor, image_tensor: torch.Tensor, detection_results: List[Dict[str, Any]], polygon_refinement: bool = False ) -> List[DetectionResult]: device = image_tensor.device boxes = get_boxes(detection_results) # Convert image tensor to float32 for processing image_tensor_float32 = image_tensor.to(torch.float32) inputs = processor(images=image_tensor_float32, input_boxes=boxes, return_tensors="pt", torch_dtype=torch.float32) # Process inputs and get outputs outputs = segmentator(**inputs) # Convert masks to bfloat16 if needed masks = outputs.pred_masks.to(torch.bfloat16) masks = processor.post_process_masks( masks=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: Union[Image.Image, str], labels: List[str], threshold: float = 0.3, polygon_refinement: bool = False, detector_id: Optional[str] = None, segmenter_id: Optional[str] = None ) -> Tuple[np.ndarray, List[DetectionResult]]: if isinstance(image, str): image = load_image(image) # Convert image to tensor and to float32 for processing image_tensor = torch.tensor(np.array(image), dtype=torch.float32, device="cuda").permute(2, 0, 1).unsqueeze(0) / 255.0 detections = detect(detect_pipeline, image, labels, threshold, detector_id) detections = segment(segmentator, segment_processor, image_tensor, detections, polygon_refinement) # Convert image tensor back to numpy array for return image_array = image_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255 image_array = image_array.astype(np.uint8) return image_array, detections class CustomFluxAttnProcessor2_0: def __init__(self, height=44, width=88, attn_enforce=1.0): if not hasattr(torch.nn.functional, "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 = key.shape[-1] head_dim = inner_dim // attn.heads query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) 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_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( batch_size, -1, attn.heads, head_dim ).transpose(1, 2) encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( batch_size, -1, attn.heads, head_dim ).transpose(1, 2) encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( batch_size, -1, attn.heads, head_dim ).transpose(1, 2) if attn.norm_added_q is not None: encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) if attn.norm_added_k is not None: encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) 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:] img_attn_probs = img_attn_probs.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 = torch.nn.functional.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) hidden_states = hidden_states.to(query.dtype) if encoder_hidden_states is not None: encoder_hidden_states, hidden_states = ( hidden_states[:, : encoder_hidden_states.shape[1]], hidden_states[:, encoder_hidden_states.shape[1] :], ) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) encoder_hidden_states = attn.to_add_out(encoder_hidden_states) return hidden_states, encoder_hidden_states else: return hidden_states def segment_image(image, object_name): image_array, detections = grounded_segmentation( object_detector, segmentator, segment_processor, image=image, labels=object_name, threshold=0.3, polygon_refinement=True, ) segment_result = image_array * np.expand_dims((255 - detections[0].mask) / 255, axis=-1) segmented_image = Image.fromarray(segment_result.astype(np.uint8)) return segmented_image def make_diptych(image): ref_image = np.array(image) ref_image = np.concatenate([ref_image, np.zeros_like(ref_image)], axis=1) ref_image = Image.fromarray(ref_image) return ref_image @spaces.GPU() def inpaint_image(image, prompt, object_name): width = 512 height = 512 size = (width * 2, height) diptych_text_prompt = f"A diptych with two side-by-side images of same {object_name}. On the left, a photo of {object_name}. On the right, {prompt}" reference_image = image.resize((width, height)).convert("RGB") segmented_image = segment_image(reference_image, object_name) mask_image = np.concatenate([np.zeros((height, width, 3)), np.ones((height, width, 3))*255], axis=1) mask_image = Image.fromarray(mask_image.astype(np.uint8)) diptych_image_prompt = make_diptych(segmented_image) base_attn_procs = pipe.transformer.attn_processors.copy() new_attn_procs = base_attn_procs.copy() for i, (k, v) in enumerate(new_attn_procs.items()): new_attn_procs[k] = CustomFluxAttnProcessor2_0(height=height // 16, width=width // 16 * 2, attn_enforce=1.3) pipe.transformer.set_attn_processor(new_attn_procs) generator = torch.Generator(device="cuda").manual_seed(42) with torch.no_grad(): result = pipe( prompt=diptych_text_prompt, height=size[1], width=size[0], control_image=diptych_image_prompt, control_mask=mask_image, num_inference_steps=20, generator=generator, controlnet_conditioning_scale=0.95, guidance_scale=3.5, negative_prompt="", true_guidance_scale=3.5 ).images[0] result = result.crop((width, 0, width*2, height)) torch.cuda.empty_cache() return result, diptych_image_prompt # Create Gradio interface iface = gr.Interface( fn=inpaint_image, inputs=[ gr.Image(type="pil", label="Upload Image"), gr.Textbox(lines=3, value="replicate this {subject_name} exactly but as a photo of the {subject_name} surfing on the beach", label="Prompt"), gr.Textbox(lines=1, value="bear plushie", label="Subject Name") ], outputs=[ gr.Image(type="pil", label="Inpainted Image"), gr.Image(type="pil", label="Diptych Image") ], title="FLUX Inpainting with Diptych Prompting", description="Upload an image, specify a prompt, and provide the subject name. The app will automatically generate the inpainted image." ) # Launch the app iface.launch()