JLNTV_IPD/CogVideo
4
0
Fork 0
mirror of https://github.com/THUDM/CogVideo.git synced 2025-12-02 18:52:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

remove alpha channel from init

Browse Source
This commit is contained in:
Stephan Auerhahn 2022-07-27 01:23:00 -07:00
parent a150e9dab1
commit 903c2a98fc
1 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
predict.py
View File

@ -486,7 +486,7 @@ class Predictor(BasePredictor):
description="Run both stages (uncheck to run more quickly and output only a few frames)", default=True
),
use_guidance: bool = Input(description="Use stage 1 guidance (recommended)", default=True),
image_prompt: Path = Input(description="Starting image", default=None)
image_prompt: Path = Input(description="Starting image (optional, prompt has little effect when used)", default=None)
) -> typing.Iterator[Path]:
if translate:
prompt = self.translator.translate(prompt.strip())
@ -500,11 +500,15 @@ class Predictor(BasePredictor):
self.image_prompt = None
if os.path.exists(str(image_prompt)):
try:
Image.open(str(image_prompt))
image = Image.open(str(image_prompt)).convert("RGBA")
# Remove alpha channel if present
bg = Image.new("RGBA", image.size, (255, 255, 255))
image = Image.alpha_composite(bg, image).convert("RGB")
imagefile = f'{tempfile.mkdtemp()}/input.png'
image.save(imagefile, format="png")
self.image_prompt = imagefile
except (FileNotFoundError, UnidentifiedImageError):
logging.debug("Bad image prompt; ignoring") # Is there a better way to input images?
else:
self.image_prompt = str(image_prompt)
self.args.both_stages = both_stages
for file in self.run():