Update app.py

app.py

@@ -10,11 +10,10 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-import cv2
 from PIL import Image
 from io import BytesIO
 from diffusers.utils import load_image
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting
+from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
@@ -28,7 +27,6 @@ ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
 ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
 ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_IMG2IMG", "1") == "1"
 ENABLE_USE_INPAINTING = os.getenv("ENABLE_USE_INPAINTING", "1") == "1"
-ENABLE_USE_CONTROLNET = os.getenv("ENABLE_USE_CONTROLNET", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -61,23 +59,39 @@ def generate(
     lora_scale: float = 0.7,
     use_img2img: bool = False,
     use_inpainting: bool = False,
-    use_controlnet: bool = False,
     url = '',
     img_url = '',
     mask_url = '',
 ):
     if torch.cuda.is_available():
-            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-            pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16)
+            
+        if not use_img2img:
+            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+            
+        if use_img2img:
+            pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
 
-            image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")
-            np_image = np.array(image)
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+                
+        if use_inpainting:
+            pipe = AutoPipelineForInpainting.from_pretrained(model, torch_dtype=torch.float16)
             
-            # get canny image
-            np_image = cv2.Canny(np_image, 100, 200)
-            np_image = np_image[:, :, None]
-            np_image = np.concatenate([np_image, np_image, np_image], axis=2)
-            canny_image = Image.fromarray(np_image)
+            response = requests.get(url)
+            init_image = Image.open(BytesIO(response.content)).convert("RGB")
+            init_image = init_image.resize((width, height))
+
+            image_init = load_image(img_url)
+            mask_image = load_image(mask_url)
+            
+        if use_lora:
+            pipe.load_lora_weights(lora)
+            pipe.fuse_lora(lora_scale)
                                    
         if ENABLE_CPU_OFFLOAD:
             pipe.enable_model_cpu_offload()
@@ -97,15 +111,51 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    if use_controlnet:
+    if use_inpainting:
         image = pipe(
-            "futuristic-looking woman",
-            num_inference_steps=20,
+            prompt=prompt,
+            image=image_init,
+            mask_image=mask_image,
+            strength=strength_img2img,
+            negative_prompt=negative_prompt,
+            prompt_2=prompt_2,
+            width=width,
+            height=height,
+            negative_prompt_2=negative_prompt_2,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
             generator=generator,
-            image=image,
-            control_image=canny_image,
         ).images[0]
         return image
+    elif use_img2img:
+        images = pipe(
+            prompt=prompt,
+            image=init_image,
+            strength=strength_img2img,
+            negative_prompt=negative_prompt,
+            prompt_2=prompt_2,
+            negative_prompt_2=negative_prompt_2,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
+            generator=generator,
+            output_type="pil",
+        ).images[0]
+        return images
+    else:
+        return pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            prompt_2=prompt_2,
+            negative_prompt_2=negative_prompt_2,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
+            generator=generator,
+            output_type="pil",
+        ).images[0]
 
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
@@ -139,7 +189,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         result = gr.Image(label="Result", show_label=False)
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
-            use_controlnet = gr.Checkbox(label='Use Controlnet', value=False, visible=ENABLE_USE_CONTROLNET)
             use_inpainting = gr.Checkbox(label='Use Inpainting', value=False, visible=ENABLE_USE_INPAINTING)
             use_img2img = gr.Checkbox(label='Use Img2Img', value=False, visible=ENABLE_USE_IMG2IMG)
             use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
@@ -274,13 +323,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
-    use_controlnet.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_controlnet,
-        outputs=img_url,
-        queue=False,
-        api_name=False,
-    )
 
     gr.on(
         triggers=[
@@ -322,7 +364,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
             url,
             img_url,
             mask_url,
-            use_controlnet,
         ],
         outputs=result,
         api_name="run",