Spaces:
Runtime error
Runtime error
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline | |
import torch | |
from PIL import Image | |
import lightning.pytorch as pl | |
from scripts.api_utils import accelerator | |
from typing import Optional | |
pl.seed_everything(42) | |
class ImageGenerator: | |
""" | |
A class to generate images using ControlNet and Stable Diffusion XL pipelines. | |
Attributes: | |
controlnet (ControlNetModel): The ControlNet model. | |
pipeline (StableDiffusionXLControlNetPipeline): The Stable Diffusion XL pipeline with ControlNet. | |
""" | |
def __init__(self, controlnet_model_name, sd_pipeline_model_name): | |
""" | |
Initializes the ImageGenerator with the specified models. | |
Args: | |
controlnet_model_name (str): The name of the ControlNet model. | |
sd_pipeline_model_name (str): The name of the Stable Diffusion XL pipeline model. | |
image (str): The path to the image to be used. | |
""" | |
self.controlnet = ControlNetModel.from_pretrained( | |
controlnet_model_name, torch_dtype=torch.float16, variant="fp16" | |
) | |
self.pipeline = StableDiffusionXLControlNetPipeline.from_pretrained( | |
sd_pipeline_model_name, | |
torch_dtype=torch.float16, | |
variant="fp16", | |
controlnet=self.controlnet, | |
).to(accelerator()) | |
def inference(self, prompt, negative_prompt, height, width, guidance_scale, num_images_per_prompt, num_inference_steps, image_path, controlnet_conditioning_scale, control_guidance_end,output_path:Optional[str]): | |
""" | |
Generates images based on the provided parameters. | |
Args: | |
prompt (str): The prompt for image generation. | |
negative_prompt (str): The negative prompt for image generation. | |
height (int): The height of the generated images. | |
width (int): The width of the generated images. | |
guidance_scale (float): The guidance scale for image generation. | |
num_images_per_prompt (int): The number of images to generate per prompt. | |
num_inference_steps (int): The number of inference steps. | |
image_path (str): The path to the image to be used. | |
controlnet_conditioning_scale (float): The conditioning scale for ControlNet. | |
control_guidance_end (float): The end guidance for ControlNet. | |
Returns: | |
list: A list of generated images. | |
""" | |
images_list = self.pipeline( | |
prompt=prompt, | |
negative_prompt=negative_prompt, | |
height=height, | |
width=width, | |
guidance_scale=guidance_scale, | |
num_images_per_prompt=num_images_per_prompt, | |
num_inference_steps=num_inference_steps, | |
image=Image.open(image_path), | |
controlnet_conditioning_scale=controlnet_conditioning_scale, | |
control_guidance_end=control_guidance_end, | |
).images | |
if output_path: | |
for i,image in enumerate(images_list): | |
image.save(f'{output_path}/output_{i}.png') | |
else: | |
return images_list | |
if __name__ == "__main__": | |
generator = ImageGenerator( | |
controlnet_model_name="destitech/controlnet-inpaint-dreamer-sdxl", | |
sd_pipeline_model_name="RunDiffusion/Juggernaut-XL-v9" | |
) | |
generator.inference( | |
prompt='Park', | |
negative_prompt='low Resolution , Bad Resolution', | |
height=1080, | |
width=1920, | |
guidance_scale=7.5, | |
num_images_per_prompt=4, | |
num_inference_steps=100, | |
image_path='/home/PicPilot/sample_data/example1.jpg', | |
controlnet_conditioning_scale=0.9, | |
control_guidance_end=0.9, | |
output_path='/home/PicPilot/output' | |
) | |