temp commit for entity control

examples/EntityControl/entity_control_flux.py

@@ -0,0 +1,54 @@
+import torch
+from diffsynth import ModelManager, FluxImagePipeline, download_customized_models, FluxImageLoraPipeline
+from examples.EntityControl.utils import visualize_masks
+import os
+import json
+from PIL import Image
+
+# lora_path = download_customized_models(
+#     model_id="DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1",
+#     origin_file_path="merged_lora.safetensors",
+#     local_dir="models/lora"
+# )[0]
+
+lora_path = '/root/model_bf16.safetensors'
+model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda")
+model_manager.load_models([
+    "t2i_models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
+    "t2i_models/FLUX/FLUX.1-dev/text_encoder_2",
+    "t2i_models/FLUX/FLUX.1-dev/ae.safetensors",
+    "t2i_models/FLUX/FLUX.1-dev/flux1-dev.safetensors"
+])
+model_manager.load_lora(lora_path, lora_alpha=1.)
+
+pipe = FluxImagePipeline.from_model_manager(model_manager)
+
+mask_dir = '/mnt/nas1/zhanghong/DiffSynth-Studio/workdirs/tmp_mask'
+image_shape = 1024
+guidance = 3.5
+cfg = 3.0
+negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+names = ['row_2_1']
+seeds = [0]
+# use this to apply regional attention in negative prompt prediction for better results with more time
+use_seperated_negtive_prompt = False
+for name, seed in zip(names, seeds):
+    out_dir = f'workdirs/entity_control/{name}'
+    os.makedirs(out_dir, exist_ok=True)
+    cur_dir = os.path.join(mask_dir, name)
+    metas = json.load(open(os.path.join(mask_dir, name, 'prompts.json')))
+    for seed in range(3, 10):
+        prompt = metas['global_prompt']
+        mask_prompts = metas['mask_prompts']
+        masks = [Image.open(os.path.join(mask_dir, name, f"{mask_idx}.png")).resize((image_shape, image_shape), resample=Image.NEAREST) for mask_idx in range(len(mask_prompts))]
+        torch.manual_seed(seed)
+        image = pipe(
+            prompt=prompt,
+            cfg_scale=cfg,
+            negative_prompt=negative_prompt,
+            num_inference_steps=50, embedded_guidance=guidance, height=image_shape, width=image_shape,
+            entity_prompts=mask_prompts, entity_masks=masks,
+            use_seperated_negtive_prompt=use_seperated_negtive_prompt
+        )
+        use_sep = f'_sepneg' if use_seperated_negtive_prompt else ''
+        visualize_masks(image, masks, mask_prompts, os.path.join(out_dir, f"{name}_{seed}{use_sep}.png"))

examples/EntityControl/entity_inpaint_flux.py

@@ -0,0 +1,59 @@
+import torch
+from diffsynth import ModelManager, FluxImagePipeline, download_customized_models, FluxImageLoraPipeline
+from examples.EntityControl.utils import visualize_masks
+import os
+import json
+from PIL import Image
+
+# lora_path = download_customized_models(
+#     model_id="DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1",
+#     origin_file_path="merged_lora.safetensors",
+#     local_dir="models/lora"
+# )[0]
+
+lora_path = '/root/model_bf16.safetensors'
+model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda")
+model_manager.load_models([
+    "t2i_models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
+    "t2i_models/FLUX/FLUX.1-dev/text_encoder_2",
+    "t2i_models/FLUX/FLUX.1-dev/ae.safetensors",
+    "t2i_models/FLUX/FLUX.1-dev/flux1-dev.safetensors"
+])
+model_manager.load_lora(lora_path, lora_alpha=1.)
+
+pipe = FluxImagePipeline.from_model_manager(model_manager)
+
+mask_dir = '/mnt/nas1/zhanghong/DiffSynth-Studio/workdirs/tmp_mask'
+image_shape = 1024
+guidance = 3.5
+cfg = 3.0
+negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+names = ['inpaint2']
+seeds = [0]
+use_seperated_negtive_prompt = False
+for name, seed in zip(names, seeds):
+    out_dir = f'workdirs/paper_app/inpaint/elc/{name}'
+    os.makedirs(out_dir, exist_ok=True)
+    cur_dir = os.path.join(mask_dir, name)
+    metas = json.load(open(os.path.join(mask_dir, name, 'prompts.json')))
+    inpaint_input = Image.open(os.path.join(cur_dir, 'input.png')).convert('RGB')
+    prompt = metas['global_prompt']
+    prompt = 'A person with a dog walking on the cloud. A rocket in the sky'
+    mask_prompts = metas['mask_prompts']
+    masks = [Image.open(os.path.join(mask_dir, name, f"{mask_idx}.png")).resize((image_shape, image_shape), resample=Image.NEAREST) for mask_idx in range(len(mask_prompts))]
+    torch.manual_seed(seed)
+    image = pipe(
+        prompt=prompt,
+        cfg_scale=cfg,
+        negative_prompt=negative_prompt,
+        num_inference_steps=50,
+        embedded_guidance=guidance,
+        height=image_shape,
+        width=image_shape,
+        entity_prompts=mask_prompts,
+        entity_masks=masks,
+        inpaint_input=inpaint_input,
+        use_seperated_negtive_prompt=use_seperated_negtive_prompt,
+    )
+    use_sep = f'_sepneg' if use_seperated_negtive_prompt else ''
+    visualize_masks(image, masks, mask_prompts, os.path.join(out_dir, f"{name}_{seed}{use_sep}.png"))

examples/EntityControl/utils.py

@@ -0,0 +1,59 @@
+from PIL import Image, ImageDraw, ImageFont
+import random
+
+def visualize_masks(image, masks, mask_prompts, output_path, font_size=35, use_random_colors=False):
+    # Create a blank image for overlays
+    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
+    
+    colors = [
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+    ]
+    # Generate random colors for each mask
+    if use_random_colors:
+        colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 80) for _ in range(len(masks))]
+    
+    # Font settings
+    try:
+        font = ImageFont.truetype("arial", font_size)  # Adjust as needed
+    except IOError:
+        font = ImageFont.load_default(font_size)
+
+    # Overlay each mask onto the overlay image
+    for mask, mask_prompt, color in zip(masks, mask_prompts, colors):
+        # Convert mask to RGBA mode
+        mask_rgba = mask.convert('RGBA')
+        mask_data = mask_rgba.getdata()
+        new_data = [(color if item[:3] == (255, 255, 255) else (0, 0, 0, 0)) for item in mask_data]
+        mask_rgba.putdata(new_data)
+
+        # Draw the mask prompt text on the mask
+        draw = ImageDraw.Draw(mask_rgba)
+        mask_bbox = mask.getbbox()  # Get the bounding box of the mask
+        text_position = (mask_bbox[0] + 10, mask_bbox[1] + 10)  # Adjust text position based on mask position
+        draw.text(text_position, mask_prompt, fill=(255, 255, 255, 255), font=font)
+
+        # Alpha composite the overlay with this mask
+        overlay = Image.alpha_composite(overlay, mask_rgba)
+    
+    # Composite the overlay onto the original image
+    result = Image.alpha_composite(image.convert('RGBA'), overlay)
+    
+    # Save or display the resulting image
+    result.save(output_path)
+
+    return result