[BugFix] Download Anima assets for README CI test (#7928)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
This commit is contained in:
1 parent
d38566ae80
commit
b4737d3066
4 files changed
+41
-8
No files matched your search
@@ -6,9 +6,12 @@ Offline inference tests: text-to-image.
|
||||
See examples/offline_inference/text_to_image/README.md
|
||||
"""
|
||||
|
||||
import json
|
||||
import shlex
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from huggingface_hub import hf_hub_download, snapshot_download
|
||||
|
||||
from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
|
||||
from tests.helpers.assertions import assert_image_valid
|
||||
@@ -37,12 +40,38 @@ def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool,
|
||||
README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
|
||||
|
||||
|
||||
def _prepare_anima_snippet(snippet: ReadmeSnippet) -> ReadmeSnippet:
|
||||
if "--model-class-name AnimaPipeline" not in snippet.code:
|
||||
return snippet
|
||||
|
||||
assert "/path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors" in snippet.code, (
|
||||
"Anima README checkpoint placeholder changed; update _prepare_anima_snippet to match"
|
||||
)
|
||||
|
||||
checkpoint = hf_hub_download(
|
||||
repo_id="circlestone-labs/Anima",
|
||||
filename="split_files/diffusion_models/anima-base-v1.0.safetensors",
|
||||
)
|
||||
components = snapshot_download(
|
||||
repo_id="circlestone-labs/Anima-Base-v1.0-Diffusers",
|
||||
allow_patterns=["text_encoder/*", "vae/*", "tokenizer/*", "t5_tokenizer/*", "scheduler/*"],
|
||||
)
|
||||
argv = shlex.split(snippet.code)
|
||||
argv[argv.index("--model") + 1] = checkpoint
|
||||
custom_args_index = argv.index("--custom-pipeline-args") + 1
|
||||
custom_args = json.loads(argv[custom_args_index])
|
||||
custom_args["components_path"] = components
|
||||
argv[custom_args_index] = json.dumps(custom_args)
|
||||
return snippet._replace(code=shlex.join(argv))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
|
||||
def test_text_to_image(snippet: ReadmeSnippet, example_runner: ExampleRunner):
|
||||
should_skip, reason = snippet.skip
|
||||
if should_skip:
|
||||
pytest.skip(reason)
|
||||
|
||||
snippet = _prepare_anima_snippet(snippet)
|
||||
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
|
||||
for asset in result.assets:
|
||||
assert_image_valid(asset)
|
||||
Reference in new issue
Block a user