forked from replicate/replicate-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_training.py
More file actions
84 lines (64 loc) · 2.53 KB
/
Copy pathtest_training.py
File metadata and controls
84 lines (64 loc) · 2.53 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import pytest
import replicate
from replicate.exceptions import ReplicateException
input_images_url = "https://replicate.delivery/pbxt/JMV5OrEWpBAC5gO8rre0tPOyJIOkaXvG0TWfVJ9b4zhLeEUY/data.zip"
@pytest.mark.vcr("trainings-create.yaml")
@pytest.mark.asyncio
async def test_trainings_create(mock_replicate_api_token):
training = replicate.trainings.create(
version="stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
input={
"input_images": input_images_url,
"use_face_detection_instead": True,
},
destination="replicate/dreambooth-sdxl",
)
assert training.id is not None
assert training.status == "starting"
@pytest.mark.vcr("trainings-create.yaml")
@pytest.mark.asyncio
async def test_trainings_create_with_positional_argument(mock_replicate_api_token):
training = replicate.trainings.create(
"stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
{
"input_images": input_images_url,
"use_face_detection_instead": True,
},
"replicate/dreambooth-sdxl",
)
assert training.id is not None
assert training.status == "starting"
@pytest.mark.vcr("trainings-create__invalid-destination.yaml")
@pytest.mark.asyncio
async def test_trainings_create_with_invalid_destination(mock_replicate_api_token):
with pytest.raises(ReplicateException):
replicate.trainings.create(
"stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
input={
"input_images": input_images_url,
},
destination="<invalid>",
)
@pytest.mark.vcr("trainings-get.yaml")
@pytest.mark.asyncio
async def test_trainings_get(mock_replicate_api_token):
id = "medrnz3bm5dd6ultvad2tejrte"
training = replicate.trainings.get(id)
assert training.id == id
assert training.status == "processing"
@pytest.mark.vcr("trainings-cancel.yaml")
@pytest.mark.asyncio
async def test_trainings_cancel(mock_replicate_api_token):
input = {
"input_images": input_images_url,
"use_face_detection_instead": True,
}
destination = "replicate/dreambooth-sdxl"
training = replicate.trainings.create(
version="stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
destination=destination,
input=input,
)
assert training.status == "starting"
# training = replicate.trainings.cancel(training)
training.cancel()