Update app.py

app.py

@@ -1,11 +1,16 @@
 import spaces
 from typing import Tuple, Union, List
 import os
-
+import time
 import numpy as np
 from PIL import Image
 
 import torch
+
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from diffusers.models import AutoencoderKL
+from diffusers.models.attention_processor import AttnProcessor2_0
 from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
 from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
 from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
@@ -33,6 +38,188 @@ css = """
 """
 
 
+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}")
+
+def download_models():
+    models = {
+        "MODEL": ("https://huggingface.co/dantea1118/juggernaut_reborn/resolve/main/juggernaut_reborn.safetensors?download=true", "models/models/Stable-diffusion", "juggernaut_reborn.safetensors"),
+        "UPSCALER_X2": ("https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth?download=true", "models/upscalers/", "RealESRGAN_x2.pth"),
+        "UPSCALER_X4": ("https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth?download=true", "models/upscalers/", "RealESRGAN_x4.pth"),
+        "NEGATIVE_1": ("https://huggingface.co/philz1337x/embeddings/resolve/main/verybadimagenegative_v1.3.pt?download=true", "models/embeddings", "verybadimagenegative_v1.3.pt"),
+        "NEGATIVE_2": ("https://huggingface.co/datasets/AddictiveFuture/sd-negative-embeddings/resolve/main/JuggernautNegative-neg.pt?download=true", "models/embeddings", "JuggernautNegative-neg.pt"),
+        "LORA_1": ("https://huggingface.co/philz1337x/loras/resolve/main/SDXLrender_v2.0.safetensors?download=true", "models/Lora", "SDXLrender_v2.0.safetensors"),
+        "LORA_2": ("https://huggingface.co/philz1337x/loras/resolve/main/more_details.safetensors?download=true", "models/Lora", "more_details.safetensors"),
+        "CONTROLNET": ("https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile.pth?download=true", "models/ControlNet", "control_v11f1e_sd15_tile.pth"),
+        "VAE": ("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"),
+    }
+
+    for model, (url, folder, filename) in models.items():
+        download_file(url, folder, filename)
+
+download_models()
+
+def timer_func(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(f"{func.__name__} took {end_time - start_time:.2f} seconds")
+        return result
+    return wrapper
+
+class LazyLoadPipeline:
+    def __init__(self):
+        self.pipe = None
+
+    @timer_func
+    def load(self):
+        if self.pipe is None:
+            print("Starting to load the pipeline...")
+            self.pipe = self.setup_pipeline()
+            print(f"Moving pipeline to device: {device}")
+            self.pipe.to(device)
+            if USE_TORCH_COMPILE:
+                print("Compiling the model...")
+                self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+    @timer_func
+    def setup_pipeline(self):
+        print("Setting up the pipeline...")
+        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.fuse_lora(lora_scale=1.)
+        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 __call__(self, *args, **kwargs):
+        return self.pipe(*args, **kwargs)
+
+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)
+
+lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
+lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
+
+@timer_func
+def resize_and_upscale(input_image, resolution):
+    scale = 2 if resolution <= 2048 else 4
+    input_image = input_image.convert("RGB")
+    W, H = input_image.size
+    k = float(resolution) / min(H, W)
+    H = int(round(H * k / 64.0)) * 64
+    W = int(round(W * k / 64.0)) * 64
+    img = input_image.resize((W, H), resample=Image.LANCZOS)
+    if scale == 2:
+        img = lazy_realesrgan_x2.predict(img)
+    else:
+        img = lazy_realesrgan_x4.predict(img)
+    return img
+
+@timer_func
+def create_hdr_effect(original_image, hdr):
+    if hdr == 0:
+        return original_image
+    cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
+    factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
+               1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
+               1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
+    images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
+    merge_mertens = cv2.createMergeMertens()
+    hdr_image = merge_mertens.process(images)
+    hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
+    return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
+
+lazy_pipe = LazyLoadPipeline()
+lazy_pipe.load()
+
+def prepare_image(input_image, resolution, hdr):
+    condition_image = resize_and_upscale(input_image, resolution)
+    condition_image = create_hdr_effect(condition_image, hdr)
+    return condition_image
+
+@spaces.GPU
+@timer_func
+def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
+    print("Starting image processing...")
+    torch.cuda.empty_cache()
+    
+    condition_image = prepare_image(input_image, resolution, hdr)
+    
+    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"
+    
+    options = {
+        "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.Generator(device=device).manual_seed(0),
+    }
+    
+    print("Running inference...")
+    result = lazy_pipe(**options).images[0]
+    print("Image processing completed successfully")
+    
+    # Convert input_image and result to numpy arrays
+    input_array = np.array(input_image)
+    result_array = np.array(result)
+    
+    return [input_array, result_array]
+    
+
+
 def filter_items(
     colors_list: Union[List, np.ndarray],
     items_list: Union[List, np.ndarray],
@@ -267,12 +454,13 @@ class ControlNetDepthDesignModelMulti:
 
 
 def create_demo(model):
-    gr.Markdown("### Stable Design demo")
+    gr.Markdown("### Just try zeroGPU")
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')
             input_text = gr.Textbox(label='Prompt', placeholder='Please upload your image first', lines=2)
             with gr.Accordion('Advanced options', open=False):
+                
                 num_steps = gr.Slider(label='Steps',
                                       minimum=1,
                                       maximum=50,
@@ -305,6 +493,13 @@ def create_demo(model):
                 n_prompt = gr.Textbox(
                     label='Negative Prompt',
                     value="window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner")
+
+                resolution = gr.Slider(minimum=256, maximum=2048, value=512, step=256, label="Resolution")
+                num_inference_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Inference Steps")
+                strength = gr.Slider(minimum=0, maximum=1, value=0.4, step=0.01, 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")
+        
             submit = gr.Button("Submit")
         
         with gr.Column():
@@ -326,6 +521,7 @@ def create_demo(model):
                            inputs=[input_image, input_text], cache_examples=False)
 
 
+
 controlnet_depth= ControlNetModel.from_pretrained(
     "controlnet_depth", torch_dtype=dtype, use_safetensors=True)
 controlnet_seg = ControlNetModel.from_pretrained(
@@ -353,18 +549,18 @@ depth_feature_extractor, depth_estimator = get_depth_pipeline()
 depth_estimator = depth_estimator.to(device)
 
 
+
 def main():
     model = ControlNetDepthDesignModelMulti()
     print('Models uploaded successfully')
     
-    title = "# StableDesign"
+    title = "# Just try zeroGPU"
     description = """
-    Download RoomAI on the App Store to experience the full functionality
+    For test only
     """
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)        
-        create_demo(model)
 
 
     demo.queue().launch(share=False)