Update app.py

app.py

@@ -15,7 +15,11 @@ import gradio as gr
 USE_TORCH_COMPILE = 0
 ENABLE_CPU_OFFLOAD = 0
 
-# Function to download files
+# Set up the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# Function to download files (from the example)
 def download_file(url, folder_path, filename):
     if not os.path.exists(folder_path):
         os.makedirs(folder_path)
@@ -34,73 +38,26 @@ def download_file(url, folder_path, filename):
             print(f"Error downloading the file. Status code: {response.status_code}")
 
 # Download necessary models and files
+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()
 
 
-    
-# 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):
@@ -215,26 +172,7 @@ class ImageProcessor:
 
         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()
 
 @spaces.GPU