Initial commit

README.md

@@ -0,0 +1,13 @@
+---
+title: NativDemo
+emoji: 🏆
+colorFrom: gray
+colorTo: blue
+sdk: gradio
+sdk_version: 5.6.0
+app_file: app.py
+pinned: false
+short_description: In place text localization
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,626 @@
+import spaces
+import os
+import gradio as gr
+import easyocr
+import numpy as np
+import cv2
+import base64
+import torch
+from shapely import Polygon
+from ultralytics import YOLO
+
+from io import BytesIO
+from openai import OpenAI
+from PIL import Image, ImageDraw, ImageFont
+
+from diffusers.utils import load_image, check_min_version
+from controlnet_flux import FluxControlNetModel
+from transformer_flux import FluxTransformer2DModel
+from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
+
+import huggingface_hub
+huggingface_hub.login(os.getenv('HF_TOKEN_FLUX'))
+
+bubble_detection_model = YOLO("speech_bubble_model.pt")
+
+language_to_ocr = {
+    'Simplified Chinese': 'ch_sim',
+    'Traditional Chinese': 'ch_tra',
+    'Korean': 'ko',
+    'Japanese': 'ja',
+    'English': 'en',
+}
+
+OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
+
+MARKDOWN = """
+# Made by Nativ
+"""
+
+check_min_version("0.30.2")
+transformer = FluxTransformer2DModel.from_pretrained(
+        "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dytpe=torch.bfloat16
+    )
+
+cuda_device =torch.device("cuda")
+# Build pipeline
+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,
+    transformer=transformer,
+    torch_dtype=torch.bfloat16
+).to(cuda_device)
+pipe.transformer.to(torch.bfloat16)
+pipe.controlnet.to(torch.bfloat16)
+
+
+def hex_to_rgba(hex_color):
+    print(hex_color)
+    """Convert hex color to RGBA tuple."""
+    hex_color = hex_color.lstrip('#')  # Remove '#' if present
+    if len(hex_color) == 6:  # Handle `#RRGGBB`
+        r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
+        return r, g, b, 255  # Add full opacity (alpha = 255)
+    elif len(hex_color) == 8:  # Handle `#RRGGBBAA` if alpha is included
+        r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
+        return r, g, b, a
+    else:
+        raise ValueError(f"Invalid hex color format: {hex_color}")
+
+import re
+
+def rgba_to_tuple(rgba_color):
+    """Convert rgba(r, g, b, a) string to an (R, G, B, A) tuple."""
+    # Match the rgba format: rgba(r, g, b, a)
+    match = re.match(r'rgba\(([\d.]+),\s*([\d.]+),\s*([\d.]+),\s*([\d.]+)\)', rgba_color)
+    if not match:
+        raise ValueError(f"Invalid RGBA color format: {rgba_color}")
+    
+    r, g, b, a = map(float, match.groups())
+    r, g, b = int(r), int(g), int(b)
+    a = int(a * 255)  # Scale alpha from [0, 1] to [0, 255]
+    return r, g, b, a
+
+
+def color_to_rgba(color):
+    """Convert a color string (hex or rgba) to an RGBA tuple."""
+    if color.startswith("#"):  # Hex format
+        return hex_to_rgba(color)
+    elif color.startswith("rgba"):  # rgba(r, g, b, a) format
+        return rgba_to_tuple(color)
+    else:
+        raise ValueError(f"Unsupported color format: {color}")
+
+
+def localize_boxes(merged_results, img_boxes, source_language, target_language):
+    # Convert image to base64
+    buffered = BytesIO()
+    img_boxes.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    
+    print(merged_results)
+    
+    prompt = f"""You are an expert translator and localization specialist with deep understanding of both {source_language} and {target_language} cultures.
+
+    Task: Translate the detected text while preserving the cultural context and maintaining visual harmony. Make the results in capital letters.
+
+    Source Text and Coordinates:
+    {merged_results}
+
+    Requirements:
+    1. Maintain the original meaning and tone while adapting to {target_language} cultural context
+    2. Keep translations concise and visually balanced (similar character length when possible)
+    3. Preserve any:
+    - Brand names
+    - Product names
+    - Technical terms
+    - Numbers and units
+    4. Consider the visual context from the provided image
+    5. Use appropriate formality level for {target_language}
+    6. Maintain any special formatting (if present)
+
+    Format your response EXACTLY as a JSON-like list of dictionaries. Keep the box coordinates EXACTLY as they are, do not change them, only translate the text.
+    [{{'box': [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], 'text': 'translated_text'}}]
+
+    Important: Only output the JSON format above, no explanations or additional text."""
+
+    client = OpenAI(api_key=OPENAI_API_KEY)
+    
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{img_str}"
+                        }
+                    }
+                ]
+            }
+        ],
+        max_tokens=1000,
+        temperature=0
+    )
+    
+    try:
+        translation_text = response.choices[0].message.content
+        translation_text = translation_text.replace("```json", "").replace("```", "").strip()
+        translated_results = eval(translation_text)
+        return translated_results
+    except Exception as e:
+        print(f"Error parsing GPT-4o response: {e}")
+        return merged_results  
+    
+def merge_boxes(boxes, image_shape, distance_threshold=10):
+    """Merge boxes that are close to each other and return their associated text"""
+    if not boxes:
+        return []
+    
+    # Extract boxes and create mapping to original data
+    boxes_only = [box[0] for box in boxes]
+    texts = [box[1] for box in boxes]  # Extract the text content
+    
+    # Create a binary mask of all boxes
+    height, width = image_shape[:2]
+    mask = np.zeros((height, width), dtype=np.uint8)
+    
+    # Draw all boxes on mask and create a mapping of pixel positions to box indices
+    box_indices_map = {}  # Will store which original box each pixel belongs to
+    for idx, coords in enumerate(boxes_only):
+        pts = np.array(coords, dtype=np.int32)
+        cv2.fillPoly(mask, [pts], 255)
+        # Store the indices of boxes for each filled pixel
+        y_coords, x_coords = np.where(mask == 255)
+        for y, x in zip(y_coords, x_coords):
+            if (y, x) not in box_indices_map:
+                box_indices_map[(y, x)] = []
+            box_indices_map[(y, x)].append(idx)
+    
+    # Dilate to connect nearby components
+    kernel = np.ones((distance_threshold, distance_threshold), np.uint8)
+    dilated = cv2.dilate(mask, kernel, iterations=1)
+    
+    # Find connected components
+    num_labels, labels = cv2.connectedComponents(dilated)
+    
+    # Create new merged boxes with their associated text
+    merged_results = []
+    for label in range(1, num_labels):  # Skip background (0)
+        points = np.where(labels == label)
+        if len(points[0]):  # If component is not empty
+            y0, x0 = points[0].min(), points[1].min()
+            y1, x1 = points[0].max(), points[1].max()
+            # Add small padding
+            x0 = max(0, x0 - 2)
+            y0 = max(0, y0 - 2)
+            x1 = min(width, x1 + 2)
+            y1 = min(height, y1 + 2)
+            
+            # Find all original boxes that overlap with this merged box
+            box_indices = set()
+            for y in range(y0, y1+1):
+                for x in range(x0, x1+1):
+                    if (y, x) in box_indices_map:
+                        box_indices.update(box_indices_map[(y, x)])
+            
+            # Combine text from all overlapping boxes
+            combined_text = ' '.join([texts[idx] for idx in box_indices])
+            
+            merged_results.append({
+                'box': [[x0, y0], [x1, y0], [x1, y1], [x0, y1]],
+                'text': combined_text
+            })    
+    return merged_results
+
+def is_box_inside_yolo(box, yolo_boxes, overlap_threshold=0.5):
+    """
+    Check if a text box is inside any of the YOLO-detected speech bubbles.
+    box: [[x0,y0], [x1,y0], [x1,y1], [x0,y1]]
+    yolo_boxes: list of YOLO boxes in xywh format
+    overlap_threshold: minimum overlap ratio required to consider the text inside bubble
+    """
+    text_poly = Polygon(box)
+    text_area = text_poly.area
+    
+    for yolo_box in yolo_boxes:
+        x_center, y_center, width, height = yolo_box
+        x1, y1 = x_center - width / 2, y_center - height / 2
+        x2, y2 = x_center + width / 2, y_center + height / 2
+        bubble_box = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
+        bubble_poly = Polygon(bubble_box)
+        
+        # Calculate intersection
+        if text_poly.intersects(bubble_poly):
+            intersection = text_poly.intersection(bubble_poly)
+            overlap_ratio = intersection.area / text_area
+            if overlap_ratio >= overlap_threshold:
+                return True
+    
+    return False
+
+def remove_text_regions(image, boxes, yolo_boxes):
+    """Fill detected text regions with white"""
+    img_removed = image.copy()
+    mask = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    
+    # Fill all detected boxes with white
+    for box in boxes:
+        pts = np.array(box[0], dtype=np.int32)
+        if is_box_inside_yolo(box[0], yolo_boxes):
+            cv2.fillPoly(img_removed, [pts], (255, 255, 255, 255))
+        cv2.fillPoly(mask, [pts], (255, 255, 255, 255))
+    
+    img_removed_rgb = cv2.cvtColor(img_removed, cv2.COLOR_BGR2RGB)
+    
+    return img_removed_rgb, mask
+        
+def fit_text_to_box(text, merged_coordinates, font_path, font_color, angle=0):
+    """
+    Adjusts the text to fit optimally inside the given box dimensions.
+    
+    Args:
+        text (str): The text to fit.
+        box_size (tuple): A tuple (width, height) specifying the box dimensions.
+        font_path (str): Path to the font file to be used.
+    
+    Returns:
+        PIL.Image: An image with the text fitted inside the box.
+    """
+    width, height = merged_coordinates[1][0] - merged_coordinates[0][0], merged_coordinates[2][1] - merged_coordinates[1][1]
+    font_size = 1
+
+    # Create a dummy image to measure text size
+    dummy_image = Image.new('RGB', (width, height))
+    draw = ImageDraw.Draw(dummy_image)
+
+    # Load a small font initially
+    font = ImageFont.truetype(font_path, font_size)
+
+    while True:
+        # Break text into lines that fit within the width
+        words = text.split()
+        lines = []
+        current_line = []
+        for word in words:
+            test_line = " ".join(current_line + [word])
+            test_width = draw.textlength(test_line, font=font)
+            if test_width <= width:
+                current_line.append(word)
+            else:
+                lines.append(" ".join(current_line))
+                current_line = [word]
+        if current_line:
+            lines.append(" ".join(current_line))
+
+        # Calculate total height required for the lines
+        line_height = font.getbbox('A')[3] + 5  # Add line spacing
+        total_height = len(lines) * line_height
+
+        # Check if text fits within the height
+        if total_height > height or any(draw.textlength(line, font=font) > width for line in lines):
+            break
+
+        # Increment font size
+        font_size += 1
+        font = ImageFont.truetype(font_path, font_size)
+
+    # Use the last fitting font
+    font_size -= 1
+    font = ImageFont.truetype(font_path, font_size)
+
+    # Create the final image with a transparent background
+    image = Image.new('RGBA', (width, height), (255, 255, 255, 0))
+    draw = ImageDraw.Draw(image)
+
+    # Center the text vertically and horizontally
+    lines = []
+    current_line = []
+    for word in text.split():
+        test_line = " ".join(current_line + [word])
+        if draw.textlength(test_line, font=font) <= width:
+            current_line.append(word)
+        else:
+            lines.append(" ".join(current_line))
+            current_line = [word]
+    if current_line:
+        lines.append(" ".join(current_line))
+
+    line_height = font.getbbox('A')[3] + 5
+    total_text_height = len(lines) * line_height
+    y_offset = (height - total_text_height) // 2
+
+    for line in lines:
+        text_width = draw.textlength(line, font=font)
+        x_offset = (width - text_width) // 2
+        draw.text((x_offset, y_offset), line, font=font, fill=font_color)
+        y_offset += line_height
+
+    rotated_image = image.rotate(0, expand=True)
+
+    return rotated_image
+
+def shorten_box(merged_coordinates, pct=0):
+    # Calculate the center of the box
+    center_x = (merged_coordinates[0][0] + merged_coordinates[2][0]) / 2
+    center_y = (merged_coordinates[0][1] + merged_coordinates[2][1]) / 2
+    
+    # Calculate the width and height of the box
+    width = merged_coordinates[1][0] - merged_coordinates[0][0]
+    height = merged_coordinates[2][1] - merged_coordinates[1][1]
+    
+    # Shrink width and height by 10%
+    new_width = width * 1-pct/100.
+    new_height = height * 1-pct/100.
+    
+    # Calculate the new coordinates
+    merged_coordinates_new = np.array([
+        [center_x - new_width / 2, center_y - new_height / 2],  # Top-left
+        [center_x + new_width / 2, center_y - new_height / 2],  # Top-right
+        [center_x + new_width / 2, center_y + new_height / 2],  # Bottom-right
+        [center_x - new_width / 2, center_y + new_height / 2]   # Bottom-left
+    ], dtype=int)
+
+    return merged_coordinates_new
+
+
+def detect_and_show_text(reader, image):
+    """Detect text and show bounding boxes"""
+    if isinstance(image, Image.Image):
+        img_array = np.array(image)
+    else:
+        img_array = image
+        
+    # Get YOLO results first
+    yolo_results = bubble_detection_model(img_array, conf=0.7)[0]
+    yolo_boxes = yolo_results.boxes.xywh.cpu().numpy()  # Get YOLO boxes in xywh format
+    
+    # Detect text
+    results = reader.readtext(img_array, text_threshold=0.6)
+    
+    # Create visualization
+    img_boxes = img_array.copy()
+    
+    # Ensure we're working with RGB
+    if len(img_array.shape) == 3:
+        if img_array.shape[2] == 3:  # If it's a 3-channel image
+            img_boxes = cv2.cvtColor(img_boxes, cv2.COLOR_BGR2RGB)
+    
+    # Draw original EasyOCR boxes on img_boxes
+    for result in results:
+        pts = np.array(result[0], dtype=np.int32)
+        cv2.polylines(img_boxes, [pts], isClosed=True, color=(0, 255, 0), thickness=2)  # Draw original boxes in green
+    
+    # Remove text and merge boxes for visualization
+    img_removed, mask = remove_text_regions(img_array, results, yolo_boxes)
+    merged_results = merge_boxes(results, img_array.shape)
+    
+    # Draw merged detection boxes and text (if needed)
+    for merged_result in merged_results:
+        pts = np.array(merged_result['box'], dtype=np.int32)
+        # Color the box red if inside bubble, blue if outside
+        color = (0, 0, 255) if is_box_inside_yolo(merged_result['box'], yolo_boxes) else (255, 0, 0)
+        cv2.polylines(img_boxes, [pts], True, color, 2)  # Draw merged boxes in red or blue
+    
+    # Convert to RGB
+    img_boxes_rgb = cv2.cvtColor(img_boxes, cv2.COLOR_BGR2RGB)
+    img_removed_rgb = cv2.cvtColor(img_removed, cv2.COLOR_BGR2RGB)
+    mask_rgba = cv2.cvtColor(mask, cv2.COLOR_RGB2RGBA)
+    
+    # Get YOLO visualization without labels
+    bubbles_img = yolo_results.plot(labels=False)
+    
+    # Convert to PIL Images
+    img_boxes_pil = Image.fromarray(img_boxes_rgb)
+    img_removed_pil = Image.fromarray(img_removed_rgb)
+    bubbles_img_pil = Image.fromarray(bubbles_img)
+    mask_pil = Image.fromarray(mask_rgba)
+
+    return img_boxes_pil, bubbles_img_pil, img_removed_pil, merged_results, mask_pil
+
+
+def position_text_back(text, merged_coordinates, inpainted_image, font_path, font_color):
+    coords = shorten_box(merged_coordinates)
+    top_left_coords = coords[0]
+    text_image = fit_text_to_box(text, coords, font_path, font_color)
+
+    # Create a transparent layer to blend
+    layer = Image.new("RGBA", inpainted_image.size, (0, 0, 0, 0))
+
+    # Paste the text image onto the transparent layer at the specified position
+    layer.paste(text_image, tuple(top_left_coords), mask=text_image)
+
+    # Ensure both images are in "RGBA" mode
+    if inpainted_image.mode != "RGBA":
+        inpainted_image = inpainted_image.convert("RGBA")
+    if layer.mode != "RGBA":
+        layer = layer.convert("RGBA")
+
+    # Blend the transparent layer with the inpainted image
+    blended_image = Image.alpha_composite(inpainted_image, layer)
+
+    return blended_image
+
+@spaces.GPU()
+def process(image, mask,
+            prompt="background",
+            negative_prompt="text",
+            num_inference_steps=10,
+            controlnet_conditioning_scale=0.9,
+            guidance_scale=3.5,
+            seed=124,
+            true_guidance_scale=3.5            
+            ):
+    size = (768, 768)
+    image_pil = Image.fromarray(image)
+    image_or = image_pil.copy()
+    
+    image_pil = image_pil.convert("RGB").resize(size)
+    mask = mask.convert("RGB").resize(size)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    result = pipe(
+    prompt=prompt,
+    height=size[1],
+    width=size[0],
+    control_image=image_pil,
+    control_mask=mask,
+    num_inference_steps=num_inference_steps,
+    generator=generator,
+    controlnet_conditioning_scale=controlnet_conditioning_scale,
+    guidance_scale=guidance_scale,
+    negative_prompt=negative_prompt,
+    true_guidance_scale=true_guidance_scale
+    ).images[0]
+
+    return result.resize((image_or.size[:2]))
+
+
+@spaces.GPU()
+def process_image(image, source_language, target_language, mode, font, font_color, num_inference_steps):
+    """Main processing function for Gradio"""
+    if image is None:
+        return None, None, None, []
+    
+    # Initialize reader (equivalent to what handle_localization did)
+    easy_ocr_lan = language_to_ocr.get(source_language, 'en')
+    reader = easyocr.Reader([easy_ocr_lan], model_storage_directory='.', gpu=False)
+    
+    # Detect text and get results
+    img_with_boxes, img_bubbles, img_removed_text, merged_results, mask = detect_and_show_text(reader, image)
+
+    if mode == "Basic":
+        img_inpainted = img_removed_text
+    else:
+        img_inpainted = process(image, mask, num_inference_steps=num_inference_steps)
+
+    font_rgba = color_to_rgba(font_color)  # Convert hex to RGBA
+    
+    # Get translations
+    translations = localize_boxes(
+        merged_results,
+        img_with_boxes,
+        source_language,
+        target_language
+    )
+    
+    # Create initial result with translations
+    final_result = img_inpainted.copy()
+    for translation in translations:
+        box = translation['box']
+        text = translation['text']
+        final_result = position_text_back(text, box, final_result, font_path=f"fonts/{font}.ttf", font_color=font_rgba)
+    
+    # Return all results directly (no need to store in session state)
+    return img_with_boxes, img_bubbles, img_inpainted, final_result, translations, np.array(mask)
+
+
+def update_translations(image, edited_texts, translations_list, img_removed_text, font, font_color):
+    """Update the image with edited translations"""
+    if image is None or img_removed_text is None:
+        return None
+    
+    # Convert numpy array back to PIL Image
+    img_removed = Image.fromarray(img_removed_text)
+    final_result = img_removed.copy()
+    
+    font_rgba = color_to_rgba(font_color)  # Convert hex to RGBA
+
+    # Update the translations with edited texts
+    for trans, new_text in zip(translations_list, edited_texts.split('\n')):
+        trans['text'] = new_text.strip()
+        box = trans['box']
+        final_result = position_text_back(new_text, box, final_result, font_path=f"fonts/{font}.ttf", font_color=font_rgba)
+    
+    return np.array(final_result)
+
+
+
+with gr.Blocks(title="Nativ - Demo") as demo:
+    # Store translations list in state
+    translations_state = gr.State([])
+    
+    gr.Markdown("# Nativ - Demo")
+    
+    with gr.Row():
+        with gr.Column():
+            # Input components
+            input_image = gr.Image(type="numpy", label="Upload Image")
+            source_language = gr.Dropdown(
+                choices=['Simplified Chinese', 'Traditional Chinese', 'Korean', 'Japanese', 'English'],
+                value='Simplified Chinese',
+                label="Source Language"
+            )
+            target_language = gr.Dropdown(
+                choices=['English', 'Spanish', 'Chinese', 'Korean', 'French', 'Japanese'],
+                value='English',
+                label="Target Language"
+            )
+            # Toggle for mode selection
+            localization_mode = gr.Radio(
+                choices=["Basic", "Advanced"],
+                value="Basic",
+                label="Localization Mode"
+            )
+            font_selector_i = gr.Dropdown(
+            choices=['Arial', 'Ldfcomicsansbold', 'Times New Roman', 'georgia', 'calibri', 'Verdana', 'omniscript_bold', 'helvetica'],  # Add more fonts as needed
+            value='omniscript_bold',
+                label="Select Font"
+            )
+            num_inference_steps = gr.Slider(minimum=1, step=1, maximum=50, value=10, label="num_inference_steps")
+            font_color_picker_i = gr.ColorPicker(
+                value="#000000",  # Default color: black
+                label="Select Font Color"
+            )
+            process_btn = gr.Button("Localize")
+        
+        with gr.Column():
+            # Output components
+            speech_bubbles = gr.Image(type="numpy", label="Detected Speech Bubbles", interactive=False)
+            detected_boxes = gr.Image(type="numpy", label="Detected Text Regions", interactive=False)
+            removed_text = gr.Image(type="numpy", label="Removed Text", interactive=False)
+            final_output = gr.Image(type="numpy", label="Final Result", interactive=False)
+    
+    # Translation editing section
+    with gr.Row():
+        translations_text = gr.Textbox(
+            label="Edit Translations (one per line)", 
+            lines=5,
+            placeholder="Edit translations here..."
+        )
+        font_selector_f = gr.Dropdown(
+            choices=['Arial', 'Ldfcomicsansbold', 'Times New Roman', 'georgia', 'calibri', 'Verdana', 'omniscript_bold', 'helvetica'],  # Add more fonts as needed
+            value='omniscript_bold',
+            label="Select Font"
+        )
+        font_color_picker_f = gr.ColorPicker(
+            value="#000000",  # Default color: black
+            label="Select Font Color"
+        )
+        update_btn = gr.Button("Apply Changes")
+    
+    def process_and_show_translations(image, source_lang, target_lang, mode, font, font_color, num_inference_steps):
+        boxes, bubbles, removed, final, translations, mask = process_image(image, source_lang, target_lang, mode, font, font_color, num_inference_steps)
+        # Extract just the texts and join with newlines
+        texts = '\n'.join(t['text'] for t in translations)
+        return boxes, bubbles, removed, final, texts, translations
+    
+    # Process button click
+    process_btn.click(
+        fn=process_and_show_translations,
+        inputs=[input_image, source_language, target_language, localization_mode, font_selector_i, font_color_picker_i, num_inference_steps],
+        outputs=[detected_boxes, speech_bubbles, removed_text, final_output, translations_text, translations_state]
+    )
+    
+    # Update translations button click
+    update_btn.click(
+        fn=update_translations,
+        inputs=[input_image, translations_text, translations_state, removed_text, font_selector_f, font_color_picker_f],
+        outputs=final_output
+    )
+
+
+demo.launch(debug=False, show_error=True,share=True)

controlnet_flux.py

@@ -0,0 +1,418 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.attention_processor import AttentionProcessor
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_version,
+    logging,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.models.controlnet import BaseOutput, zero_module
+from diffusers.models.embeddings import (
+    CombinedTimestepGuidanceTextProjEmbeddings,
+    CombinedTimestepTextProjEmbeddings,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from transformer_flux import (
+    EmbedND,
+    FluxSingleTransformerBlock,
+    FluxTransformerBlock,
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FluxControlNetOutput(BaseOutput):
+    controlnet_block_samples: Tuple[torch.Tensor]
+    controlnet_single_block_samples: Tuple[torch.Tensor]
+
+
+class FluxControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False,
+        axes_dims_rope: List[int] = [16, 56, 56],
+        extra_condition_channels: int = 1 * 4,
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pos_embed = EmbedND(
+            dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
+        )
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings
+            if guidance_embeds
+            else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        # controlnet_blocks
+        self.controlnet_blocks = nn.ModuleList([])
+        for _ in range(len(self.transformer_blocks)):
+            self.controlnet_blocks.append(
+                zero_module(nn.Linear(self.inner_dim, self.inner_dim))
+            )
+
+        self.controlnet_single_blocks = nn.ModuleList([])
+        for _ in range(len(self.single_transformer_blocks)):
+            self.controlnet_single_blocks.append(
+                zero_module(nn.Linear(self.inner_dim, self.inner_dim))
+            )
+
+        self.controlnet_x_embedder = zero_module(
+            torch.nn.Linear(in_channels + extra_condition_channels, self.inner_dim)
+        )
+
+        self.gradient_checkpointing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self):
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    @classmethod
+    def from_transformer(
+        cls,
+        transformer,
+        num_layers: int = 4,
+        num_single_layers: int = 10,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        load_weights_from_transformer=True,
+    ):
+        config = transformer.config
+        config["num_layers"] = num_layers
+        config["num_single_layers"] = num_single_layers
+        config["attention_head_dim"] = attention_head_dim
+        config["num_attention_heads"] = num_attention_heads
+
+        controlnet = cls(**config)
+
+        if load_weights_from_transformer:
+            controlnet.pos_embed.load_state_dict(transformer.pos_embed.state_dict())
+            controlnet.time_text_embed.load_state_dict(
+                transformer.time_text_embed.state_dict()
+            )
+            controlnet.context_embedder.load_state_dict(
+                transformer.context_embedder.state_dict()
+            )
+            controlnet.x_embedder.load_state_dict(transformer.x_embedder.state_dict())
+            controlnet.transformer_blocks.load_state_dict(
+                transformer.transformer_blocks.state_dict(), strict=False
+            )
+            controlnet.single_transformer_blocks.load_state_dict(
+                transformer.single_transformer_blocks.state_dict(), strict=False
+            )
+
+            controlnet.controlnet_x_embedder = zero_module(
+                controlnet.controlnet_x_embedder
+            )
+
+        return controlnet
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        controlnet_cond: torch.Tensor,
+        conditioning_scale: float = 1.0,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if (
+                joint_attention_kwargs is not None
+                and joint_attention_kwargs.get("scale", None) is not None
+            ):
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+        hidden_states = self.x_embedder(hidden_states)
+
+        # add condition
+        hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        image_rotary_emb = self.pos_embed(ids)
+
+        block_samples = ()
+        for _, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
+                (
+                    encoder_hidden_states,
+                    hidden_states,
+                ) = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            block_samples = block_samples + (hidden_states,)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        single_block_samples = ()
+        for _, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            single_block_samples = single_block_samples + (
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+        # controlnet block
+        controlnet_block_samples = ()
+        for block_sample, controlnet_block in zip(
+            block_samples, self.controlnet_blocks
+        ):
+            block_sample = controlnet_block(block_sample)
+            controlnet_block_samples = controlnet_block_samples + (block_sample,)
+
+        controlnet_single_block_samples = ()
+        for single_block_sample, controlnet_block in zip(
+            single_block_samples, self.controlnet_single_blocks
+        ):
+            single_block_sample = controlnet_block(single_block_sample)
+            controlnet_single_block_samples = controlnet_single_block_samples + (
+                single_block_sample,
+            )
+
+        # scaling
+        controlnet_block_samples = [
+            sample * conditioning_scale for sample in controlnet_block_samples
+        ]
+        controlnet_single_block_samples = [
+            sample * conditioning_scale for sample in controlnet_single_block_samples
+        ]
+
+        #
+        controlnet_block_samples = (
+            None if len(controlnet_block_samples) == 0 else controlnet_block_samples
+        )
+        controlnet_single_block_samples = (
+            None
+            if len(controlnet_single_block_samples) == 0
+            else controlnet_single_block_samples
+        )
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (controlnet_block_samples, controlnet_single_block_samples)
+
+        return FluxControlNetOutput(
+            controlnet_block_samples=controlnet_block_samples,
+            controlnet_single_block_samples=controlnet_single_block_samples,
+        )

craft_mlt_25k.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a5efbfb48b4081100544e75e1e2b57f8de3d84f213004b14b85fd4b3748db17
+size 83152330

english_g2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2272681d9d67a04e2dff396b6e95077bc19001f8f6d3593c307b9852e1c29e8
+size 15143997

fonts/app.py

@@ -0,0 +1,575 @@
+import spaces
+import os
+import gradio as gr
+import easyocr
+import numpy as np
+import cv2
+import base64
+import torch
+from shapely import Polygon
+from ultralytics import YOLO
+
+from io import BytesIO
+from openai import OpenAI
+from PIL import Image, ImageDraw, ImageFont
+
+from diffusers.utils import load_image, check_min_version
+from controlnet_flux import FluxControlNetModel
+from transformer_flux import FluxTransformer2DModel
+from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
+
+import huggingface_hub
+huggingface_hub.login(os.getenv('HF_TOKEN_FLUX'))
+
+bubble_detection_model = YOLO("speech_bubble_model.pt")
+
+language_to_ocr = {
+    'Simplified Chinese': 'ch_sim',
+    'Traditional Chinese': 'ch_tra',
+    'Korean': 'ko',
+    'Japanese': 'ja',
+    'English': 'en',
+}
+
+OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
+
+MARKDOWN = """
+# Made by Nativ
+"""
+
+check_min_version("0.30.2")
+transformer = FluxTransformer2DModel.from_pretrained(
+        "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dytpe=torch.bfloat16
+    )
+
+cuda_device =torch.device("cuda")
+# Build pipeline
+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,
+    transformer=transformer,
+    torch_dtype=torch.bfloat16
+).to(cuda_device)
+pipe.transformer.to(torch.bfloat16)
+pipe.controlnet.to(torch.bfloat16)
+
+
+def localize_boxes(merged_results, img_boxes, source_language, target_language):
+    # Convert image to base64
+    buffered = BytesIO()
+    img_boxes.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    
+    print(merged_results)
+    
+    prompt = f"""You are an expert translator and localization specialist with deep understanding of both {source_language} and {target_language} cultures.
+
+    Task: Translate the detected text while preserving the cultural context and maintaining visual harmony. Make the results in capital letters.
+
+    Source Text and Coordinates:
+    {merged_results}
+
+    Requirements:
+    1. Maintain the original meaning and tone while adapting to {target_language} cultural context
+    2. Keep translations concise and visually balanced (similar character length when possible)
+    3. Preserve any:
+    - Brand names
+    - Product names
+    - Technical terms
+    - Numbers and units
+    4. Consider the visual context from the provided image
+    5. Use appropriate formality level for {target_language}
+    6. Maintain any special formatting (if present)
+
+    Format your response EXACTLY as a JSON-like list of dictionaries. Keep the box coordinates EXACTLY as they are, do not change them, only translate the text.
+    [{{'box': [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], 'text': 'translated_text'}}]
+
+    Important: Only output the JSON format above, no explanations or additional text."""
+
+    client = OpenAI(api_key=OPENAI_API_KEY)
+    
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{img_str}"
+                        }
+                    }
+                ]
+            }
+        ],
+        max_tokens=1000,
+        temperature=0
+    )
+    
+    try:
+        translation_text = response.choices[0].message.content
+        translation_text = translation_text.replace("```json", "").replace("```", "").strip()
+        translated_results = eval(translation_text)
+        return translated_results
+    except Exception as e:
+        print(f"Error parsing GPT-4o response: {e}")
+        return merged_results  
+    
+def merge_boxes(boxes, image_shape, distance_threshold=10):
+    """Merge boxes that are close to each other and return their associated text"""
+    if not boxes:
+        return []
+    
+    # Extract boxes and create mapping to original data
+    boxes_only = [box[0] for box in boxes]
+    texts = [box[1] for box in boxes]  # Extract the text content
+    
+    # Create a binary mask of all boxes
+    height, width = image_shape[:2]
+    mask = np.zeros((height, width), dtype=np.uint8)
+    
+    # Draw all boxes on mask and create a mapping of pixel positions to box indices
+    box_indices_map = {}  # Will store which original box each pixel belongs to
+    for idx, coords in enumerate(boxes_only):
+        pts = np.array(coords, dtype=np.int32)
+        cv2.fillPoly(mask, [pts], 255)
+        # Store the indices of boxes for each filled pixel
+        y_coords, x_coords = np.where(mask == 255)
+        for y, x in zip(y_coords, x_coords):
+            if (y, x) not in box_indices_map:
+                box_indices_map[(y, x)] = []
+            box_indices_map[(y, x)].append(idx)
+    
+    # Dilate to connect nearby components
+    kernel = np.ones((distance_threshold, distance_threshold), np.uint8)
+    dilated = cv2.dilate(mask, kernel, iterations=1)
+    
+    # Find connected components
+    num_labels, labels = cv2.connectedComponents(dilated)
+    
+    # Create new merged boxes with their associated text
+    merged_results = []
+    for label in range(1, num_labels):  # Skip background (0)
+        points = np.where(labels == label)
+        if len(points[0]):  # If component is not empty
+            y0, x0 = points[0].min(), points[1].min()
+            y1, x1 = points[0].max(), points[1].max()
+            # Add small padding
+            x0 = max(0, x0 - 2)
+            y0 = max(0, y0 - 2)
+            x1 = min(width, x1 + 2)
+            y1 = min(height, y1 + 2)
+            
+            # Find all original boxes that overlap with this merged box
+            box_indices = set()
+            for y in range(y0, y1+1):
+                for x in range(x0, x1+1):
+                    if (y, x) in box_indices_map:
+                        box_indices.update(box_indices_map[(y, x)])
+            
+            # Combine text from all overlapping boxes
+            combined_text = ' '.join([texts[idx] for idx in box_indices])
+            
+            merged_results.append({
+                'box': [[x0, y0], [x1, y0], [x1, y1], [x0, y1]],
+                'text': combined_text
+            })    
+    return merged_results
+
+def is_box_inside_yolo(box, yolo_boxes, overlap_threshold=0.5):
+    """
+    Check if a text box is inside any of the YOLO-detected speech bubbles.
+    box: [[x0,y0], [x1,y0], [x1,y1], [x0,y1]]
+    yolo_boxes: list of YOLO boxes in xywh format
+    overlap_threshold: minimum overlap ratio required to consider the text inside bubble
+    """
+    text_poly = Polygon(box)
+    text_area = text_poly.area
+    
+    for yolo_box in yolo_boxes:
+        x_center, y_center, width, height = yolo_box
+        x1, y1 = x_center - width / 2, y_center - height / 2
+        x2, y2 = x_center + width / 2, y_center + height / 2
+        bubble_box = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
+        bubble_poly = Polygon(bubble_box)
+        
+        # Calculate intersection
+        if text_poly.intersects(bubble_poly):
+            intersection = text_poly.intersection(bubble_poly)
+            overlap_ratio = intersection.area / text_area
+            if overlap_ratio >= overlap_threshold:
+                return True
+    
+    return False
+
+def remove_text_regions(image, boxes, yolo_boxes):
+    """Fill detected text regions with white"""
+    img_removed = image.copy()
+    mask = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    
+    # Fill all detected boxes with white
+    for box in boxes:
+        pts = np.array(box[0], dtype=np.int32)
+        if is_box_inside_yolo(box[0], yolo_boxes):
+            cv2.fillPoly(img_removed, [pts], (255, 255, 255, 255))
+        cv2.fillPoly(mask, [pts], (255, 255, 255, 255))
+    
+    img_removed_rgb = cv2.cvtColor(img_removed, cv2.COLOR_BGR2RGB)
+    
+    return img_removed_rgb, mask
+        
+def fit_text_to_box(text, merged_coordinates, angle=0, font_path):
+    """
+    Adjusts the text to fit optimally inside the given box dimensions.
+    
+    Args:
+        text (str): The text to fit.
+        box_size (tuple): A tuple (width, height) specifying the box dimensions.
+        font_path (str): Path to the font file to be used.
+    
+    Returns:
+        PIL.Image: An image with the text fitted inside the box.
+    """
+    width, height = merged_coordinates[1][0] - merged_coordinates[0][0], merged_coordinates[2][1] - merged_coordinates[1][1]
+    font_size = 1
+
+    # Create a dummy image to measure text size
+    dummy_image = Image.new('RGB', (width, height))
+    draw = ImageDraw.Draw(dummy_image)
+
+    # Load a small font initially
+    font = ImageFont.truetype(font_path, font_size)
+
+    while True:
+        # Break text into lines that fit within the width
+        words = text.split()
+        lines = []
+        current_line = []
+        for word in words:
+            test_line = " ".join(current_line + [word])
+            test_width = draw.textlength(test_line, font=font)
+            if test_width <= width:
+                current_line.append(word)
+            else:
+                lines.append(" ".join(current_line))
+                current_line = [word]
+        if current_line:
+            lines.append(" ".join(current_line))
+
+        # Calculate total height required for the lines
+        line_height = font.getbbox('A')[3] + 5  # Add line spacing
+        total_height = len(lines) * line_height
+
+        # Check if text fits within the height
+        if total_height > height or any(draw.textlength(line, font=font) > width for line in lines):
+            break
+
+        # Increment font size
+        font_size += 1
+        font = ImageFont.truetype(font_path, font_size)
+
+    # Use the last fitting font
+    font_size -= 1
+    font = ImageFont.truetype(font_path, font_size)
+
+    # Create the final image with a transparent background
+    image = Image.new('RGBA', (width, height), (255, 255, 255, 0))
+    draw = ImageDraw.Draw(image)
+
+    # Center the text vertically and horizontally
+    lines = []
+    current_line = []
+    for word in text.split():
+        test_line = " ".join(current_line + [word])
+        if draw.textlength(test_line, font=font) <= width:
+            current_line.append(word)
+        else:
+            lines.append(" ".join(current_line))
+            current_line = [word]
+    if current_line:
+        lines.append(" ".join(current_line))
+
+    line_height = font.getbbox('A')[3] + 5
+    total_text_height = len(lines) * line_height
+    y_offset = (height - total_text_height) // 2
+
+    for line in lines:
+        text_width = draw.textlength(line, font=font)
+        x_offset = (width - text_width) // 2
+        draw.text((x_offset, y_offset), line, font=font, fill="black")
+        y_offset += line_height
+
+    rotated_image = image.rotate(0, expand=True)
+
+    return rotated_image
+
+def shorten_box(merged_coordinates, pct=0):
+    # Calculate the center of the box
+    center_x = (merged_coordinates[0][0] + merged_coordinates[2][0]) / 2
+    center_y = (merged_coordinates[0][1] + merged_coordinates[2][1]) / 2
+    
+    # Calculate the width and height of the box
+    width = merged_coordinates[1][0] - merged_coordinates[0][0]
+    height = merged_coordinates[2][1] - merged_coordinates[1][1]
+    
+    # Shrink width and height by 10%
+    new_width = width * 1-pct/100.
+    new_height = height * 1-pct/100.
+    
+    # Calculate the new coordinates
+    merged_coordinates_new = np.array([
+        [center_x - new_width / 2, center_y - new_height / 2],  # Top-left
+        [center_x + new_width / 2, center_y - new_height / 2],  # Top-right
+        [center_x + new_width / 2, center_y + new_height / 2],  # Bottom-right
+        [center_x - new_width / 2, center_y + new_height / 2]   # Bottom-left
+    ], dtype=int)
+
+    return merged_coordinates_new
+
+
+def detect_and_show_text(reader, image):
+    """Detect text and show bounding boxes"""
+    if isinstance(image, Image.Image):
+        img_array = np.array(image)
+    else:
+        img_array = image
+        
+    # Get YOLO results first
+    yolo_results = bubble_detection_model(img_array, conf=7)[0]
+    yolo_boxes = yolo_results.boxes.xywh.cpu().numpy()  # Get YOLO boxes in xywh format
+    
+    # Detect text
+    results = reader.readtext(img_array, text_threshold=0.6)
+    
+    # Create visualization
+    img_boxes = img_array.copy()
+    
+    # Ensure we're working with RGB
+    if len(img_array.shape) == 3:
+        if img_array.shape[2] == 3:  # If it's a 3-channel image
+            img_boxes = cv2.cvtColor(img_boxes, cv2.COLOR_BGR2RGB)
+    
+    # Draw original EasyOCR boxes on img_boxes
+    for result in results:
+        pts = np.array(result[0], dtype=np.int32)
+        cv2.polylines(img_boxes, [pts], isClosed=True, color=(0, 255, 0), thickness=2)  # Draw original boxes in green
+    
+    # Remove text and merge boxes for visualization
+    img_removed, mask = remove_text_regions(img_array, results, yolo_boxes)
+    merged_results = merge_boxes(results, img_array.shape)
+    
+    # Draw merged detection boxes and text (if needed)
+    for merged_result in merged_results:
+        pts = np.array(merged_result['box'], dtype=np.int32)
+        # Color the box red if inside bubble, blue if outside
+        color = (0, 0, 255) if is_box_inside_yolo(merged_result['box'], yolo_boxes) else (255, 0, 0)
+        cv2.polylines(img_boxes, [pts], True, color, 2)  # Draw merged boxes in red or blue
+    
+    # Convert to RGB
+    img_boxes_rgb = cv2.cvtColor(img_boxes, cv2.COLOR_BGR2RGB)
+    img_removed_rgb = cv2.cvtColor(img_removed, cv2.COLOR_BGR2RGB)
+    mask_rgba = cv2.cvtColor(mask, cv2.COLOR_RGB2RGBA)
+    
+    # Get YOLO visualization without labels
+    bubbles_img = yolo_results.plot(labels=False)
+    
+    # Convert to PIL Images
+    img_boxes_pil = Image.fromarray(img_boxes_rgb)
+    img_removed_pil = Image.fromarray(img_removed_rgb)
+    bubbles_img_pil = Image.fromarray(bubbles_img)
+    mask_pil = Image.fromarray(mask_rgba)
+
+    return img_boxes_pil, bubbles_img_pil, img_removed_pil, merged_results, mask_pil
+
+
+def position_text_back(text, merged_coordinates, inpainted_image, font_path):
+    coords = shorten_box(merged_coordinates)
+    top_left_coords = coords[0]
+    text_image = fit_text_to_box(text, coords, font_path)
+
+    # Create a transparent layer to blend
+    layer = Image.new("RGBA", inpainted_image.size, (0, 0, 0, 0))
+
+    # Paste the text image onto the transparent layer at the specified position
+    layer.paste(text_image, tuple(top_left_coords), mask=text_image)
+
+    # Ensure both images are in "RGBA" mode
+    if inpainted_image.mode != "RGBA":
+        inpainted_image = inpainted_image.convert("RGBA")
+    if layer.mode != "RGBA":
+        layer = layer.convert("RGBA")
+
+    # Blend the transparent layer with the inpainted image
+    blended_image = Image.alpha_composite(inpainted_image, layer)
+
+    return blended_image
+
+@spaces.GPU()
+def process(image, mask,
+            prompt="background",
+            negative_prompt="text",
+            controlnet_conditioning_scale=0.9,
+            guidance_scale=3.5,
+            seed=124,
+            num_inference_steps=10,
+            true_guidance_scale=3.5            
+            ):
+    size = (768, 768)
+    image_pil = Image.fromarray(image)
+    image_or = image_pil.copy()
+    
+    image_pil = image_pil.convert("RGB").resize(size)
+    mask = mask.convert("RGB").resize(size)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    result = pipe(
+    prompt=prompt,
+    height=size[1],
+    width=size[0],
+    control_image=image_pil,
+    control_mask=mask,
+    num_inference_steps=num_inference_steps,
+    generator=generator,
+    controlnet_conditioning_scale=controlnet_conditioning_scale,
+    guidance_scale=guidance_scale,
+    negative_prompt=negative_prompt,
+    true_guidance_scale=true_guidance_scale
+    ).images[0]
+
+    return result.resize((image_or.size[:2]))
+
+
+@spaces.GPU()
+def process_image(image, source_language, target_language, mode, font):
+    """Main processing function for Gradio"""
+    if image is None:
+        return None, None, None, []
+    
+    # Initialize reader (equivalent to what handle_localization did)
+    easy_ocr_lan = language_to_ocr.get(source_language, 'en')
+    reader = easyocr.Reader([easy_ocr_lan], model_storage_directory='.', gpu=False)
+    
+    # Detect text and get results
+    img_with_boxes, img_bubbles, img_removed_text, merged_results, mask = detect_and_show_text(reader, image)
+
+    if mode == "Basic":
+        img_inpainted = img_removed_text
+    else:
+        img_inpainted = process(image, mask)
+    
+    # Get translations
+    translations = localize_boxes(
+        merged_results,
+        img_with_boxes,
+        source_language,
+        target_language
+    )
+    
+    # Create initial result with translations
+    final_result = img_inpainted.copy()
+    for translation in translations:
+        box = translation['box']
+        text = translation['text']
+        final_result = position_text_back(text, box, final_result, font_path=f"fonts/{font}.ttf")
+    
+    # Return all results directly (no need to store in session state)
+    return img_with_boxes, img_bubbles, img_inpainted, final_result, translations, np.array(mask)
+
+
+def update_translations(image, edited_texts, translations_list, img_removed_text, font):
+    """Update the image with edited translations"""
+    if image is None or img_removed_text is None:
+        return None
+    
+    # Convert numpy array back to PIL Image
+    img_removed = Image.fromarray(img_removed_text)
+    final_result = img_removed.copy()
+    
+    # Update the translations with edited texts
+    for trans, new_text in zip(translations_list, edited_texts.split('\n')):
+        trans['text'] = new_text.strip()
+        box = trans['box']
+        final_result = position_text_back(new_text, box, final_result, font_path=f"fonts/{font}.ttf")
+    
+    return np.array(final_result)
+
+
+
+with gr.Blocks(title="Nativ - Demo") as demo:
+    # Store translations list in state
+    translations_state = gr.State([])
+    
+    gr.Markdown("# Nativ - Demo")
+    
+    with gr.Row():
+        with gr.Column():
+            # Input components
+            input_image = gr.Image(type="numpy", label="Upload Image")
+            source_language = gr.Dropdown(
+                choices=['Simplified Chinese', 'Traditional Chinese', 'Korean', 'Japanese', 'English'],
+                value='Simplified Chinese',
+                label="Source Language"
+            )
+            target_language = gr.Dropdown(
+                choices=['English', 'Spanish', 'Chinese', 'Korean', 'French', 'Japanese'],
+                value='English',
+                label="Target Language"
+            )
+            # Toggle for mode selection
+            localization_mode = gr.Radio(
+                choices=["Basic", "Advanced"],
+                value="Basic",
+                label="Localization Mode"
+            )
+            font_selector_i = gr.Dropdown(
+            choices=['Arial', 'Ldfcomicsansbold', 'Times New Roman', 'georgia', 'calibri', 'Verdana', 'omniscript_bold', 'helvetica'],  # Add more fonts as needed
+            value='omniscript_bold',
+                label="Select Font"
+            )
+            process_btn = gr.Button("Localize")
+        
+        with gr.Column():
+            # Output components
+            speech_bubbles = gr.Image(type="numpy", label="Detected Speech Bubbles", interactive=False)
+            detected_boxes = gr.Image(type="numpy", label="Detected Text Regions", interactive=False)
+            removed_text = gr.Image(type="numpy", label="Removed Text", interactive=False)
+            final_output = gr.Image(type="numpy", label="Final Result", interactive=False)
+    
+    # Translation editing section
+    with gr.Row():
+        translations_text = gr.Textbox(
+            label="Edit Translations (one per line)", 
+            lines=5,
+            placeholder="Edit translations here..."
+        )
+        font_selector_f = gr.Dropdown(
+            choices=['Arial', 'Ldfcomicsansbold', 'Times New Roman', 'georgia', 'calibri', 'Verdana', 'omniscript_bold', 'helvetica'],  # Add more fonts as needed
+            value='Arial',
+            label="Select Font"
+        )
+        update_btn = gr.Button("Apply Changes")
+    
+    def process_and_show_translations(image, source_lang, target_lang, mode, font):
+        boxes, bubbles, removed, final, translations, mask = process_image(image, source_lang, target_lang, mode, font)
+        # Extract just the texts and join with newlines
+        texts = '\n'.join(t['text'] for t in translations)
+        return boxes, bubbles, removed, final, texts, translations
+    
+    # Process button click
+    process_btn.click(
+        fn=process_and_show_translations,
+        inputs=[input_image, source_language, target_language, localization_mode, font_selector_i],
+        outputs=[detected_boxes, speech_bubbles, removed_text, final_output, translations_text, translations_state]
+    )
+    
+    # Update translations button click
+    update_btn.click(
+        fn=update_translations,
+        inputs=[input_image, translations_text, translations_state, removed_text, font_selector_f],
+        outputs=final_output
+    )
+
+
+demo.launch(debug=False, show_error=True,share=True)

korean_g2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecdead53c909748aad74bc2f2beeab3e3ad751134db6335e97ea19f879ca407e
+size 16081533

latin_g2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aaa95be1c4a9cb3496879bed7c520886ce1164f89e026f0c54488394e74e8c55
+size 15406141

pipeline_flux_controlnet_inpaint.py

@@ -0,0 +1,1049 @@
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import (
+    CLIPTextModel,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FluxLoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+
+from transformer_flux import FluxTransformer2DModel
+from controlnet_flux import FluxControlNetModel
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers.utils import load_image
+        >>> from diffusers import FluxControlNetPipeline
+        >>> from diffusers import FluxControlNetModel
+
+        >>> controlnet_model = "InstantX/FLUX.1-dev-controlnet-canny-alpha"
+        >>> controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+        >>> pipe = FluxControlNetPipeline.from_pretrained(
+        ...     base_model, controlnet=controlnet, torch_dtype=torch.bfloat16
+        ... )
+        >>> pipe.to("cuda")
+        >>> control_image = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg")
+        >>> control_mask = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg")
+        >>> prompt = "A girl in city, 25 years old, cool, futuristic"
+        >>> image = pipe(
+        ...     prompt,
+        ...     control_image=control_image,
+        ...     controlnet_conditioning_scale=0.6,
+        ...     num_inference_steps=28,
+        ...     guidance_scale=3.5,
+        ... ).images[0]
+        >>> image.save("flux.png")
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FluxControlNetInpaintingPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
+    r"""
+    The Flux pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        controlnet: FluxControlNetModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            scheduler=scheduler,
+            controlnet=controlnet,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels))
+            if hasattr(self, "vae") and self.vae is not None
+            else 16
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_resize=True, do_convert_rgb=True, do_normalize=True)
+        self.mask_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_resize=True,
+            do_convert_grayscale=True,
+            do_normalize=False,
+            do_binarize=True,
+        )
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length
+            if hasattr(self, "tokenizer") and self.tokenizer is not None
+            else 77
+        )
+        self.default_sample_size = 64
+    
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer_2.batch_decode(
+                untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(
+            text_input_ids.to(device), output_hidden_states=False
+        )[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            batch_size * num_images_per_prompt, seq_len, -1
+        )
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device), output_hidden_states=False
+        )
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        max_sequence_length: int = 512,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier-free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                negative prompt to be encoded
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                negative prompt to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is
+                used in all text-encoders
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        if do_classifier_free_guidance:
+            # 处理 negative prompt
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            
+            negative_pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            negative_prompt_embeds = self._get_t5_prompt_embeds(
+                negative_prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+        else:
+            negative_pooled_prompt_embeds = None
+            negative_prompt_embeds = None            
+
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(
+            device=device, dtype=self.text_encoder.dtype
+        )
+
+        return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds,text_ids
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(
+                f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs
+            for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (
+            not isinstance(prompt, str) and not isinstance(prompt, list)
+        ):
+            raise ValueError(
+                f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
+            )
+        elif prompt_2 is not None and (
+            not isinstance(prompt_2, str) and not isinstance(prompt_2, list)
+        ):
+            raise ValueError(
+                f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}"
+            )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}"
+            )
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux._prepare_latent_image_ids
+    @staticmethod
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+        latent_image_ids[..., 1] = (
+            latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+        )
+        latent_image_ids[..., 2] = (
+            latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+        )
+
+        (
+            latent_image_id_height,
+            latent_image_id_width,
+            latent_image_id_channels,
+        ) = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
+        latent_image_ids = latent_image_ids.reshape(
+            batch_size,
+            latent_image_id_height * latent_image_id_width,
+            latent_image_id_channels,
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux._pack_latents
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(
+            batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+        )
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size, (height // 2) * (width // 2), num_channels_latents * 4
+        )
+
+        return latents
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux._unpack_latents
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        height = height // vae_scale_factor
+        width = width // vae_scale_factor
+
+        latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(
+            batch_size, channels // (2 * 2), height * 2, width * 2
+        )
+
+        return latents
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.prepare_latents
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        height = 2 * (int(height) // self.vae_scale_factor)
+        width = 2 * (int(width) // self.vae_scale_factor)
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = self._prepare_latent_image_ids(
+                batch_size, height, width, device, dtype
+            )
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(
+            latents, batch_size, num_channels_latents, height, width
+        )
+
+        latent_image_ids = self._prepare_latent_image_ids(
+            batch_size, height, width, device, dtype
+        )
+
+        return latents, latent_image_ids
+
+    # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image
+    def prepare_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+    ):
+        if isinstance(image, torch.Tensor):
+            pass
+        else:
+            image = self.image_processor.preprocess(image, height=height, width=width)
+
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        return image
+
+    def prepare_image_with_mask(
+        self,
+        image,
+        mask,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance = False,
+    ):
+        # Prepare image
+        if isinstance(image, torch.Tensor):
+            pass
+        else:
+            image = self.image_processor.preprocess(image, height=height, width=width)
+
+        image_batch_size = image.shape[0]
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt
+        image = image.repeat_interleave(repeat_by, dim=0)
+        image = image.to(device=device, dtype=dtype)
+
+        # Prepare mask
+        if isinstance(mask, torch.Tensor):
+            pass
+        else:
+            mask = self.mask_processor.preprocess(mask, height=height, width=width)
+        mask = mask.repeat_interleave(repeat_by, dim=0)
+        mask = mask.to(device=device, dtype=dtype)
+
+        # Get masked image
+        masked_image = image.clone()
+        masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1
+
+        # Encode to latents
+        image_latents = self.vae.encode(masked_image.to(self.vae.dtype)).latent_dist.sample()
+        image_latents = (
+            image_latents - self.vae.config.shift_factor
+        ) * self.vae.config.scaling_factor
+        image_latents = image_latents.to(dtype)
+
+        mask = torch.nn.functional.interpolate(
+            mask, size=(height // self.vae_scale_factor * 2, width // self.vae_scale_factor * 2)
+        )
+        mask = 1 - mask
+
+        control_image = torch.cat([image_latents, mask], dim=1)
+
+        # Pack cond latents
+        packed_control_image = self._pack_latents(
+            control_image,
+            batch_size * num_images_per_prompt,
+            control_image.shape[1],
+            control_image.shape[2],
+            control_image.shape[3],
+        )
+        
+        if do_classifier_free_guidance:
+            packed_control_image = torch.cat([packed_control_image] * 2)
+
+        return packed_control_image, height, width
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 7.0,
+        true_guidance_scale: float = 3.5 ,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        control_image: PipelineImageInput = None,
+        control_mask: PipelineImageInput = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = true_guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None)
+            if self.joint_attention_kwargs is not None
+            else None
+        )
+        (            
+            prompt_embeds,
+            pooled_prompt_embeds,
+            negative_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            text_ids
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            do_classifier_free_guidance = self.do_classifier_free_guidance,
+            negative_prompt = negative_prompt,
+            negative_prompt_2 = negative_prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        
+        # 在 encode_prompt 之后
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim = 0)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim = 0)
+            text_ids = torch.cat([text_ids, text_ids], dim = 0)
+
+        # 3. Prepare control image
+        num_channels_latents = self.transformer.config.in_channels // 4
+        if isinstance(self.controlnet, FluxControlNetModel):
+            control_image, height, width = self.prepare_image_with_mask(
+                image=control_image,
+                mask=control_mask,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=dtype,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        
+        if self.do_classifier_free_guidance:
+            latent_image_ids = torch.cat([latent_image_ids] * 2)
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+
+                # handle guidance
+                if self.transformer.config.guidance_embeds:
+                    guidance = torch.tensor([guidance_scale], device=device)
+                    guidance = guidance.expand(latent_model_input.shape[0])
+                else:
+                    guidance = None
+
+                # controlnet
+                (
+                    controlnet_block_samples,
+                    controlnet_single_block_samples,
+                ) = self.controlnet(
+                    hidden_states=latent_model_input,
+                    controlnet_cond=control_image,
+                    conditioning_scale=controlnet_conditioning_scale,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    controlnet_block_samples=[
+                        sample.to(dtype=self.transformer.dtype)
+                        for sample in controlnet_block_samples
+                    ],
+                    controlnet_single_block_samples=[
+                        sample.to(dtype=self.transformer.dtype)
+                        for sample in controlnet_single_block_samples
+                    ] if controlnet_single_block_samples is not None else controlnet_single_block_samples,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+                
+                # 在生成循环中
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + true_guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, return_dict=False
+                )[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = self._unpack_latents(
+                latents, height, width, self.vae_scale_factor
+            )
+            latents = (
+                latents / self.vae.config.scaling_factor
+            ) + self.vae.config.shift_factor
+            latents = latents.to(self.vae.dtype)
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

requirements.txt

@@ -0,0 +1,15 @@
+easyocr==1.7.2
+openai==1.51.2
+pillow==10.4.0
+opencv-contrib-python==4.10.0.84
+opencv-python==4.10.0.84
+opencv-python-headless==4.10.0.84
+diffusers==0.30.2
+torch
+transformers
+accelerate
+huggingface_hub
+sentencepiece
+ultralytics==8.3.27
+ultralytics-thop==2.0.10
+shapely==2.0.6

speech_bubble_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:200141c44d4153662050531905ec3354d1d6302c28939dbbb019c037250ed610
+size 54809749

transformer_flux.py

@@ -0,0 +1,525 @@
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import (
+    Attention,
+    FluxAttnProcessor2_0,
+    FluxSingleAttnProcessor2_0,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import (
+    AdaLayerNormContinuous,
+    AdaLayerNormZero,
+    AdaLayerNormZeroSingle,
+)
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_version,
+    logging,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import (
+    CombinedTimestepGuidanceTextProjEmbeddings,
+    CombinedTimestepTextProjEmbeddings,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+# YiYi to-do: refactor rope related functions/classes
+def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+    assert dim % 2 == 0, "The dimension must be even."
+
+    scale = torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+
+    batch_size, seq_length = pos.shape
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    cos_out = torch.cos(out)
+    sin_out = torch.sin(out)
+
+    stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
+    out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
+    return out.float()
+
+
+# YiYi to-do: refactor rope related functions/classes
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: List[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+@maybe_allow_in_graph
+class FluxSingleTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        processor = FluxSingleAttnProcessor2_0()
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=processor,
+            qk_norm="rms_norm",
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return hidden_states
+
+
+@maybe_allow_in_graph
+class FluxTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(
+        self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6
+    ):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = FluxAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(
+            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
+        )
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, emb=temb
+        )
+
+        (
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1_context(encoder_hidden_states, emb=temb)
+
+        # Attention.
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = (
+            norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = (
+            norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+            + c_shift_mlp[:, None]
+        )
+
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = (
+            encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+        )
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
+):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False,
+        axes_dims_rope: List[int] = [16, 56, 56],
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = (
+            self.config.num_attention_heads * self.config.attention_head_dim
+        )
+
+        self.pos_embed = EmbedND(
+            dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
+        )
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings
+            if guidance_embeds
+            else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim,
+            pooled_projection_dim=self.config.pooled_projection_dim,
+        )
+
+        self.context_embedder = nn.Linear(
+            self.config.joint_attention_dim, self.inner_dim
+        )
+        self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = nn.Linear(
+            self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
+        )
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if (
+                joint_attention_kwargs is not None
+                and joint_attention_kwargs.get("scale", None) is not None
+            ):
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        image_rotary_emb = self.pos_embed(ids)
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
+                (
+                    encoder_hidden_states,
+                    hidden_states,
+                ) = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(
+                    controlnet_block_samples
+                )
+                interval_control = int(np.ceil(interval_control))
+                hidden_states = (
+                    hidden_states
+                    + controlnet_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(
+                    controlnet_single_block_samples
+                )
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

zh_sim_g2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb678fdef09d651e7763ca551ad790dc89f0b2e3d2a640484330e338fb574c7a
+size 21951421