Update app.py

app.py

@@ -102,43 +102,6 @@ download_file(
 # Set up the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load ControlNet model
-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")
-
-# 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,
-    safety_checker=safety_checker
-)
-
-# 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
-
-# Enable FreeU
-pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
-
 class LazyRealESRGAN:
     def __init__(self, device, scale):
         self.device = device
@@ -217,51 +180,90 @@ def create_hdr_effect(original_image, hdr):
 
     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
+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()
 
 @spaces.GPU
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
-    pipe = pipe.to(device)
-    pipe.unet.set_attn_processor(AttnProcessor2_0())
+    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 = process_image(input_image, prompt, negative_prompt, resolution, num_inference_steps, guidance_scale, strength, hdr)
+    result = image_processor.process_image(input_image, prompt, negative_prompt, resolution, num_inference_steps, guidance_scale, strength, hdr)
     return result
-    
-# Simple options
-simple_options = [
-    gr.Image(type="pil", label="Input Image"),
-    gr.Slider(minimum=2048, maximum=3072, step=512, value=2048, label="Resolution"),
-    gr.Slider(minimum=10, maximum=100, step=10, value=20, label="Inference Steps"),
-    gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.35, label="Strength"),
-    gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0, label="HDR"),
-    gr.Slider(minimum=1, maximum=10, step=0.1, value=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()
+# 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)