| |
|
|
| import torch |
| import os |
|
|
| from diffusers import DiffusionPipeline, ControlNetModel, DDIMScheduler |
|
|
| from PIL import Image |
|
|
| test_prompt = "best quality, extremely detailed" |
| test_negative_prompt = "blur, lowres, bad anatomy, worst quality, low quality" |
|
|
| def resize_for_condition_image(input_image: Image, resolution: int): |
| input_image = input_image.convert("RGB") |
| W, H = input_image.size |
| k = float(resolution) / min(H, W) |
| H *= k |
| W *= k |
| H = int(round(H / 64.0)) * 64 |
| W = int(round(W / 64.0)) * 64 |
| img = input_image.resize((W, H), resample=Image.LANCZOS if k > 1 else Image.AREA) |
| return img |
|
|
| def generate_image(seed, prompt, negative_prompt, control, guess_mode=False): |
| latent = torch.randn( |
| (1, 4, 64, 64), |
| device="cpu", |
| generator=torch.Generator(device="cpu").manual_seed(seed), |
| ).cuda() |
| image = pipe( |
| prompt=prompt, |
| negative_prompt=negative_prompt, |
| guidance_scale=4.0 if guess_mode else 9.0, |
| num_inference_steps=50 if guess_mode else 20, |
| latents=latent, |
| image=control, |
| controlnet_conditioning_image=control, |
| strength=1.0, |
| |
| ).images[0] |
| return image |
|
|
|
|
| if __name__ == "__main__": |
| model_name = "f1e_sd15_tile" |
| original_image_folder = "./control_images/" |
| control_image_folder = "./control_images/converted/" |
| output_image_folder = "./output_images/diffusers/" |
| os.makedirs(output_image_folder, exist_ok=True) |
|
|
| |
| |
| controlnet = ControlNetModel.from_pretrained('takuma104/control_v11', |
| subfolder='control_v11f1e_sd15_tile') |
|
|
| if model_name == "p_sd15s2_lineart_anime": |
| base_model_id = "Linaqruf/anything-v3.0" |
| base_model_revision = None |
| else: |
| base_model_id = "runwayml/stable-diffusion-v1-5" |
| base_model_revision = "non-ema" |
|
|
| pipe = DiffusionPipeline.from_pretrained( |
| base_model_id, |
| revision=base_model_revision, |
| custom_pipeline="stable_diffusion_controlnet_img2img", |
| controlnet=controlnet, |
| safety_checker=None, |
| ).to("cuda") |
| pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) |
|
|
| original_image_filenames = [ |
| "dog_64x64.png", |
| ] |
|
|
| image_conditions = [ |
| resize_for_condition_image( |
| Image.open(f"{original_image_folder}{fn}"), |
| resolution=512, |
| ) |
| for fn in original_image_filenames |
| ] |
|
|
| for i, control in enumerate(image_conditions): |
| for seed in range(4): |
| image = generate_image( |
| seed=seed, |
| prompt=test_prompt, |
| negative_prompt=test_negative_prompt, |
| control=control, |
| ) |
| image.save(f"{output_image_folder}output_{model_name}_{i}_{seed}.png") |
|
|
