File size: 1,863 Bytes
a09cfc1 | 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 | from text_diffusion_pipeline import TextConditionalDDPMPipeline
from latent_diffusion_pipeline import UnconditionalDDPMPipeline
import os
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
def get_pipeline(model_path):
# If model_path is a local directory, use the original logic
if os.path.isdir(model_path):
#Diffusion models
if os.path.exists(os.path.join(model_path, "unet")):
if os.path.exists(os.path.join(model_path, "text_encoder")):
#If it has a text encoder and a unet, it's text conditional diffusion
pipe = TextConditionalDDPMPipeline.from_pretrained(model_path)
else:
#If it has no text encoder, use the unconditional diffusion model
pipe = UnconditionalDDPMPipeline.from_pretrained(model_path)
else:
# Assume it's a Hugging Face Hub model ID
# Try to load config to determine if it's text-conditional
try:
config, _ = DiffusionPipeline.load_config(model_path)
components = config.get("components", {})
except Exception:
components = {}
if "text_encoder" in components or "text_encoder" in str(components):
# Use the local pipeline file for custom_pipeline
pipe = DiffusionPipeline.from_pretrained(
model_path,
custom_pipeline="models.text_diffusion_pipeline.TextConditionalDDPMPipeline",
trust_remote_code=True,
)
else:
# Fallback: try unconditional
pipe = DiffusionPipeline.from_pretrained(
model_path,
custom_pipeline="models.latent_diffusion_pipeline.UnconditionalDDPMPipeline",
trust_remote_code=True,
)
return pipe |