Update app.py

app.py

@@ -115,16 +115,16 @@ async def generate_new_img(hex_code: str, prompt: str = Query(..., description="
             # Generate the image using the text-to-image model
             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 hex code {hex_to_rgb(hex_color)} as a tone in the image while image is still recognized as it is original."""
+                            Also, use this RGB color {hex_to_rgb(hex_color)} as a tone in the image while image is still recognized as it is original."""
             if device=="cpu":
-                controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32, use_safetensors=True)
-                pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float32, use_safetensors=True)
+                pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
                     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32, use_safetensors=True
                 ).to(device)
             else:
-                controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16, use_safetensors=True)
-                pipe = StableDiffusionControlNetPipeline.from_pretrained(
-                    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32, use_safetensors=True
+                controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16, use_safetensors=True)
+                pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
                 ).to(device)
                 pipe.enable_model_cpu_offload()
                 
@@ -156,7 +156,7 @@ async def generate_new_img(hex_code: str, prompt: str = Query(..., description="
 async def create_ad_template(punchline: str, punchline_color:str, button_text:str,button_color:str):
     if uploaded_image is not None and uploaded_logo is not None and created_image is not None:
         # Create LOGO
-        
+        print("Drawing")
         #logo_width, logo_height = logo_image.size
         logo_width, logo_height = 100, 100  # Desired size for the logo
         logo_image = logo_image.resize((logo_width, logo_height))
@@ -171,16 +171,20 @@ async def create_ad_template(punchline: str, punchline_color:str, button_text:st
 
         # Create a blank canvas for the ad template
         ad_template = Image.new("RGB", (800, 600), "#FFFFFF")
+        print("Template created")
         # Add logo
         ad_template.paste(logo_image, logo_image_location)
+        print("Logo added")
         # Add generated image
         ad_template.paste(generated_image, generated_image_location)
+        print("Image added")
         # Add the text at the bottom of the ad template
         draw = ImageDraw.Draw(ad_template)
         font = ImageFont.truetype("arial.ttf", 24)
         text_width, text_height = draw.textsize(punchline, font=font)
         text_position = (400 - text_width / 2, 500 - text_height / 2)
         draw.text(text_position, punchline, font=font, fill=punchline_color)
+        print("Punchline added")
         # Add the button at the bottom of the ad template
         button_width, button_height = 200, 50
         button_position = ((400 - button_width / 2), (550 - button_height / 2))
@@ -190,15 +194,15 @@ async def create_ad_template(punchline: str, punchline_color:str, button_text:st
         text_position = (button_position[0] + (button_width - text_width) / 2, button_position[1] + (button_height - text_height) / 2)
         rect_text_color = (255, 255, 255)  # Text color within the rectangle
         draw.text(text_position, button_text, fill=rect_text_color, font=font)
-
+        
         border_width = 2  # Border width
         border_positions = [0, 0, 800, 600]
         draw.rectangle(border_positions, outline="#D60D0D", width=border_width)
-
+        print("Button added")
         ad_template_data = io.BytesIO()
         ad_template.save(ad_template_data, format="PNG")
         ad_template_data.seek(0)
-        
+        print("Template finished")
         return StreamingResponse(ad_template_data, media_type="image/png")
         
     else: