DeepFloyd IF for around 12GB of VRAM

DeepFloyd_IF_VRAM12GB.py

# huggingface login
from huggingface_hub import login

login()


# load textencorder in 8bit quantized
from transformers import T5EncoderModel
from diffusers import DiffusionPipeline
import datetime
import gc
import torch
from diffusers.utils import pt_to_pil

def flush():
    gc.collect()
    torch.cuda.empty_cache()

default_prompt = "anime girl, silver hair and blue eyes, smile"
default_negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
default_steps = 50

while(True):
    prompt = input("prompt: ")
    if(prompt == ""):prompt = default_prompt
    elif(prompt == "exit()"):break
    negative_prompt = input("negative_prompt: ")
    if(negative_prompt == ""):negative_prompt = default_negative_prompt
    elif(negative_prompt == " "):negative_prompt = ""
    steps = input(f"stage1 steps(default: {default_steps}): ")
    if(steps == ""):steps = default_steps
    else:steps = int(steps)

    start_dt = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

    print("Loading Text Encoder...",end="")
    text_encoder = T5EncoderModel.from_pretrained(
        "DeepFloyd/IF-I-XL-v1.0",
        subfolder="text_encoder", 
        device_map="auto", 
        load_in_8bit=True, 
        variant="8bit"
    )
    print("finish!")

    # load stage1 model for text encording
    print("Encording Text...",end="")
    pipe = DiffusionPipeline.from_pretrained(
        "DeepFloyd/IF-I-XL-v1.0", 
        text_encoder=text_encoder, # pass the previously instantiated 8bit text encoder
        unet=None, 
        device_map="auto"
    )
    prompt_embeds, negative_embeds = pipe.encode_prompt(prompt=prompt, negative_prompt=negative_prompt)
    print("finish!")

    # unload textencorder
    del text_encoder
    del pipe

    flush()
    # load stage1 model for generation
    print("Loading stage1 model...",end="")
    pipe = DiffusionPipeline.from_pretrained(
        "DeepFloyd/IF-I-XL-v1.0", 
        text_encoder=None, 
        variant="fp16", 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    print("finish!")
    print("Generating images...",end="")
    generator = torch.Generator().manual_seed(1)
    image = pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_embeds, 
        output_type="pt",
        generator=generator,
        num_inference_steps=steps,
    ).images
    print("finish!")

    #apply watermarks

    pil_image = pt_to_pil(image)
    pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

    pil_image[0].save(f"./{start_dt}_DeepFloyd_IFstage1.png")
    # unload stage1 model
    del pipe
    flush()

    proceed_check = input("Proceed to the next stage?: ")
    if(proceed_check == "no" or proceed_check == "NO" or proceed_check == "N"):
        continue

    # load stage2 model
    print("Loading stage2 model...",end="")
    pipe = DiffusionPipeline.from_pretrained(
        "DeepFloyd/IF-II-L-v1.0", 
        text_encoder=None, # no use of text encoder => memory savings!
        variant="fp16", 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    print("Generating...")
    image = pipe(
        image=image, 
        prompt_embeds=prompt_embeds, 
        negative_prompt_embeds=negative_embeds, 
        output_type="pt",
        generator=generator,
    ).images
    print("finish!")
    pil_image = pt_to_pil(image)
    pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

    pil_image[0].save(f"./{start_dt}_DeepFloyd_IFstage2.png")

    #unload stage2 model
    del pipe
    flush()

    proceed_check = input("Proceed to the next stage?: ")
    if(proceed_check == "no" or proceed_check == "NO" or proceed_check == "N"):
        continue

    # load stage3 model
    print("Loading stage3 model...",end="")
    pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler", 
    torch_dtype=torch.float16, 
    device_map="auto"
    )
    print("finish!")

    pil_image = pipe(prompt, generator=generator, image=image).images

    #apply watermarks
    from diffusers.pipelines.deepfloyd_if import IFWatermarker

    watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
    watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

    pil_image[0].save(f"./{start_dt}_DeepFloyd_IFstage3.png")
    del watermarker
    flush()

LICENSE

Copyright (c) 2023 Stella

This code is based on the original code developed by DeepFloyd and StabilityAI.

The original code is licensed under the following terms:

Copyright (c) 2023 DeepFloyd, StabilityAI

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

1. The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

2. All persons obtaining a copy or substantial portion of the Software,
a modified version of the Software (or substantial portion thereof), or
a derivative work based upon this Software (or substantial portion thereof)
must not delete, remove, disable, diminish, or circumvent any inference filters or
inference filter mechanisms in the Software, or any portion of the Software that
implements any such filters or filter mechanisms.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

This code is subject to the same terms and conditions as the original code.

Thank you for your feedback, and I'm glad to hear that my code was helpful for you!
You can run this script without relying on bitsandbytes by making the following changes to lines 41 and 42:

load_in_8bit=False, 
variant="fp16"

This should allow the script to work without issue, though it may run a bit slower. Please let me know if you have any further questions or encounter any issues!