-
Notifications
You must be signed in to change notification settings - Fork 368
Description
not run on colab t4
from OmniGen import OmniGenPipeline
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import transformers
transformers.logging.set_verbosity_error()
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1", device_map=device)
Text to Image
images = pipe(
prompt="A curly-haired man in a red shirt is drinking tea.",
height=768,
width=512,
guidance_scale=1,
seed=0,
separate_cfg_infer=True,
num_inference_steps=1,
num_images_per_prompt=1,
use_kv_cache=True
)
images[0].save("example_t2i.png") # save output PIL Image
Text to Image
images = pipe(
prompt="A curly-haired man in a red shirt is drinking tea.",
height=768,
width=512,
guidance_scale=1,
seed=0,
separate_cfg_infer=True,
num_inference_steps=1,
num_images_per_prompt=1,
use_kv_cache=True
)
images[0].save("example_t2i.png") # save output PIL Image
TypeError Traceback (most recent call last)
in <cell line: 8>()
6 transformers.logging.set_verbosity_error()
7 device = "cuda" if torch.cuda.is_available() else "cpu"
----> 8 pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1", device_map=device)
9
10 # Text to Image
TypeError: OmniGenPipeline.from_pretrained() got an unexpected keyword argument 'device_map'