app.py

import spaces
import os
import requests
import torch
from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
from diffusers.models import AutoencoderKL
from PIL import Image
from RealESRGAN import RealESRGAN
import cv2
import numpy as np
from diffusers.models.attention_processor import AttnProcessor2_0
import gradio as gr

# Constants
SD15_WEIGHTS = "weights"
CONTROLNET_CACHE = "controlnet-cache"
SCHEDULERS = {
    "DDIM": DDIMScheduler,
    "DPMSolverMultistep": DPMSolverMultistepScheduler,
    "K_EULER_ANCESTRAL": EulerAncestralDiscreteScheduler,
    "K_EULER": EulerDiscreteScheduler,
}

# Function to download files
def download_file(url, folder_path, filename):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    file_path = os.path.join(folder_path, filename)

    if os.path.isfile(file_path):
        print(f"File already exists: {file_path}")
    else:
        response = requests.get(url, stream=True)
        if response.status_code == 200:
            with open(file_path, 'wb') as file:
                for chunk in response.iter_content(chunk_size=1024):
                    file.write(chunk)
            print(f"File successfully downloaded and saved: {file_path}")
        else:
            print(f"Error downloading the file. Status code: {response.status_code}")

# Download necessary models and files

@spaces.GPU
def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
    prompt = "masterpiece, best quality, highres"
    negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
    result = process_image(input_image, prompt, negative_prompt, resolution, num_inference_steps, guidance_scale, strength, hdr)
    return result
    
# MODEL
download_file(
    "https://huggingface.co/dantea1118/juggernaut_reborn/resolve/main/juggernaut_reborn.safetensors?download=true",
    "models/models/Stable-diffusion",
    "juggernaut_reborn.safetensors"
)

# UPSCALER

download_file(
    "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth?download=true",
    "models/upscalers/",
    "RealESRGAN_x2.pth"
)

download_file(
    "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth?download=true",
    "models/upscalers/",
    "RealESRGAN_x4.pth"
)

# NEGATIVE
download_file(
    "https://huggingface.co/philz1337x/embeddings/resolve/main/verybadimagenegative_v1.3.pt?download=true",
    "models/embeddings",
    "verybadimagenegative_v1.3.pt"
)
download_file(
    "https://huggingface.co/datasets/AddictiveFuture/sd-negative-embeddings/resolve/main/JuggernautNegative-neg.pt?download=true",
    "models/embeddings",
    "JuggernautNegative-neg.pt"
)

# LORA

download_file(
    "https://huggingface.co/philz1337x/loras/resolve/main/SDXLrender_v2.0.safetensors?download=true",
    "models/Lora",
    "SDXLrender_v2.0.safetensors"
)
download_file(
    "https://huggingface.co/philz1337x/loras/resolve/main/more_details.safetensors?download=true",
    "models/Lora",
    "more_details.safetensors"
)

# CONTROLNET

download_file(
    "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile.pth?download=true",
    "models/ControlNet",
    "control_v11f1e_sd15_tile.pth"
)

# VAE

download_file(
    "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors?download=true",
    "models/VAE",
    "vae-ft-mse-840000-ema-pruned.safetensors"
)

# Set up the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load ControlNet model
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/control_v11f1e_sd15_tile", torch_dtype=torch.float16
)

# Load the Stable Diffusion pipeline with Juggernaut Reborn model
model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
    model_path,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    use_safetensors=True
)

# Load and set VAE
vae = AutoencoderKL.from_single_file(
    "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
    torch_dtype=torch.float16
)
pipe.vae = vae


# Load embeddings and LoRA models
pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
pipe.load_lora_weights("models/Lora/SDXLrender_v2.0.safetensors")
pipe.fuse_lora(lora_scale=0.5)
pipe.load_lora_weights("models/Lora/more_details.safetensors")
# Set up the scheduler
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# Move the pipeline to the device and enable memory efficient attention
pipe = pipe.to(device)
pipe.unet.set_attn_processor(AttnProcessor2_0())

# Enable FreeU
pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)

def resize_and_upscale(input_image, resolution):
    scale = 2
    if resolution == 2048:
        init_w = 1024
    elif resolution == 2560:
        init_w = 1280
    elif resolution == 3072:
        init_w = 1536
    else:
        init_w = 1024
        scale = 4

    input_image = input_image.convert("RGB")
    W, H = input_image.size
    k = float(init_w) / 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)
    model = RealESRGAN(device, scale=scale)
    model.load_weights(f'models/upscalers/RealESRGAN_x{scale}.pth', download=False)
    img = model.predict(img)
    return img

def calculate_brightness_factors(hdr_intensity):
    factors = [1.0] * 9
    if hdr_intensity > 0:
        factors = [1.0 - 0.9 * hdr_intensity, 1.0 - 0.7 * hdr_intensity, 1.0 - 0.45 * hdr_intensity,
                   1.0 - 0.25 * hdr_intensity, 1.0, 1.0 + 0.2 * hdr_intensity,
                   1.0 + 0.4 * hdr_intensity, 1.0 + 0.6 * hdr_intensity, 1.0 + 0.8 * hdr_intensity]
    return factors

def pil_to_cv(pil_image):
    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

def adjust_brightness(cv_image, factor):
    hsv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv_image)
    v = np.clip(v * factor, 0, 255).astype('uint8')
    adjusted_hsv = cv2.merge([h, s, v])
    return cv2.cvtColor(adjusted_hsv, cv2.COLOR_HSV2BGR)

def create_hdr_effect(original_image, hdr):
    cv_original = pil_to_cv(original_image)
    brightness_factors = calculate_brightness_factors(hdr)
    images = [adjust_brightness(cv_original, factor) for factor in brightness_factors]

    merge_mertens = cv2.createMergeMertens()
    hdr_image = merge_mertens.process(images)
    hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
    hdr_image_pil = Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))

    return hdr_image_pil

def process_image(input_image, prompt, negative_prompt, resolution=2048, num_inference_steps=50, guidance_scale=3, strength=0.35, hdr=0):
    condition_image = resize_and_upscale(input_image, resolution)
    condition_image = create_hdr_effect(condition_image, hdr)

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=condition_image,
        control_image=condition_image,
        width=condition_image.size[0],
        height=condition_image.size[1],
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=torch.manual_seed(0),
    ).images[0]

    return result



# Simple options
simple_options = [
    gr.Image(type="pil", label="Input Image"),
    gr.Slider(minimum=2048, maximum=3072, step=512, default=2048, label="Resolution"),
    gr.Slider(minimum=10, maximum=100, step=10, default=20, label="Inference Steps"),
    gr.Slider(minimum=0.0, maximum=1.0, step=0.05, default=0.35, label="Strength"),
    gr.Slider(minimum=0.0, maximum=1.0, step=0.1, default=0, label="HDR"),
    gr.Slider(minimum=1, maximum=10, step=0.1, default=3, label="Guidance Scale")
]

# Create the Gradio interface
iface = gr.Interface(
    fn=gradio_process_image,
    inputs=simple_options,
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Image Processing with Stable Diffusion",
    description="Upload an image and adjust the settings to process it using Stable Diffusion."
)

iface.launch()