model_bbox/MIGC/inference_single_image.py

import os
import sys
import torch

migc_path = os.path.dirname(os.path.abspath(__file__))
print(migc_path)
if migc_path not in sys.path:
    sys.path.append(migc_path)
import yaml
from diffusers import EulerDiscreteScheduler
from migc.migc_utils import seed_everything
from migc.migc_pipeline import StableDiffusionMIGCPipeline, MIGCProcessor, AttentionStore

def normalize_bbox(bboxes, img_width, img_height):
    normalized_bboxes = []
    for box in bboxes:
        x_min, y_min, x_max, y_max = box
        
        x_min = x_min / img_width
        y_min = y_min / img_height
        x_max = x_max / img_width
        y_max = y_max / img_height
        
        normalized_bboxes.append([x_min, y_min, x_max, y_max])
    
    return [normalized_bboxes]

def create_simple_prompt(input_str):
    # 先将输入字符串按分号分割，并去掉空字符串
    objects = [obj for obj in input_str.split(';') if obj.strip()]
    
    # 创建详细描述字符串
    prompt_description = "masterpiece, best quality, " + ", ".join(objects)
    
    # 创建最终结构
    prompt_final = [[prompt_description] + objects]
    
    return prompt_final


def inference_single_image(prompt, grounding_instruction, state):
    print(prompt)
    print(grounding_instruction)
    bbox = state['boxes']
    print(bbox)
    bbox = normalize_bbox(bbox, 600, 600)
    print(bbox)
    simple_prompt = create_simple_prompt(grounding_instruction)
    print(simple_prompt)
    migc_ckpt_path = 'pretrained_weights/MIGC_SD14.ckpt'
    migc_ckpt_path_all = os.path.join(migc_path, migc_ckpt_path)
    print(migc_ckpt_path_all)
    assert os.path.isfile(migc_ckpt_path_all), "Please download the ckpt of migc and put it in the pretrained_weighrs/ folder!"


    sd1x_path = '/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b' if os.path.isdir('/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b') else "CompVis/stable-diffusion-v1-4"
    # MIGC is a plug-and-play controller.
    # You can go to https://civitai.com/search/models?baseModel=SD%201.4&baseModel=SD%201.5&sortBy=models_v5 find a base model with better generation ability to achieve better creations.
    
    # Construct MIGC pipeline
    pipe = StableDiffusionMIGCPipeline.from_pretrained(
        sd1x_path)
    pipe.attention_store = AttentionStore()
    from migc.migc_utils import load_migc
    load_migc(pipe.unet , pipe.attention_store,
            migc_ckpt_path_all, attn_processor=MIGCProcessor)
    pipe = pipe.to("cuda")
    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
    

    # prompt_final = [['masterpiece, best quality,black colored ball,gray colored cat,white colored  bed,\
    #                  green colored plant,red colored teddy bear,blue colored wall,brown colored vase,orange colored book,\
    #                  yellow colored hat', 'black colored ball', 'gray colored cat', 'white colored  bed', 'green colored plant', \
    #                     'red colored teddy bear', 'blue colored wall', 'brown colored vase', 'orange colored book', 'yellow colored hat']]
    
    # bboxes = [[[0.3125, 0.609375, 0.625, 0.875], [0.5625, 0.171875, 0.984375, 0.6875], \
    #            [0.0, 0.265625, 0.984375, 0.984375], [0.0, 0.015625, 0.21875, 0.328125], \
    #             [0.171875, 0.109375, 0.546875, 0.515625], [0.234375, 0.0, 1.0, 0.3125], \
    #                 [0.71875, 0.625, 0.953125, 0.921875], [0.0625, 0.484375, 0.359375, 0.8125], \
    #                     [0.609375, 0.09375, 0.90625, 0.28125]]]
    negative_prompt = 'worst quality, low quality, bad anatomy, watermark, text, blurry'
    seed = 7351007268695528845
    seed_everything(seed)
    print("Start inference: ")
    image = pipe(simple_prompt, bbox, num_inference_steps=50, guidance_scale=7.5, 
                    MIGCsteps=25, aug_phase_with_and=False, negative_prompt=negative_prompt).images[0]
    return image




# def MIGC_Pipe():
#     migc_ckpt_path = 'pretrained_weights/MIGC_SD14.ckpt'
#     migc_ckpt_path_all = os.path.join(migc_path, migc_ckpt_path)
#     print(migc_ckpt_path_all)
#     assert os.path.isfile(migc_ckpt_path_all), "Please download the ckpt of migc and put it in the pretrained_weighrs/ folder!"
#     sd1x_path = '/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b' if os.path.isdir('/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b') else "CompVis/stable-diffusion-v1-4"
#     pipe = StableDiffusionMIGCPipeline.from_pretrained(
#         sd1x_path)
#     pipe.attention_store = AttentionStore()
#     from migc.migc_utils import load_migc
#     load_migc(pipe.unet , pipe.attention_store,
#             migc_ckpt_path_all, attn_processor=MIGCProcessor)
#     pipe = pipe.to("cuda")
#     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
#     return pipe

def MIGC_Pipe():
    migc_ckpt_path = 'pretrained_weights/MIGC_SD14.ckpt'
    migc_ckpt_path_all = os.path.join(migc_path, migc_ckpt_path)
    print(f"加载 MIGC 权重文件路径: {migc_ckpt_path_all}")
    
    assert os.path.isfile(migc_ckpt_path_all), f"请下载 MIGC 的 ckpt 文件并将其放在 'pretrained_weights/' 文件夹中: {migc_ckpt_path_all}"
    
    sd1x_path = '/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b' if os.path.isdir('/share/bcy/cache/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b') else "CompVis/stable-diffusion-v1-4"
    print(f"加载 StableDiffusion 模型: {sd1x_path}")
    
    # 加载 StableDiffusionMIGCPipeline
    print("load sd:")
    pipe = StableDiffusionMIGCPipeline.from_pretrained(sd1x_path)
    pipe.attention_store = AttentionStore()
    
    # 导入并加载 MIGC 权重
    print("load migc")
    from migc.migc_utils import load_migc
    load_migc(pipe.unet, pipe.attention_store, migc_ckpt_path_all, attn_processor=MIGCProcessor)
    
    # 确保模型和 attention_store 被正确加载
    assert pipe.unet is not None, "unet 模型未正确加载！"
    assert pipe.attention_store is not None, "attention_store 未正确加载！"
    
    # 转移到 CUDA
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("使用 CUDA 设备")
    else:
        device = torch.device("cpu")
        print("使用 CPU")
    
    pipe = pipe.to(device)
    
    # 设置调度器
    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
    
    return pipe


def create_simple_prompt(input_str):
    # 先将输入字符串按分号分割，并去掉空字符串
    objects = [obj for obj in input_str.split(';') if obj.strip()]
    
    # 创建详细描述字符串
    prompt_description = "masterpiece, best quality, " + ", ".join(objects)
    
    # 创建最终结构
    prompt_final = [[prompt_description] + objects]
    
    return prompt_final


def inference_image(pipe, prompt, grounding_instruction, state):
    print(prompt)
    print(grounding_instruction)
    bbox = state['boxes']
    print(bbox)
    bbox = normalize_bbox(bbox, 600, 600)
    print(bbox)
    simple_prompt = create_simple_prompt(grounding_instruction)
    print(simple_prompt)
    negative_prompt = 'worst quality, low quality, bad anatomy, watermark, text, blurry'
    seed = 7351007268695528845
    seed_everything(seed)
    print("Start inference: ")
    image = pipe(simple_prompt, bbox, num_inference_steps=50, guidance_scale=7.5, 
                    MIGCsteps=25, aug_phase_with_and=False, negative_prompt=negative_prompt).images[0]
    return image



if __name__ == "__main__":
    prompt_final = [['masterpiece, best quality,black colored ball,gray colored cat,white colored  bed,\
                     green colored plant,red colored teddy bear,blue colored wall,brown colored vase,orange colored book,\
                     yellow colored hat', 'black colored ball', 'gray colored cat', 'white colored  bed', 'green colored plant', \
                        'red colored teddy bear', 'blue colored wall', 'brown colored vase', 'orange colored book', 'yellow colored hat']]
    
    bboxes = [[[0.3125, 0.609375, 0.625, 0.875], [0.5625, 0.171875, 0.984375, 0.6875], \
               [0.0, 0.265625, 0.984375, 0.984375], [0.0, 0.015625, 0.21875, 0.328125], \
                [0.171875, 0.109375, 0.546875, 0.515625], [0.234375, 0.0, 1.0, 0.3125], \
                    [0.71875, 0.625, 0.953125, 0.921875], [0.0625, 0.484375, 0.359375, 0.8125], \
                        [0.609375, 0.09375, 0.90625, 0.28125]]]
    image = inference_single_image("a cat", prompt_final, bboxes)
    image.save("output.png")
    print("done")