Spaces:
Sleeping
Sleeping
| import os | |
| os.environ ['HF_ENDPOINT'] = 'https://hf-mirror.com' | |
| from ast import main | |
| from numpy import imag | |
| import torch | |
| from diffusers import StableDiffusionPipeline | |
| import os | |
| from PIL import Image | |
| 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_reco_prompt( | |
| caption: str = '', | |
| phrases=[], | |
| boxes=[], | |
| normalize_boxes=True, | |
| image_resolution=512, | |
| num_bins=1000, | |
| ): | |
| """ | |
| method to create ReCo prompt | |
| caption: global caption | |
| phrases: list of regional captions | |
| boxes: list of regional coordinates (unnormalized xyxy) | |
| """ | |
| SOS_token = '<|startoftext|>' | |
| EOS_token = '<|endoftext|>' | |
| box_captions_with_coords = [] | |
| box_captions_with_coords += [caption] | |
| box_captions_with_coords += [EOS_token] | |
| for phrase, box in zip(phrases, boxes): | |
| if normalize_boxes: | |
| box = [float(x) / image_resolution for x in box] | |
| # quantize into bins | |
| quant_x0 = int(round((box[0] * (num_bins - 1)))) | |
| quant_y0 = int(round((box[1] * (num_bins - 1)))) | |
| quant_x1 = int(round((box[2] * (num_bins - 1)))) | |
| quant_y1 = int(round((box[3] * (num_bins - 1)))) | |
| # ReCo format | |
| # Add SOS/EOS before/after regional captions | |
| box_captions_with_coords += [ | |
| f"<bin{str(quant_x0).zfill(3)}>", | |
| f"<bin{str(quant_y0).zfill(3)}>", | |
| f"<bin{str(quant_x1).zfill(3)}>", | |
| f"<bin{str(quant_y1).zfill(3)}>", | |
| SOS_token, | |
| phrase, | |
| EOS_token | |
| ] | |
| text = " ".join(box_captions_with_coords) | |
| return text | |
| def inference_image(pipe, prompt, grounding_instruction, state): | |
| print(prompt) | |
| print(grounding_instruction) | |
| bbox = state['boxes'] | |
| # bbox = state | |
| print(bbox) | |
| bbox = normalize_bbox(bbox, 600, 600) | |
| print(bbox) | |
| objects = [obj for obj in grounding_instruction.split(';') if obj.strip()] | |
| print(objects) | |
| prompt_reco = create_reco_prompt(prompt, objects, bbox, normalize_boxes=False) | |
| print(prompt_reco) | |
| image = pipe(prompt_reco, guidance_scale=4).images[0] | |
| return image | |
| if __name__ == "__main__": | |
| path = '/home/bcy/cache/.cache/huggingface/hub/models--j-min--reco_sd14_coco/snapshots/11a062da5a0a84501047cb19e113f520eb610415' if os.path.isdir('/home/bcy/cache/.cache/huggingface/hub/models--j-min--reco_sd14_coco/snapshots/11a062da5a0a84501047cb19e113f520eb610415') else "CompVis/stable-diffusion-v1-4" | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| "j-min/reco_sd14_coco", | |
| torch_dtype=torch.float16 | |
| ) | |
| pipe = pipe.to("cuda") | |
| # caption = "A box contains six donuts with varying types of glazes and toppings." | |
| # phrases = ["chocolate donut.", "dark vanilla donut.", "donut with sprinkles.", "donut with powdered sugar.", "pink donut.", "brown donut."] | |
| # boxes = [[263.68, 294.912, 380.544, 392.832], [121.344, 265.216, 267.392, 401.92], [391.168, 294.912, 506.368, 381.952], [120.064, 143.872, 268.8, 270.336], [264.192, 132.928, 393.216, 263.68], [386.048, 148.48, 490.688, 259.584]] | |
| # prompt = create_reco_prompt(caption, phrases, boxes) | |
| # print(prompt) | |
| # generated_image = pipe( | |
| # prompt, | |
| # guidance_scale=4).images[0] | |
| # generated_image.save("output1.jpg") | |
| prompt = "a dog and a cat;" | |
| grounding_instruction = "cut dog; big cat;" | |
| bbox = [(136, 252, 280, 455), (284, 205, 480, 500)] | |
| inference_image(pipe, prompt, grounding_instruction, bbox) | |