| import torch |
| try: |
| import fal_client |
| except: |
| fal_client = None |
|
|
| from diffusers import AutoPipelineForText2Image, HunyuanVideoPipeline, DiffusionPipeline |
| import json |
| import diffusers |
| from functools import partial |
| import os |
| |
| os.environ['FAL_KEY'] = 'YOUR_API_KEY' |
|
|
| def init_multiple_pipelines(pipe_name, pipe_init_kwargs, num_devices, device_id=None): |
| pipelines_dict = [] |
|
|
| if device_id is not None: |
| assert num_devices == 1 |
|
|
| for i in range(num_devices): |
| actual_device_id = device_id if device_id is not None else i |
| try: |
| pipeline = AutoPipelineForText2Image.from_pretrained(pipe_name, **pipe_init_kwargs).to(f'cuda:{actual_device_id}') |
| except Exception as e: |
| |
| 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(f'cuda:{actual_device_id}') |
| |
| |
| |
| pipelines_dict.append(pipeline) |
| return pipelines_dict |
|
|
|
|
| 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 |
