run format

This commit is contained in:
Toran Bruce Richards
2024-09-29 20:57:05 +01:00
parent ca54d06090
commit deb64d8da1

View File

@@ -1,7 +1,9 @@
import replicate
from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema
from backend.data.model import BlockSecret, SchemaField, SecretField
class ReplicateFluxBasicModelBlock(Block):
class Input(BlockSchema):
api_key: BlockSecret = SecretField(
@@ -67,7 +69,7 @@ class ReplicateFluxBasicModelBlock(Block):
# Run the model
output = client.run(
f"{model_name}",
input={"prompt": prompt,"output_format": "png"},
input={"prompt": prompt, "output_format": "png"},
)
# Check if output is a list or a string and extract accordingly; otherwise, assign a default message
@@ -76,5 +78,7 @@ class ReplicateFluxBasicModelBlock(Block):
elif isinstance(output, str):
result_url = output # If output is a string, use it directly
else:
result_url = "No output received" # Fallback message if output is not as expected
result_url = (
"No output received" # Fallback message if output is not as expected
)
return result_url