Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,635 Bytes
3427608 |
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 |
import imagen_hub
class ImagenHubModel():
def __init__(self, model_name):
self.model = imagen_hub.load(model_name)
def __call__(self, *args, **kwargs):
return self.model.infer_one_image(*args, **kwargs)
class PNP(ImagenHubModel):
def __init__(self):
super().__init__('PNP')
def __call__(self, *args, **kwargs):
if "num_inversion_steps" not in kwargs:
kwargs["num_inversion_steps"] = 200
return super().__call__(*args, **kwargs)
class Prompt2prompt(ImagenHubModel):
def __init__(self):
super().__init__('Prompt2prompt')
def __call__(self, *args, **kwargs):
if "num_inner_steps" not in kwargs:
kwargs["num_inner_steps"] = 3
return super().__call__(*args, **kwargs)
def load_imagenhub_model(model_name, model_type=None):
if model_name == 'PNP':
return PNP()
if model_name == 'Prompt2prompt':
return Prompt2prompt()
return ImagenHubModel(model_name)
# for name in ['DeepFloydIF', 'PixArtAlpha', 'Kandinsky']: #, 'OpenJourney', 'LCM', 'SD' 'SDXL'
# #
# pipe = ImagenHubModel(name)
# result = pipe(prompt='a cute dog is playing a ball')
# print(result)
# for name in ['SD']:
# from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
# import torch
# pipe = DiffusionPipeline.from_pretrained(
# "stabilityai/stable-diffusion-2-base",
# torch_dtype=torch.float16,
# safety_checker=None,
# ).to("cuda")
# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) |