Skip to content

Commit

Permalink
🔧 fix: Fix data structure of long polling response
Browse files Browse the repository at this point in the history
  • Loading branch information
@bryantanjw
bryantanjw committed Jun 8, 2024
1 parent 4a957d1 commit aef0a65
Show file tree
Hide file tree
Showing 7 changed files with 223 additions and 93 deletions.
80 changes: 48 additions & 32 deletions README.md
View file

Large diffs are not rendered by default.

68 changes: 55 additions & 13 deletions civitai/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ async def async_create():

base_model = "SDXL" if "sdxl" in validated_input.model else "SD_1_5"
job_input = {"$type": "textToImage",
"baseModel": base_model, **validated_input.dict(exclude_unset=True)}
"baseModel": base_model, **validated_input.model_dump(exclude_unset=True)}

response = await post_v1consumerjobs(job_input, wait=False, api_config_override=self.civitai)
modified_response = {
Expand All @@ -62,13 +62,20 @@ async def async_create():
}

if wait:
job_result = await self._poll_for_job_completion(response.token)
if job_result and modified_response["jobs"]:
modified_response["jobs"][0]["result"] = job_result
response = await self._poll_for_job_completion(response.token)
return response

return modified_response

return asyncio.run(async_create())
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None

if loop and loop.is_running():
return asyncio.ensure_future(async_create())
else:
return asyncio.run(async_create())

async def _poll_for_job_completion(self, token, interval=1, timeout=300):
start_time = time.time()
Expand All @@ -78,14 +85,25 @@ async def _poll_for_job_completion(self, token, interval=1, timeout=300):
if response and response.jobs:
last_response = response
logging.info(f"Job status: {last_response}")
job = response.jobs[0]
if job.result and job.result.get("blobUrl"):
return job.result
completed_jobs = []
for job in response.jobs:
if job.result and job.result.get("blobUrl"):
completed_jobs.append({
"jobId": job.jobId,
"cost": job.cost,
"result": job.result,
"scheduled": job.scheduled,
})
if completed_jobs:
return {
"token": response.token,
"jobs": completed_jobs
}
await asyncio.sleep(interval)

if last_response:
logging.warning(f"Job {token} did not complete within {
timeout} seconds. Returning the last response.")
logging.warning(
"Job {token} did not complete within 5 minutes.")
return {
"token": last_response.token,
"jobs": [{
Expand Down Expand Up @@ -127,21 +145,45 @@ async def async_get():

return modified_response

return asyncio.run(async_get())
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None

if loop and loop.is_running():
return asyncio.ensure_future(async_get())
else:
return asyncio.run(async_get())

def query(self, query, detailed=False):
async def async_query():
response = await post_v1consumerjobsquery(query, detailed=detailed, api_config_override=self.civitai)
return response

return asyncio.run(async_query())
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None

if loop and loop.is_running():
return asyncio.ensure_future(async_query())
else:
return asyncio.run(async_query())

def cancel(self, job_id):
async def async_cancel():
response = await delete_v1consumerjobsjobId(job_id, api_config_override=self.civitai)
return response

return asyncio.run(async_cancel())
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None

if loop and loop.is_running():
return asyncio.ensure_future(async_cancel())
else:
return asyncio.run(async_cancel())


# Expose the 'Civitai' class at the module level
Expand Down
14 changes: 9 additions & 5 deletions examples/streamlit_demo/app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ def generate_image(input_data):
:return: The job result
"""
try:
output = civitai.image.create(input_data)
output = civitai.image.create(input_data, wait=True)
return output
except Exception as e:
st.error(f"Failed to generate image: {str(e)}")
Expand All @@ -27,14 +27,18 @@ def generate_image(input_data):

with col_input:
with st.form("image_gen_form"):
model = st.text_input("Model URN", value="urn:air:sd1:checkpoint:civitai:4201@130072")
model = st.text_input(
"Model URN", value="urn:air:sdxl:checkpoint:civitai:101055@128078")
prompt = st.text_area("Prompt", value="A cat")
negative_prompt = st.text_area("Negative Prompt", value="A dog")
scheduler = st.selectbox("Scheduler", ["EulerA", "HeunC", "Dopri5", "Bosh3"], index=0)
scheduler = st.selectbox(
"Scheduler", ["EulerA", "HeunC", "Dopri5", "Bosh3"], index=0)
steps = st.slider("Steps", min_value=1, max_value=100, value=30)
cfg_scale = st.slider("CFG Scale", min_value=1, max_value=20, value=10)
width = st.number_input("Width", min_value=1, max_value=1024, value=768)
height = st.number_input("Height", min_value=1, max_value=1024, value=512)
width = st.number_input("Width", min_value=1,
max_value=1024, value=768)
height = st.number_input(
"Height", min_value=1, max_value=1024, value=512)
seed = st.number_input("Seed", value=-1)
submit_button = st.form_submit_button("Generate Image")

Expand Down
139 changes: 104 additions & 35 deletions examples/text2img.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,42 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 10,
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Requirement already satisfied: civitai-py==0.1.62 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (0.1.62)\n",
"Requirement already satisfied: pydantic>=2 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from civitai-py==0.1.62) (2.7.3)\n",
"Requirement already satisfied: typing-extensions>=4.7.1 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from civitai-py==0.1.62) (4.10.0)\n",
"Requirement already satisfied: urllib3>=1.25.3 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from civitai-py==0.1.62) (2.0.7)\n",
"Requirement already satisfied: annotated-types>=0.4.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from pydantic>=2->civitai-py==0.1.62) (0.6.0)\n",
"Requirement already satisfied: pydantic-core==2.18.4 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from pydantic>=2->civitai-py==0.1.62) (2.18.4)\n",
"Note: you may need to restart the kernel to use updated packages.\n",
"Requirement already satisfied: ipython in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (8.22.2)\n",
"Requirement already satisfied: decorator in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (5.1.1)\n",
"Requirement already satisfied: jedi>=0.16 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (0.19.1)\n",
"Requirement already satisfied: matplotlib-inline in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (0.1.6)\n",
"Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (3.0.43)\n",
"Requirement already satisfied: pygments>=2.4.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (2.17.2)\n",
"Requirement already satisfied: stack-data in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (0.6.3)\n",
"Requirement already satisfied: traitlets>=5.13.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (5.14.2)\n",
"Requirement already satisfied: pexpect>4.3 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from ipython) (4.9.0)\n",
"Requirement already satisfied: parso<0.9.0,>=0.8.3 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from jedi>=0.16->ipython) (0.8.3)\n",
"Requirement already satisfied: ptyprocess>=0.5 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from pexpect>4.3->ipython) (0.7.0)\n",
"Requirement already satisfied: wcwidth in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython) (0.2.13)\n",
"Requirement already satisfied: executing>=1.2.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from stack-data->ipython) (2.0.1)\n",
"Requirement already satisfied: asttokens>=2.1.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from stack-data->ipython) (2.4.1)\n",
"Requirement already satisfied: pure-eval in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from stack-data->ipython) (0.2.2)\n",
"Requirement already satisfied: six>=1.12.0 in /Users/bryan/civitai-python/venv/lib/python3.12/site-packages (from asttokens>=2.1.0->stack-data->ipython) (1.16.0)\n",
"Note: you may need to restart the kernel to use updated packages.\n"
]
}
],
"source": [
"%pip install civitai-py\n",
"%pip install civitai-py==0.1.62\n",
"%pip install ipython"
]
},
Expand All @@ -39,22 +70,30 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 11,
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Response: {'token': 'eyJKb2JzIjpbImZhOTQyM2IxLTlkODEtNDhkZi1iYjc3LTkwOTNlNWEwMzUyMSJdfQ==', 'jobs': [{'jobId': 'fa9423b1-9d81-48df-bb77-9093e5a03521', 'cost': 2.0, 'result': {'blobKey': 'D4A2F1AA82E6639148D0C889C510C2C8A4F5C2E0D267191FDAA5D80957FCA18B', 'available': False}, 'scheduled': True}]}\n"
]
}
],
"source": [
"# Import the Civitai SDK\n",
"import civitai\n",
"\n",
"# Define the input parameters for the image creation\n",
"input = {\n",
" \"model\": \"urn:air:sd1:checkpoint:civitai:4201@130072\",\n",
" \"model\": \"urn:air:sdxl:checkpoint:civitai:101055@128078\",\n",
" \"params\": {\n",
" \"prompt\": \"A cat\",\n",
" \"negativePrompt\": \"A dog\",\n",
" \"prompt\": \"RAW photo, face portrait photo of woman, wearing black dress, happy face, hard shadows, cinematic shot, dramatic lighting\",\n",
" \"negativePrompt\": \"(deformed, distorted, disfigured:1.3)\",\n",
" \"scheduler\": \"EulerA\",\n",
" \"steps\": 30,\n",
" \"cfgScale\": 10,\n",
" \"steps\": 20,\n",
" \"cfgScale\": 7,\n",
" \"width\": 768,\n",
" \"height\": 512,\n",
" \"seed\": -1,\n",
Expand All @@ -63,7 +102,8 @@
"}\n",
"\n",
"# Generate the image\n",
"response = civitai.image.create(input)\n",
"# We're using await in the notebook to handle the coroutine\n",
"response = await civitai.image.create(input)\n",
"print(\"Response:\", response)"
]
},
Expand All @@ -76,31 +116,54 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 7,
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Job Status Response: {'token': 'eyJKb2JzIjpbIjZmZGY5OWM5LTc4NTEtNDAzYy1iODI3LTlhZTdmNzc0MDY1NiJdfQ==', 'jobs': [{'jobId': '6fdf99c9-7851-403c-b827-9ae7f7740656', 'cost': 2.0, 'result': {'blobKey': '1E74613D6580100FFE70C550EA024885900E758846C4F59A54DFA67B6649F153', 'available': True, 'blobUrl': 'https://blobs-temp.sfo3.digitaloceanspaces.com/1E74613D6580100FFE70C550EA024885900E758846C4F59A54DFA67B6649F153?X-Amz-Expires=3600&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=DO00F84RAAYEUTBJ6D9L%2F20240608%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240608T185601Z&X-Amz-SignedHeaders=host&X-Amz-Signature=6df5a825a67bfb6f3dd28ff602b7a06964f752c6cbdbe2e5f88c9422502f3ea8', 'blobUrlExpirationDate': '2024-06-08T19:56:01.8499363Z'}, 'scheduled': False}]}\n"
]
},
{
"data": {
"text/html": [
"<img src=\"https://blobs-temp.sfo3.digitaloceanspaces.com/1E74613D6580100FFE70C550EA024885900E758846C4F59A54DFA67B6649F153?X-Amz-Expires=3600&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=DO00F84RAAYEUTBJ6D9L%2F20240608%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240608T185601Z&X-Amz-SignedHeaders=host&X-Amz-Signature=6df5a825a67bfb6f3dd28ff602b7a06964f752c6cbdbe2e5f88c9422502f3ea8\"/>"
],
"text/plain": [
"<IPython.core.display.Image object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"job_token = response['token']\n",
"job_id = response['jobs'][0]['jobId']\n",
"\n",
"# Retrieve job status and image\n",
"job_response = civitai.jobs.get(token=job_token, job_id=job_id)\n",
"print(\"Job Status Response:\", job_response)\n",
"response = await civitai.jobs.get(token=job_token, job_id=job_id)\n",
"print(\"Job Status Response:\", response)\n",
"\n",
"if job_response and 'jobs' in job_response and 'result' in job_response['jobs'][0]:\n",
" image_url = job_response['jobs'][0]['result'].get('blobUrl')\n",
" display(Image(url=image_url))\n",
"if response['jobs'][0]['result'].get('available'):\n",
" image_url = response['jobs'][0]['result'].get('blobUrl')\n",
" if image_url:\n",
" display(Image(url=image_url))\n",
" else:\n",
" print(\"Image URL not found in the job result.\")\n",
"else:\n",
" print(\"No image was created, or the job is not yet complete.\")"
" print(\"No image was created, the job is not yet complete, or the result is not available.\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Run image generation in the background."
"### Example 2: Run image generation in the background."
]
},
{
Expand All @@ -112,32 +175,38 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<img src=\"https://blobs-temp.sfo3.digitaloceanspaces.com/D0694EBA47EE72F362A42D27F8E71325AA562FB28A96D124FF28972B9C982A46?X-Amz-Expires=3600&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=DO00F84RAAYEUTBJ6D9L%2F20240330%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240330T015524Z&X-Amz-SignedHeaders=host&X-Amz-Signature=ca20e3fe683beef5276bd470cd12cbfbdc670921f112297a766750dccb341c0b\"/>"
],
"text/plain": [
"<IPython.core.display.Image object>"
]
},
"metadata": {},
"output_type": "display_data"
"name": "stderr",
"output_type": "stream",
"text": [
"WARNING:root:Job eyJKb2JzIjpbIjE3ZTA5MzRjLTIwOWYtNGQ1Zi1iOGQwLWE4YzBiNzM4MDU5YiJdfQ== did not complete within 20 seconds. Returning the last response.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Response: {'token': 'eyJKb2JzIjpbIjE3ZTA5MzRjLTIwOWYtNGQ1Zi1iOGQwLWE4YzBiNzM4MDU5YiJdfQ==', 'jobs': [{'jobId': '17e0934c-209f-4d5f-b8d0-a8c0b738059b', 'cost': 2.0, 'result': {'token': 'eyJKb2JzIjpbIjE3ZTA5MzRjLTIwOWYtNGQ1Zi1iOGQwLWE4YzBiNzM4MDU5YiJdfQ==', 'jobs': [{'jobId': '17e0934c-209f-4d5f-b8d0-a8c0b738059b', 'cost': 2.0, 'result': {'blobKey': 'EF944D56AF81AF2E5C3D3FFEE7C1AAEBFE0F826663771AA0A0E3627201B7732F', 'available': False}, 'scheduled': True}]}, 'scheduled': True}]}\n",
"No image was created, the job is not yet complete, or the result is not available.\n"
]
}
],
"source": [
"# Generate the image and wait for completion\n",
"response = civitai.image.create(input, wait=True)\n",
"response = await civitai.image.create(input, wait=True)\n",
"print(\"Response:\", response)\n",
"\n",
"if response and 'jobs' in response and 'result' in response['jobs'][0]:\n",
"if response['jobs'][0]['result'].get('available'):\n",
" image_url = response['jobs'][0]['result'].get('blobUrl')\n",
" display(Image(url=image_url))\n",
" if image_url:\n",
" display(Image(url=image_url))\n",
" else:\n",
" print(\"Image URL not found in the job result.\")\n",
"else:\n",
" print(\"No image was created, or the job is not yet complete.\")"
" print(\"No image was created, the job is not yet complete, or the result is not available.\")"
]
},
{
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "civitai-py"
version = "0.1.5"
version = "0.1.7"
description = "Civitai Python SDK"
authors = ["Civitai <hello@civitai.com>"]
license = "MIT"
Expand Down
11 changes: 5 additions & 6 deletions tests/test_create_image.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,15 @@ def test_create_from_text_job(self):
input_data = {
"model": "urn:air:sdxl:checkpoint:civitai:101055@128078",
"params": {
"prompt": "A cat",
"negativePrompt": "A dog",
"prompt": "RAW photo, face portrait photo of 26 y.o woman, wearing black dress, happy face, hard shadows, cinematic shot, dramatic lighting",
"negativePrompt": "(deformed, distorted, disfigured:1.3)",
"scheduler": "EulerA",
"steps": 20,
"cfgScale": 7,
"width": 768,
"width": 512,
"height": 512,
"seed": -1,
"clipSkip": 1
},
"clipSkip": 2
}
}

output = civitai.image.create(input_data, wait=True)
Expand Down
2 changes: 1 addition & 1 deletion tests/test_get_jobs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
class TestGetJobStatus(unittest.TestCase):
def test_token_precedence_over_id(self):
"""Test that token takes precedence over job ID when both are provided."""
job_token = 'eyJKb2JzIjpbImNlZDM5NTk3LWVhNDgtNDM4Yi1iNGQzLWQ1YzJmZjk5MTBhZCJdfQ==' # Replace with a valid job token
job_token = 'eyJKb2JzIjpbImEyNWVhNzg4LWE0YjQtNDI4MC1iZjQxLWJjYjc2NWVhOGMzNSJdfQ==' # Replace with a valid job token
job_id = 'ced39597-ea48-438b-b4d3-d5c2ff9910ad' # Replace with a valid job ID

response = civitai.jobs.get(token=job_token, job_id=job_id)
Expand Down

0 comments on commit aef0a65

Please sign in to comment.