hi-res-upscale diffusers example

hi-res-upscale-example.py

from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
from diffusers import EulerAncestralDiscreteScheduler
import torch
import torch.nn.functional
import gc
from PIL import Image

from mechanisms.tokenizer_utils import encode_from_pipe

@torch.no_grad()
def run():
    torch_dtype = torch.bfloat16

    model_path = "/home/blackroot/Desktop/anime_character_studio/generation/models/autismmixSDXL_autismmixConfetti.safetensors"
    pipe = StableDiffusionXLPipeline.from_single_file(
        model_path,
        torch_dtype=torch_dtype,
        variant="bf16", 
        use_safetensors=True,
        add_watermarker=False,
        custom_pipeline="lpw_stable_diffusion_xl",
        ).to('cuda')
    
    scheduler_config = {
        "beta_end": 0.012,
        "beta_schedule": "scaled_linear",
        "beta_start": 0.00085,
        "clip_sample": False,
        "interpolation_type": "linear",
        "prediction_type": "epsilon",
        "sample_max_value": 1.0,
        "set_alpha_to_one": False,
        "skip_prk_steps": True,
        "steps_offset": 1,
        "timestep_spacing": "leading",
        "trained_betas": None,
        "use_karras_sigmas": False
    }
    
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler_config)

    prompt = "score_9, score_8_up, score_7_up, score_6_up, source_anime, source_safe, Anime Catgirl."
    generator = torch.Generator("cuda")

    image = pipe(
        prompt=prompt,
        output_type = "pil", 
        width = 1024, 
        height = 1024, 
        guidance_scale = 7.0, 
        strength = 1.0,
        num_inference_steps=20).images[0]
    
    resized = image.resize((1536, 1536), Image.NEAREST)
    resized.save('first.png')
    
    pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
        model_path,
        torch_dtype=torch_dtype,
        variant="bf16", 
        use_safetensors=True,
        add_watermarker=False,
        custom_pipeline="lpw_stable_diffusion_xl",
        ).to('cuda')
    
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler_config)
     
    images = pipe(
        image=resized,
        prompt=prompt,
        output_type = "pil", 
        guidance_scale = 7.0, 
        strength = .8,
        num_inference_steps=20).images
    
    for image in images:
        image.save('second_step.png')
        
run()