SDXL + FLUX image beautifier

This project is aimed to train or inference Diffusion models to achive low-quality photo restoration without composition and detail loss.

How to install

1. Clone this repo and install PyTorch (I used python 3.10)
2. If you want to use my features with brownian noise, you should install my fork of diffusers libriary as editable installation (because they didn't want to merge for some reasons)
3. Install Flux repo that i use FLUX .1 dev
4. Also for my experiments i'm using Compel, consider installing it too!
5. Optional. If you are interested in testing LongClip, also install it or comment out
6. Install https://github.com/xhinker/sd_embed, works the same as Compel, but usable with FLUX
7. Download juggernautXL_juggXIByRundiffusion model from CivAI somewhere, because my scripts use that model

How to use

It's research repo, but some things work out of box. For example, if you want to generate better version of your "LQ" photo, consider using that scripts:

1. baseline_brownian_karras.py - that code is without CLI, but it's easy to understand. It will generate "HQ" image with SDXL but without refiner
2. You may do refinement with flux_diffusers_pipe.py

All scripts use "Path headers" such as:

output_folder = "/home/dataset/Train_data/refined_pictures/"
control_folder = "/home/dataset/Train_data/output_pictures/"
captions_folder= "/home/dataset/Train_data/captions/"

You should consider changing them, as well as DEVICE env and PATH in those codelines:

1. DL_PictureBeauty/baseline_brownian_karras.py

   Lines 25 to 31 in 6b9bd4e

   	pipeline = StableDiffusionXLControlNetPipeline.from_single_file(
   	# "/home/web_ui/stable-diffusion-webui/models/Stable-diffusion/cyberrealisticXL_v22.safetensors",
   	"/home/web_ui/stable-diffusion-webui/models/Stable-diffusion/juggernautXL_juggXIByRundiffusion.safetensors",
   	# "/home/web_ui/stable-diffusion-webui/models/Stable-diffusion/jibMixRealisticXL_v160Aphrodite.safetensors",
   	use_safetensors=True,
   	torch_dtype=torch.bfloat16,
   	controlnet = controlnets)

2. DL_PictureBeauty/baseline_brownian_karras.py

   Lines 20 to 22 in 6b9bd4e

   	controlnets = ControlNetModel.from_pretrained(
   	"/home/web_ui/stable-diffusion-webui/models/ControlNet/controlnet_tile",
   	torch_dtype=torch.bfloat16, use_safetensors=True)

Training scripts

I've tried to distill that pipeline, but it didn't work :(

Scripts in that folder are "fixed" version of diffusers examples. You can look up SDXL training as well. Exaples of .sh scipts are provided.

Main complications in traing were:

1. SDXL Controlnet on 1000+ pictures gives bad results
2. FLUX Lora Controlnet (as they released in their HF repo) training was giving once again bad results
3. FLUX original Controlnet (as done in X-Labs repo) training was not working at all with LogNorm (0,1) dist, also tried U[0,1]

Reusable changes:

1. Check latest commit in FLUX .1 dev, there's gradient checkpoiting implemented for FLUX training