Views
No views yet


1from PIL import Image
2from transformers import AutoTokenizer, AutoModel, AutoImageProcessor, AutoModelForCausalLM
3from transformers.generation.configuration_utils import GenerationConfig
4from transformers.generation import LogitsProcessorList, PrefixConstrainedLogitsProcessor, UnbatchedClassifierFreeGuidanceLogitsProcessor
5import torch
6
7import sys
8sys.path.append(PATH_TO_BAAI_Emu3-Gen_MODEL)
9from processing_emu3 import Emu3Processor
10
11# model path
12EMU_HUB = "BAAI/Emu3-Gen"
13VQ_HUB = "BAAI/Emu3-VisionTokenizer"
14
15# prepare model and processor
16model = AutoModelForCausalLM.from_pretrained(
17 EMU_HUB,
18 device_map="cuda:0",
19 torch_dtype=torch.bfloat16,
20 attn_implementation="flash_attention_2",
21 trust_remote_code=True,
22)
23
24tokenizer = AutoTokenizer.from_pretrained(EMU_HUB, trust_remote_code=True, padding_side="left")
25image_processor = AutoImageProcessor.from_pretrained(VQ_HUB, trust_remote_code=True)
26image_tokenizer = AutoModel.from_pretrained(VQ_HUB, device_map="cuda:0", trust_remote_code=True).eval()
27processor = Emu3Processor(image_processor, image_tokenizer, tokenizer)
28
29# prepare input
30POSITIVE_PROMPT = " masterpiece, film grained, best quality."
31NEGATIVE_PROMPT = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry."
32
33classifier_free_guidance = 3.0
34prompt = "a portrait of young girl."
35prompt += POSITIVE_PROMPT
36
37kwargs = dict(
38 mode='G',
39 ratio="1:1",
40 image_area=model.config.image_area,
41 return_tensors="pt",
42 padding="longest",
43)
44pos_inputs = processor(text=prompt, **kwargs)
45neg_inputs = processor(text=NEGATIVE_PROMPT, **kwargs)
46
47# prepare hyper parameters
48GENERATION_CONFIG = GenerationConfig(
49 use_cache=True,
50 eos_token_id=model.config.eos_token_id,
51 pad_token_id=model.config.pad_token_id,
52 max_new_tokens=40960,
53 do_sample=True,
54 top_k=2048,
55)
56
57h = pos_inputs.image_size[:, 0]
58w = pos_inputs.image_size[:, 1]
59constrained_fn = processor.build_prefix_constrained_fn(h, w)
60logits_processor = LogitsProcessorList([
61 UnbatchedClassifierFreeGuidanceLogitsProcessor(
62 classifier_free_guidance,
63 model,
64 unconditional_ids=neg_inputs.input_ids.to("cuda:0"),
65 ),
66 PrefixConstrainedLogitsProcessor(
67 constrained_fn ,
68 num_beams=1,
69 ),
70])
71
72# generate
73outputs = model.generate(
74 pos_inputs.input_ids.to("cuda:0"),
75 GENERATION_CONFIG,
76 logits_processor=logits_processor,
77 attention_mask=pos_inputs.attention_mask.to("cuda:0"),
78)
79
80mm_list = processor.decode(outputs[0])
81for idx, im in enumerate(mm_list):
82 if not isinstance(im, Image.Image):
83 continue
84 im.save(f"result_{idx}.png")