Update app.py

app.py

@@ -46,6 +46,15 @@ else:
 def hex_to_rgb(hex_code):
     hex_code = hex_code.lstrip('#')
     return tuple(int(hex_code[i:i+2], 16) for i in (0, 2, 4))
+
+def get_depth_map(image, depth_estimator):
+    image = depth_estimator(image)["depth"]
+    image = np.array(image)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    detected_map = torch.from_numpy(image).float() / 255.0
+    depth_map = detected_map.permute(2, 0, 1)
+    return depth_map
     
 @app.get("/")
 async def root():
@@ -107,11 +116,15 @@ async def generate_new_img(hex_code: str, prompt: str = Query(..., description="
             ad_prompt = f"""Your system prompt is this: {prompt} Consider your system prompt first. 
                             Then from the initial image create a new image that will attract customers to put in an (ad template) 
                             Also, use this RGB color {hex_to_rgb(hex_code)} as a tone in the image while image is still recognized as it is original."""
+            print(f"uploaded image type: {type(uploaded_image)}")
+            depth_estimator = pipeline("depth-estimation")
+            depth_map = get_depth_map(uploaded_image, depth_estimator).unsqueeze(0).half().to(device)
             if device=="cpu":
                 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float32)
                 pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
                     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32
                 ).to(device)
+
             else:
                 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16)
                 pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
@@ -130,7 +143,7 @@ async def generate_new_img(hex_code: str, prompt: str = Query(..., description="
             print(f"init_image type: {type(init_image)}")
 
             image = pipe(
-                ad_prompt, image=init_image
+                ad_prompt, image=init_image, control_image=depth_map
             ).images[0]
             print(f"image type: {type(image)}")