Spaces:
Runtime error
Runtime error
| import spaces | |
| import os | |
| import requests | |
| import torch | |
| from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler | |
| from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker | |
| 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 | |
| USE_TORCH_COMPILE = 0 | |
| ENABLE_CPU_OFFLOAD = 0 | |
| # 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 | |
| # 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") | |
| class LazyRealESRGAN: | |
| def __init__(self, device, scale): | |
| self.device = device | |
| self.scale = scale | |
| self.model = None | |
| def load_model(self): | |
| if self.model is None: | |
| self.model = RealESRGAN(self.device, scale=self.scale) | |
| self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False) | |
| def predict(self, img): | |
| self.load_model() | |
| return self.model.predict(img) | |
| # Initialize the lazy models | |
| lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2) | |
| lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=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) | |
| if scale == 2: | |
| img = lazy_realesrgan_x2.predict(img) | |
| else: | |
| img = lazy_realesrgan_x4.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 | |
| class ImageProcessor: | |
| def __init__(self): | |
| self.pipe = self.setup_pipeline() | |
| def setup_pipeline(self): | |
| controlnet = ControlNetModel.from_single_file( | |
| "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16 | |
| ) | |
| safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker") | |
| 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, | |
| safety_checker=safety_checker | |
| ) | |
| vae = AutoencoderKL.from_single_file( | |
| "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors", | |
| torch_dtype=torch.float16 | |
| ) | |
| pipe.vae = vae | |
| 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") | |
| pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) | |
| pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4) | |
| return pipe | |
| def process_image(self, 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 = self.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 | |
| # Create an instance of ImageProcessor | |
| image_processor = ImageProcessor() | |
| def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale): | |
| image_processor.pipe = image_processor.pipe.to(device) | |
| image_processor.pipe.unet.set_attn_processor(AttnProcessor2_0()) | |
| 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 = image_processor.process_image(input_image, prompt, negative_prompt, resolution, num_inference_steps, guidance_scale, strength, hdr) | |
| return result | |
| # Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Image Enhancement with Stable Diffusion") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(type="pil", label="Input Image") | |
| run_button = gr.Button("Enhance Image") | |
| with gr.Column(): | |
| output_image = gr.Image(type="pil", label="Enhanced Image") | |
| with gr.Accordion("Advanced Options", open=False): | |
| resolution = gr.Slider(minimum=512, maximum=2048, value=1024, step=64, label="Resolution") | |
| num_inference_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of Inference Steps") | |
| strength = gr.Slider(minimum=0, maximum=1, value=0.35, step=0.05, label="Strength") | |
| hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect") | |
| guidance_scale = gr.Slider(minimum=0, maximum=20, value=3, step=0.5, label="Guidance Scale") | |
| run_button.click(fn=gradio_process_image, | |
| inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale], | |
| outputs=output_image) | |
| demo.launch(share=True) |