Spaces:
Runtime error
Runtime error
| from diffusers import ControlNetModel,StableDiffusionControlNetInpaintPipeline,AutoPipelineForInpainting | |
| import torch | |
| class PipelineFetcher: | |
| """ | |
| A class that fetches different pipelines for image processing. | |
| Args: | |
| controlnet_adapter_model_name (str): The name of the controlnet adapter model. | |
| controlnet_base_model_name (str): The name of the controlnet base model. | |
| kandinsky_model_name (str): The name of the Kandinsky model. | |
| image (str): The image to be processed. | |
| """ | |
| def __init__(self, controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image: str): | |
| self.controlnet_adapter_model_name = controlnet_adapter_model_name | |
| self.controlnet_base_model_name = controlnet_base_model_name | |
| self.kandinsky_model_name = kandinsky_model_name | |
| self.image = image | |
| def ControlNetInpaintPipeline(self): | |
| """ | |
| Fetches the ControlNet inpainting pipeline. | |
| Returns: | |
| pipe (StableDiffusionControlNetInpaintPipeline): The ControlNet inpainting pipeline. | |
| """ | |
| controlnet = ControlNetModel.from_pretrained(self.controlnet_adapter_model_name, torch_dtype=torch.float16) | |
| pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( | |
| self.controlnet_base_model_name, controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| pipe.to('cuda') | |
| return pipe | |
| def KandinskyPipeline(self): | |
| """ | |
| Fetches the Kandinsky pipeline. | |
| Returns: | |
| pipe (AutoPipelineForInpainting): The Kandinsky pipeline. | |
| """ | |
| pipe = AutoPipelineForInpainting.from_pretrained(self.kandinsky_model_name, torch_dtype=torch.float16) | |
| pipe.to('cuda') | |
| return pipe | |
| def fetch_control_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image): | |
| """ | |
| Fetches the control pipeline for image processing. | |
| Args: | |
| controlnet_adapter_model_name (str): The name of the controlnet adapter model. | |
| controlnet_base_model_name (str): The name of the controlnet base model. | |
| kandinsky_model_name (str): The name of the Kandinsky model. | |
| image: The input image for processing. | |
| Returns: | |
| pipe: The control pipeline for image processing. | |
| """ | |
| pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image) | |
| pipe = pipe_fetcher.ControlNetInpaintPipeline() | |
| return pipe | |
| def fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image): | |
| """ | |
| Fetches the Kandinsky pipeline. | |
| Args: | |
| controlnet_adapter_model_name (str): The name of the controlnet adapter model. | |
| controlnet_base_model_name (str): The name of the controlnet base model. | |
| kandinsky_model_name (str): The name of the Kandinsky model. | |
| image: The input image. | |
| Returns: | |
| pipe: The Kandinsky pipeline. | |
| """ | |
| pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image) | |
| pipe = pipe_fetcher.KandinskyPipeline() | |
| return pipe | |