File size: 1,543 Bytes
9b57ce7 |
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 |
import torch
try:
import fal_client
except:
fal_client = None
try:
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
except:
AutoPipelineForText2Image = None
DiffusionPipeline = None
import json
import diffusers
from functools import partial
import os
# export FAL_KEY="YOUR_API_KEY"
os.environ['FAL_KEY'] = 'YOUR_API_KEY'
def init_pipelines(pipe_name, pipe_init_kwargs, device=None):
try:
pipeline = AutoPipelineForText2Image.from_pretrained(pipe_name, **pipe_init_kwargs).to(device)
except Exception as e:
# try:
config = json.load(open(os.path.join(pipe_name, 'model_index.json')))
class_name_str = config['_class_name']
pipeline_class = getattr(diffusers, class_name_str)
pipeline = pipeline_class.from_pretrained(pipe_name, **pipe_init_kwargs).to(device)
return pipeline
def init_pipeline_from_names(pipe_names, weight_dtype):
pipelines_dict = {}
for name in pipe_names:
pipeline = AutoPipelineForText2Image.from_pretrained(name, torch_dtype=weight_dtype)
pipelines_dict[name] = pipeline
return pipelines_dict
def on_queue_update(update):
if isinstance(update, fal_client.InProgress):
for log in update.logs:
print(log["message"])
def gen_with_api(pipe_names, generation_kwargs):
result = fal_client.subscribe(
pipe_names,
arguments=generation_kwargs,
with_logs=True,
on_queue_update=on_queue_update,
)
return result |