deepseekmirror/Janus
1
0
Fork 0
mirror of https://github.com/deepseek-ai/Janus.git synced 2025-04-19 10:09:00 -04:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 7628bb8174 into 1daa72fa40

Browse Source
This commit is contained in:
Vadzim Belski 2025-02-01 12:56:27 +02:00 committed by GitHub
commit df7eb21b84
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 82 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

119
demo/app_januspro.py
View File

@ -10,35 +10,48 @@ import os
import time import time
# import spaces # Import spaces for ZeroGPU compatibility # import spaces # Import spaces for ZeroGPU compatibility
# 1. Load model and processor
# Load model and processor
model_path = "deepseek-ai/Janus-Pro-7B" model_path = "deepseek-ai/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path) config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config language_config = config.language_config
language_config._attn_implementation = 'eager' language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
language_config=language_config, # 2. Check for MPS availability, otherwise fall back to CPU
trust_remote_code=True) if torch.backends.mps.is_available():
if torch.cuda.is_available(): device = torch.device('mps')
vl_gpt = vl_gpt.to(torch.bfloat16).cuda() print("Using MPS (Metal Performance Shaders)")
else: else:
vl_gpt = vl_gpt.to(torch.float16) device = torch.device('cpu')
print("Using CPU")
# 3. Load model in float32, then move to MPS or CPU
vl_gpt = AutoModelForCausalLM.from_pretrained(
model_path,
language_config=language_config,
trust_remote_code=True,
torch_dtype=torch.float32 # Attempt to load everything in float32
)
vl_gpt = vl_gpt.float().to(device)
for name, module in vl_gpt.named_modules():
if isinstance(module, torch.nn.Module):
module.float()
vl_gpt.to(device)
vl_chat_processor = VLChatProcessor.from_pretrained(model_path) vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu' cuda_device = device
@torch.inference_mode() @torch.inference_mode()
# @spaces.GPU(duration=120)
# Multimodal Understanding function
def multimodal_understanding(image, question, seed, top_p, temperature): def multimodal_understanding(image, question, seed, top_p, temperature):
# Clear CUDA cache before generating # Clear cache if using CUDA
if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
# set seed # set seed
torch.manual_seed(seed) torch.manual_seed(seed)
np.random.seed(seed) np.random.seed(seed)
torch.cuda.manual_seed(seed)
conversation = [ conversation = [
{ {
@ -50,11 +63,17 @@ def multimodal_understanding(image, question, seed, top_p, temperature):
] ]
pil_images = [Image.fromarray(image)] pil_images = [Image.fromarray(image)]
prepare_inputs = vl_chat_processor( prepare_inputs = vl_chat_processor(
conversations=conversation, images=pil_images, force_batchify=True conversations=conversation, images=pil_images, force_batchify=True
).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16) ).to(cuda_device, dtype=torch.float32)
# Option 1: Just remove the autocast context entirely
# with torch.autocast("mps", dtype=torch.float32"):
# inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
# OR Option 2: explicitly disable autocast
with torch.autocast("mps", enabled=False):
inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs) inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
outputs = vl_gpt.language_model.generate( outputs = vl_gpt.language_model.generate(
@ -64,7 +83,7 @@ def multimodal_understanding(image, question, seed, top_p, temperature):
bos_token_id=tokenizer.bos_token_id, bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id,
max_new_tokens=512, max_new_tokens=512,
do_sample=False if temperature == 0 else True, do_sample=(temperature != 0),
use_cache=True, use_cache=True,
temperature=temperature, temperature=temperature,
top_p=top_p, top_p=top_p,
@ -73,7 +92,6 @@ def multimodal_understanding(image, question, seed, top_p, temperature):
answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True) answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
return answer return answer
def generate(input_ids, def generate(input_ids,
width, width,
height, height,
@ -82,7 +100,8 @@ def generate(input_ids,
cfg_weight: float = 5, cfg_weight: float = 5,
image_token_num_per_image: int = 576, image_token_num_per_image: int = 576,
patch_size: int = 16): patch_size: int = 16):
# Clear CUDA cache before generating # Clear cache if using CUDA
if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).to(cuda_device) tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).to(cuda_device)
@ -113,10 +132,10 @@ def generate(input_ids,
img_embeds = vl_gpt.prepare_gen_img_embeds(next_token) img_embeds = vl_gpt.prepare_gen_img_embeds(next_token)
inputs_embeds = img_embeds.unsqueeze(dim=1) inputs_embeds = img_embeds.unsqueeze(dim=1)
patches = vl_gpt.gen_vision_model.decode_code(
generated_tokens.to(dtype=torch.int),
patches = vl_gpt.gen_vision_model.decode_code(generated_tokens.to(dtype=torch.int), shape=[parallel_size, 8, width // patch_size, height // patch_size]
shape=[parallel_size, 8, width // patch_size, height // patch_size]) )
return generated_tokens.to(dtype=torch.int), patches return generated_tokens.to(dtype=torch.int), patches
@ -126,54 +145,60 @@ def unpack(dec, width, height, parallel_size=5):
visual_img = np.zeros((parallel_size, width, height, 3), dtype=np.uint8) visual_img = np.zeros((parallel_size, width, height, 3), dtype=np.uint8)
visual_img[:, :, :] = dec visual_img[:, :, :] = dec
return visual_img return visual_img
@torch.inference_mode() @torch.inference_mode()
# @spaces.GPU(duration=120) # Specify a duration to avoid timeout
def generate_image(prompt, def generate_image(prompt,
seed=None, seed=None,
guidance=5, guidance=5,
t2i_temperature=1.0): t2i_temperature=1.0):
# Clear CUDA cache and avoid tracking gradients if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
# Set the seed for reproducible results # Set the seed for reproducible results
if seed is not None: if seed is not None:
torch.manual_seed(seed) torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed) np.random.seed(seed)
width = 384 width = 384
height = 384 height = 384
parallel_size = 5 parallel_size = 5
with torch.no_grad(): with torch.no_grad():
messages = [{'role': '<|User|>', 'content': prompt}, messages = [
{'role': '<|Assistant|>', 'content': ''}] {'role': '<|User|>', 'content': prompt},
text = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(conversations=messages, {'role': '<|Assistant|>', 'content': ''}
]
text = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
conversations=messages,
sft_format=vl_chat_processor.sft_format, sft_format=vl_chat_processor.sft_format,
system_prompt='') system_prompt=''
)
text = text + vl_chat_processor.image_start_tag text = text + vl_chat_processor.image_start_tag
input_ids = torch.LongTensor(tokenizer.encode(text)) input_ids = torch.LongTensor(tokenizer.encode(text))
output, patches = generate(input_ids, output, patches = generate(
input_ids,
width // 16 * 16, width // 16 * 16,
height // 16 * 16, height // 16 * 16,
cfg_weight=guidance, cfg_weight=guidance,
parallel_size=parallel_size, parallel_size=parallel_size,
temperature=t2i_temperature) temperature=t2i_temperature
images = unpack(patches, )
images = unpack(
patches,
width // 16 * 16, width // 16 * 16,
height // 16 * 16, height // 16 * 16,
parallel_size=parallel_size) parallel_size=parallel_size
)
return [Image.fromarray(images[i]).resize((768, 768), Image.LANCZOS) for i in range(parallel_size)] return [Image.fromarray(images[i]).resize((768, 768), Image.LANCZOS)
for i in range(parallel_size)]
# Gradio interface # Gradio interface
with gr.Blocks() as demo: with gr.Blocks() as demo:
gr.Markdown(value="# Multimodal Understanding") gr.Markdown("# Multimodal Understanding")
with gr.Row(): with gr.Row():
image_input = gr.Image() image_input = gr.Image()
with gr.Column(): with gr.Column():
@ -188,22 +213,13 @@ with gr.Blocks() as demo:
examples_inpainting = gr.Examples( examples_inpainting = gr.Examples(
label="Multimodal Understanding examples", label="Multimodal Understanding examples",
examples=[ examples=[
[ ["explain this meme", "images/doge.png"],
"explain this meme", ["Convert the formula into latex code.", "images/equation.png"],
"images/doge.png",
],
[
"Convert the formula into latex code.",
"images/equation.png",
],
], ],
inputs=[question_input, image_input], inputs=[question_input, image_input],
) )
gr.Markdown("# Text-to-Image Generation")
gr.Markdown(value="# Text-to-Image Generation")
with gr.Row(): with gr.Row():
cfg_weight_input = gr.Slider(minimum=1, maximum=10, value=5, step=0.5, label="CFG Weight") cfg_weight_input = gr.Slider(minimum=1, maximum=10, value=5, step=0.5, label="CFG Weight")
@ -224,7 +240,7 @@ with gr.Blocks() as demo:
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
"A glass of red wine on a reflective surface.", "A glass of red wine on a reflective surface.",
"A cute and adorable baby fox with big brown eyes, autumn leaves in the background enchanting,immortal,fluffy, shiny mane,Petals,fairyism,unreal engine 5 and Octane Render,highly detailed, photorealistic, cinematic, natural colors.", "A cute and adorable baby fox with big brown eyes, autumn leaves in the background enchanting,immortal,fluffy, shiny mane,Petals,fairyism,unreal engine 5 and Octane Render,highly detailed, photorealistic, cinematic, natural colors.",
"The image features an intricately designed eye set against a circular backdrop adorned with ornate swirl patterns that evoke both realism and surrealism. At the center of attention is a strikingly vivid blue iris surrounded by delicate veins radiating outward from the pupil to create depth and intensity. The eyelashes are long and dark, casting subtle shadows on the skin around them which appears smooth yet slightly textured as if aged or weathered over time.\n\nAbove the eye, there's a stone-like structure resembling part of classical architecture, adding layers of mystery and timeless elegance to the composition. This architectural element contrasts sharply but harmoniously with the organic curves surrounding it. Below the eye lies another decorative motif reminiscent of baroque artistry, further enhancing the overall sense of eternity encapsulated within each meticulously crafted detail. \n\nOverall, the atmosphere exudes a mysterious aura intertwined seamlessly with elements suggesting timelessness, achieved through the juxtaposition of realistic textures and surreal artistic flourishes. Each component\u2014from the intricate designs framing the eye to the ancient-looking stone piece above\u2014contributes uniquely towards creating a visually captivating tableau imbued with enigmatic allure.", "The image features an intricately designed eye set against a circular backdrop adorned with ornate swirl patterns...",
], ],
inputs=prompt_input, inputs=prompt_input,
) )
@ -241,5 +257,4 @@ with gr.Blocks() as demo:
outputs=image_output outputs=image_output
) )
demo.launch(share=True) demo.launch(share=False)
# demo.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", server_port=37906, root_path="/path")

2
requirements.txt
View File

@ -16,4 +16,4 @@ tqdm==4.64.0
colorama==0.4.5 colorama==0.4.5
Pygments==2.12.0 Pygments==2.12.0
markdown==3.4.1 markdown==3.4.1
SentencePiece==0.1.96 SentencePiece==0.1.99