v2

app.py

@@ -106,13 +106,6 @@ if gr.NO_RELOAD:
         "runwayml/stable-diffusion-inpainting",
         torch_dtype=torch.float16,
     ).to("cuda")
-    
-    print('loading controlnet inpainting pipe')
-    controlnet_inpaint_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-        "runwayml/stable-diffusion-inpainting",
-        controlnet=controlnet,
-        torch_dtype=torch.float16,
-    ).to("cuda")
 
     print("loading preprocessor")
     preprocessor = Preprocessor()
@@ -432,35 +425,21 @@ def process_image(
     guidance_scale,
     seed,
 ):
-    preprocess_start = time.time()
-    print("processing image")
-
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)
+    
     preprocessor.load("NormalBae")
     control_image = preprocessor(
         image=image,
         image_resolution=image_resolution,
         detect_resolution=preprocess_resolution,
     )
-    preprocess_time = time.time() - preprocess_start
-    if style_selection is not None or style_selection != "None":
-        prompt = "Photo from Pinterest of " + apply_style(style_selection) + " " + prompt + "," + a_prompt
+    
+    if style_selection is not None and style_selection != "None":
+        prompt = f"Photo from Pinterest of {apply_style(style_selection)} {prompt},{a_prompt}"
     else:
-        prompt=str(get_prompt(prompt, a_prompt))
-    negative_prompt=str(n_prompt)
-    print(prompt)
-    print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
-    start = time.time()
-    # results = pipe(
-    #     prompt=prompt,
-    #     negative_prompt=negative_prompt,
-    #     guidance_scale=guidance_scale,
-    #     num_images_per_prompt=num_images,
-    #     num_inference_steps=num_steps,
-    #     generator=generator,
-    #     image=control_image,
-    # ).images[0]
+        prompt = str(get_prompt(prompt, a_prompt))
+    negative_prompt = str(n_prompt)
     
     initial_result = pipe(
         prompt=prompt,
@@ -479,28 +458,44 @@ def process_image(
     if furniture_type != "None":
         furniture_mask = generate_furniture_mask(initial_result, furniture_type)
         furniture_prompt = f"A {furniture_type} in the style of {style_selection}"
-        inpainted_image = controlnet_inpaint_pipe(
+        
+        # Apply ControlNet to get the control image for inpainting
+        control_image_inpaint = preprocessor(
+            image=initial_result,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+        
+        # Use the regular inpainting pipeline with ControlNet guidance
+        inpainted_image = inpaint_pipe(
             prompt=furniture_prompt,
             image=initial_result,
             mask_image=furniture_mask,
-            control_image=control_image,
             negative_prompt=negative_prompt,
             num_inference_steps=num_steps,
             guidance_scale=guidance_scale,
             generator=generator,
         ).images[0]
+        
+        # Apply ControlNet guidance on the inpainted result
+        final_result = pipe(
+            prompt=furniture_prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=1,
+            num_inference_steps=num_steps,
+            generator=generator,
+            image=control_image_inpaint,
+            controlnet_conditioning_scale=0.5,  # Adjust this value as needed
+        ).images[0]
     else:
-        inpainted_image = initial_result
-    
-    print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
-    torch.cuda.empty_cache()
+        final_result = initial_result
     
-    # upload block
     timestamp = int(time.time())
     img_path = f"{timestamp}.jpg"
     results_path = f"{timestamp}_out.jpg"
     imageio.imsave(img_path, image)
-    imageio.imsave(results_path, results)
+    imageio.imsave(results_path, final_result)
     api.upload_file(
         path_or_fileobj=img_path,
         path_in_repo=img_path,
@@ -517,7 +512,7 @@ def process_image(
         token=API_KEY,
         run_as_future=True,
     )
-    return inpainted_image
+    return final_result
 
 if prod:
     demo.queue(max_size=20).launch(server_name="localhost", server_port=port)