schrum2 commited on
Commit
b84e862
·
verified ·
1 Parent(s): fc4946f

change where to look for pipeline

Browse files
Files changed (1) hide show
  1. models/pipeline_loader.py +2 -2
models/pipeline_loader.py CHANGED
@@ -29,14 +29,14 @@ def get_pipeline(model_path):
29
  # Use the local pipeline file for custom_pipeline
30
  pipe = DiffusionPipeline.from_pretrained(
31
  model_path,
32
- custom_pipeline="models.text_diffusion_pipeline.TextConditionalDDPMPipeline",
33
  trust_remote_code=True,
34
  )
35
  else:
36
  # Fallback: try unconditional
37
  pipe = DiffusionPipeline.from_pretrained(
38
  model_path,
39
- custom_pipeline="models.latent_diffusion_pipeline.UnconditionalDDPMPipeline",
40
  trust_remote_code=True,
41
  )
42
 
 
29
  # Use the local pipeline file for custom_pipeline
30
  pipe = DiffusionPipeline.from_pretrained(
31
  model_path,
32
+ custom_pipeline=model_path,
33
  trust_remote_code=True,
34
  )
35
  else:
36
  # Fallback: try unconditional
37
  pipe = DiffusionPipeline.from_pretrained(
38
  model_path,
39
+ custom_pipeline=model_path,
40
  trust_remote_code=True,
41
  )
42