Views
No views yet
1import torch
2from diffusers import StableDiffusionPipeline
3
4pipe = StableDiffusionPipeline.from_pretrained(
5 "j-min/reco_sd14_coco",
6 torch_dtype=torch.float16
7)
8pipe = pipe.to("cuda")
9
10
11prompt = "A box contains six donuts with varying types of glazes and toppings. <|endoftext|> <bin514> <bin575> <bin741> <bin765> <|startoftext|> chocolate donut. <|endoftext|> <bin237> <bin517> <bin520> <bin784> <|startoftext|> dark vanilla donut. <|endoftext|> <bin763> <bin575> <bin988> <bin745> <|startoftext|> donut with sprinkles. <|endoftext|> <bin234> <bin281> <bin524> <bin527> <|startoftext|> donut with powdered sugar. <|endoftext|> <bin515> <bin259> <bin767> <bin514> <|startoftext|> pink donut. <|endoftext|> <bin753> <bin289> <bin958> <bin506> <|startoftext|> brown donut. <|endoftext|>"
12generated_image = pipe(
13 prompt,
14 guidance_scale=4).images[0]
15generated_image1def create_reco_prompt(
2 caption: str = '',
3 phrases=[],
4 boxes=[],
5 normalize_boxes=True,
6 image_resolution=512,
7 num_bins=1000,
8 ):
9 """
10 method to create ReCo prompt
11
12 caption: global caption
13 phrases: list of regional captions
14 boxes: list of regional coordinates (unnormalized xyxy)
15 """
16
17 SOS_token = '<|startoftext|>'
18 EOS_token = '<|endoftext|>'
19
20 box_captions_with_coords = []
21
22 box_captions_with_coords += [caption]
23 box_captions_with_coords += [EOS_token]
24
25 for phrase, box in zip(phrases, boxes):
26
27 if normalize_boxes:
28 box = [float(x) / image_resolution for x in box]
29
30 # quantize into bins
31 quant_x0 = int(round((box[0] * (num_bins - 1))))
32 quant_y0 = int(round((box[1] * (num_bins - 1))))
33 quant_x1 = int(round((box[2] * (num_bins - 1))))
34 quant_y1 = int(round((box[3] * (num_bins - 1))))
35
36 # ReCo format
37 # Add SOS/EOS before/after regional captions
38 box_captions_with_coords += [
39 f"<bin{str(quant_x0).zfill(3)}>",
40 f"<bin{str(quant_y0).zfill(3)}>",
41 f"<bin{str(quant_x1).zfill(3)}>",
42 f"<bin{str(quant_y1).zfill(3)}>",
43 SOS_token,
44 phrase,
45 EOS_token
46 ]
47
48 text = " ".join(box_captions_with_coords)
49 return text
50
51caption = "a photo of bus and boat; boat is left to bus."
52phrases = ["a photo of a bus.", "a photo of a boat."]
53boxes = [[0.702, 0.404, 0.927, 0.601], [0.154, 0.383, 0.311, 0.487]]
54prompt = create_reco_prompt(caption, phrases, boxes, normalize_boxes=False)
55prompt
56>>> 'a photo of bus and boat; boat is left to bus. <|endoftext|> <bin701> <bin404> <bin926> <bin600> <|startoftext|> a photo of a bus. <|endoftext|> <bin154> <bin383> <bin311> <bin487> <|startoftext|> a photo of a boat. <|endoftext|>'
57
58
59caption = "A box contains six donuts with varying types of glazes and toppings."
60phrases = ["chocolate donut.", "dark vanilla donut.", "donut with sprinkles.", "donut with powdered sugar.", "pink donut.", "brown donut."]
61boxes = [[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]]
62prompt = create_reco_prompt(caption, phrases, boxes)
63prompt
64>>> 'A box contains six donuts with varying types of glazes and toppings. <|endoftext|> <bin514> <bin575> <bin743> <bin766> <|startoftext|> chocolate donut. <|endoftext|> <bin237> <bin517> <bin522> <bin784> <|startoftext|> dark vanilla donut. <|endoftext|> <bin763> <bin575> <bin988> <bin745> <|startoftext|> donut with sprinkles. <|endoftext|> <bin234> <bin281> <bin524> <bin527> <|startoftext|> donut with powdered sugar. <|endoftext|> <bin515> <bin259> <bin767> <bin514> <|startoftext|> pink donut. <|endoftext|> <bin753> <bin290> <bin957> <bin506> <|startoftext|> brown donut. <|endoftext|>'