Update app.py

app.py

@@ -1,13 +1,16 @@
-# from controlnet_aux import OpenposeDetector
-# from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-# from diffusers import UniPCMultistepScheduler
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import UniPCMultistepScheduler
 import gradio as gr
-# import torch
+import numpy as np
+import torch
 import base64
+import cv2
 from io import BytesIO
 from PIL import Image, ImageFilter
-import cv2
-import numpy as np
+
+# Constants
+low_threshold = 100
+high_threshold = 200
 
 canvas_html = '<pose-maker/>'
 load_js = """
@@ -32,27 +35,24 @@ async (canvas, prompt) => {
 }
 """
 
-# # Models
-# controlnet = ControlNetModel.from_pretrained(
-#     "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
-# )
-# pipe = StableDiffusionControlNetPipeline.from_pretrained(
-#     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
-# )
-# pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-
-# # This command loads the individual model components on GPU on-demand. So, we don't
-# # need to explicitly call pipe.to("cuda").
-# pipe.enable_model_cpu_offload()
+# Models
+controlnet = ControlNetModel.from_pretrained(
+    "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16
+)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
-# # xformers
-# pipe.enable_xformers_memory_efficient_attention()
+# This command loads the individual model components on GPU on-demand. So, we don't
+# need to explicitly call pipe.to("cuda").
+pipe.enable_model_cpu_offload()
 
-# # Generator seed,
-# generator = torch.manual_seed(0)
+# xformers
+pipe.enable_xformers_memory_efficient_attention()
 
-low_threshold = 100
-high_threshold = 200
+# Generator seed,
+generator = torch.manual_seed(0)
 
 def get_canny_filter(image):
     if not isinstance(image, np.ndarray):
@@ -72,17 +72,17 @@ def generate_images(canvas, prompt):
             'RGB').resize((512, 512))
         input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2))
         input_img = get_canny_filter(input_img)
-        # output = pipe(
-        #     prompt,
-        #     input_img,
-        #     generator=generator,
-        #     num_images_per_prompt=3,
-        #     num_inference_steps=20,
-        # )
-        all_outputs = [input_img, input_img, input_img]
-        # all_outputs.append(input_img)
-        # for image in output.images:
-        #     all_outputs.append(image)
+        output = pipe(
+            f'{prompt}, best quality, extremely detailed',
+            input_img,
+            generator=generator,
+            num_images_per_prompt=2,
+            num_inference_steps=20,
+            negative_prompt="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
+        )
+        all_outputs = [input_img]
+        for image in output.images:
+            all_outputs.append(image)
         return all_outputs
     except Exception as e:
         raise gr.Error(str(e))