Stella2211/DeepFloyd_IF_VRAM12GB.py

## 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.
	# 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()
	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.