Rollback

app.py

@@ -5,6 +5,7 @@ import warnings
 import random
 import time
 import logging
+import fal_client
 import base64
 import numpy as np
 import math
@@ -53,6 +54,10 @@ CONFIG_FILE = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
 GROUNDINGDINO_CHECKPOINT = "groundingdino_swint_ogc.pth"
 SAM_CHECKPOINT = 'sam_hq_vit_l.pth'
 OUTPUT_DIR = "outputs"
+# FAL_KEY = os.getenv("FAL_KEY")
+# UPLOAD_DIR = "./tmp/images"
+
+# os.makedirs(UPLOAD_DIR, exist_ok=True)
 
 # Global variables for model caching
 _models = {
@@ -106,6 +111,7 @@ class ModelManager:
                 _models['sam_predictor'] = SamPredictor(sam)
                 
                 logger.info(f"SAM-HQ model loaded in {time.time() - start_time:.2f} seconds")
+
                 
         except Exception as e:
             logger.error(f"Error loading {model_name} model: {e}")
@@ -193,7 +199,9 @@ def get_grounding_output(
 
 def draw_mask(mask: np.ndarray, draw: ImageDraw.Draw) -> None:
     """Draw mask on image"""
+
     color = (255, 255, 255, 255)
+
     nonzero_coords = np.transpose(np.nonzero(mask))
     for coord in nonzero_coords:
         draw.point(coord[::-1], fill=color)
@@ -230,8 +238,8 @@ def run_grounded_sam(input_image):
         # Process input image
         if isinstance(input_image, dict):
             # Input from gradio sketch component
-            scribble = np.array(input_image["mask"]) if "mask" in input_image else None
-            image_pil = input_image["image"].convert("RGB") if "image" in input_image else None
+            scribble = np.array(input_image["mask"])
+            image_pil = input_image["image"].convert("RGB")
         else:
             # Direct image input
             image_pil = input_image.convert("RGB") if input_image else None
@@ -239,7 +247,7 @@ def run_grounded_sam(input_image):
             
         if image_pil is None:
             logger.error("No input image provided")
-            return Image.new('RGBA', (400, 300), color=(0, 0, 0, 0))
+            return [Image.new('RGB', (400, 300), color='gray')]
 
         # Transform image for GroundingDINO
         transformed_image = transform_image(image_pil)
@@ -254,32 +262,20 @@ def run_grounded_sam(input_image):
             transformed_image, text_prompt, box_threshold, text_threshold
         )
 
-        # Fix: Handle case when no boxes are detected
-        if len(boxes_filt) == 0 or boxes_filt.nelement() == 0:
-            logger.warning("No boxes detected")
-            # Create a simple fallback mask - a circle in the center
-            mask_image = Image.new('RGBA', size, color=(0, 0, 0, 0))
-            mask_draw = ImageDraw.Draw(mask_image)
-            center_x, center_y = W // 2, H // 2
-            radius = min(W, H) // 4
-            mask_draw.ellipse((center_x - radius, center_y - radius, 
-                               center_x + radius, center_y + radius), 
-                              fill=(255, 255, 255, 255))
-            return mask_image
-
-        # Scale boxes to image dimensions
-        for i in range(boxes_filt.size(0)):
-            boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
-            boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
-            boxes_filt[i][2:] += boxes_filt[i][:2]
-        
-        # Apply non-maximum suppression if we have multiple boxes
-        if boxes_filt.size(0) > 1:
-            logger.info(f"Before NMS: {boxes_filt.shape[0]} boxes")
-            nms_idx = torchvision.ops.nms(boxes_filt, scores, iou_threshold).numpy().tolist()
-            boxes_filt = boxes_filt[nms_idx]
-            pred_phrases = [pred_phrases[idx] for idx in nms_idx]
-            logger.info(f"After NMS: {boxes_filt.shape[0]} boxes")
+        if boxes_filt is not None:
+            # Scale boxes to image dimensions
+            for i in range(boxes_filt.size(0)):
+                boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
+                boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
+                boxes_filt[i][2:] += boxes_filt[i][:2]
+            
+            # Apply non-maximum suppression if we have multiple boxes
+            if boxes_filt.size(0) > 1:
+                logger.info(f"Before NMS: {boxes_filt.shape[0]} boxes")
+                nms_idx = torchvision.ops.nms(boxes_filt, scores, iou_threshold).numpy().tolist()
+                boxes_filt = boxes_filt[nms_idx]
+                pred_phrases = [pred_phrases[idx] for idx in nms_idx]
+                logger.info(f"After NMS: {boxes_filt.shape[0]} boxes")
         
         # Load SAM model
         ModelManager.load_model('sam')
@@ -289,18 +285,11 @@ def run_grounded_sam(input_image):
         image = np.array(image_pil)
         sam_predictor.set_image(image)
         
-        # Run SAM
+        # Run SAM 
+        # Use boxes for these task types
         if boxes_filt.size(0) == 0:
-            logger.warning("No boxes detected after NMS")
-            # Create a simple fallback mask
-            mask_image = Image.new('RGBA', size, color=(0, 0, 0, 0))
-            mask_draw = ImageDraw.Draw(mask_image)
-            center_x, center_y = W // 2, H // 2
-            radius = min(W, H) // 4
-            mask_draw.ellipse((center_x - radius, center_y - radius, 
-                               center_x + radius, center_y + radius), 
-                              fill=(255, 255, 255, 255))
-            return mask_image
+            logger.warning("No boxes detected")
+            return [image_pil, Image.new('RGBA', size, color=(0, 0, 0, 0))]
         
         transformed_boxes = sam_predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(device)
         
@@ -330,47 +319,30 @@ def run_grounded_sam(input_image):
     
     except Exception as e:
         logger.error(f"Error in run_grounded_sam: {e}")
-        # Return transparent image on error
+        # Return original image on error
         if isinstance(input_image, dict) and "image" in input_image:
-            return Image.new('RGBA', input_image["image"].size, color=(0, 0, 0, 0))
+            return [input_image["image"], Image.new('RGBA', input_image["image"].size, color=(0, 0, 0, 0))]
         elif isinstance(input_image, Image.Image):
-            return Image.new('RGBA', input_image.size, color=(0, 0, 0, 0))
+            return [input_image, Image.new('RGBA', input_image.size, color=(0, 0, 0, 0))]
         else:
-            return Image.new('RGBA', (400, 300), color=(0, 0, 0, 0))
-
+            return [Image.new('RGB', (400, 300), color='gray'), Image.new('RGBA', (400, 300), color=(0, 0, 0, 0))]
 
 def split_image_with_alpha(image):
-    """Ensure image is RGB and return a copy"""
-    if image.mode == 'RGBA':
-        # Create a white background
-        background = Image.new('RGB', image.size, (255, 255, 255))
-        # Composite the image with alpha onto the background
-        return Image.alpha_composite(background.convert('RGBA'), image).convert('RGB')
-    else:
-        return image.convert("RGB")
-
+    image = image.convert("RGB")
+    return image
 
 def gaussian_blur(image, radius=10):
     """Apply Gaussian blur to image."""
-    # Ensure image is in RGB mode
-    image = image.convert("RGB")
-    blurred = image.filter(ImageFilter.GaussianBlur(radius=radius))
+    blurred = image.filter(ImageFilter.GaussianBlur(radius=10))
     return blurred
 
-
 def invert_image(image):
-    """Invert image colors"""
-    # Ensure image is in RGB mode for inversion
-    image = image.convert("RGB")
     img_inverted = ImageOps.invert(image)
     return img_inverted
 
-
 def expand_mask(mask, expand, tapered_corners):
-    """Expand or contract a mask with proper error handling"""
-    # Ensure mask is in grayscale mode 'L'
-    if mask.mode != 'L':
-        mask = mask.convert("L")
+    # Ensure mask is in grayscale (mode 'L')
+    mask = mask.convert("L")
     
     # Convert to NumPy array
     mask_np = np.array(mask)
@@ -392,45 +364,31 @@ def expand_mask(mask, expand, tapered_corners):
     # Convert back to PIL image
     return Image.fromarray(mask_np, mode="L")
 
-
 def image_blend_by_mask(image_a, image_b, mask, blend_percentage):
-    """Blend two images using a mask with proper error handling"""
-    # Ensure both images are in RGB mode
-    image_a = image_a.convert('RGB')
-    image_b = image_b.convert('RGB')
-    
-    # Ensure mask is in grayscale mode 'L'
-    if mask.mode != 'L':
-        mask = mask.convert('L')
-    
-    # Invert mask for proper blending
-    mask = ImageOps.invert(mask)
+
+    mask = ImageOps.invert(mask.convert('L'))
 
     # Mask image
     masked_img = Image.composite(image_a, image_b, mask)
 
     # Blend image
     blend_mask = Image.new(mode="L", size=image_a.size,
-                          color=(round(blend_percentage * 255)))
+                           color=(round(blend_percentage * 255)))
     blend_mask = ImageOps.invert(blend_mask)
     img_result = Image.composite(image_a, masked_img, blend_mask)
 
-    return img_result
+    del image_a, image_b, blend_mask, mask
 
+    return img_result
 
 def blend_images(image_a, image_b, blend_percentage):
-    """Blend two images with proper format handling"""
-    # Ensure both images are in RGBA mode
+    """Blend img_b over image_a using the normal mode with a blend percentage."""
     img_a = image_a.convert("RGBA")
     img_b = image_b.convert("RGBA")
-    
-    # Fix: Check if sizes match and resize if needed
-    if img_a.size != img_b.size:
-        logger.warning(f"Image sizes don't match: {img_a.size} vs {img_b.size}. Resizing second image.")
-        img_b = img_b.resize(img_a.size, Image.LANCZOS)
 
     # Blend img_b over img_a using alpha_composite (normal blend mode)
     out_image = Image.alpha_composite(img_a, img_b)
+
     out_image = out_image.convert("RGB")
 
     # Create blend mask
@@ -438,9 +396,13 @@ def blend_images(image_a, image_b, blend_percentage):
     blend_mask = ImageOps.invert(blend_mask)  # Invert the mask
 
     # Apply composite blend
-    result = Image.composite(image_a.convert("RGB"), out_image, blend_mask)
+    result = Image.composite(image_a, out_image, blend_mask)
     return result
 
+def apply_image_levels(image, black_level, mid_level, white_level):
+    levels = AdjustLevels(black_level, mid_level, white_level)
+    adjusted_image = levels.adjust(image)
+    return adjusted_image
 
 class AdjustLevels:
     def __init__(self, min_level, mid_level, max_level):
@@ -449,238 +411,175 @@ class AdjustLevels:
         self.max_level = max_level
 
     def adjust(self, im):
-        """Adjust image levels with proper error handling"""
-        # Ensure image is in RGB mode
-        im = im.convert("RGB")
-        
-        # Convert to numpy array
+
         im_arr = np.array(im).astype(np.float32)
-        
-        # Apply levels adjustment
         im_arr[im_arr < self.min_level] = self.min_level
-        im_arr = (im_arr - self.min_level) * (255 / (self.max_level - self.min_level))
+        im_arr = (im_arr - self.min_level) * \
+            (255 / (self.max_level - self.min_level))
         im_arr = np.clip(im_arr, 0, 255)
 
-        # Apply mid-level adjustment (gamma correction)
-        # Fix: Add error handling for potential division by zero or log errors
-        try:
-            gamma_dividend = (self.mid_level - self.min_level)
-            gamma_divisor = (self.max_level - self.min_level)
-            
-            # Avoid division by zero
-            if gamma_divisor == 0:
-                gamma = 1.0
-            elif gamma_dividend <= 0:
-                gamma = 1.0
-            else:
-                gamma = math.log(0.5) / math.log(gamma_dividend / gamma_divisor)
-                
-            # Ensure gamma is reasonable
-            gamma = max(0.1, min(5.0, gamma))
-            
-            im_arr = np.power(im_arr / 255, gamma) * 255
-        except Exception as e:
-            logger.error(f"Error in gamma calculation: {e}")
-            # Fall back to no gamma adjustment
-            pass
+        # mid-level adjustment
+        gamma = math.log(0.5) / math.log((self.mid_level - self.min_level) / (self.max_level - self.min_level))
+        im_arr = np.power(im_arr / 255, gamma) * 255
 
         im_arr = im_arr.astype(np.uint8)
-        im = Image.fromarray(im_arr)
-        return im
-
 
-def apply_image_levels(image, black_level, mid_level, white_level):
-    """Apply levels adjustment to an image"""
-    levels = AdjustLevels(black_level, mid_level, white_level)
-    adjusted_image = levels.adjust(image)
-    return adjusted_image
+        im = Image.fromarray(im_arr)
 
+        return im
 
 def resize_image(image, scaling_factor=1):
-    """Resize image with error handling"""
-    if scaling_factor <= 0:
-        logger.warning(f"Invalid scaling factor: {scaling_factor}, using 1.0 instead")
-        scaling_factor = 1.0
-        
-    try:
-        new_width = int(image.width * scaling_factor)
-        new_height = int(image.height * scaling_factor)
-        
-        # Ensure minimum size
-        new_width = max(1, new_width)
-        new_height = max(1, new_height)
-        
-        resized = image.resize((new_width, new_height), Image.LANCZOS)
-        return resized
-    except Exception as e:
-        logger.error(f"Error resizing image: {e}")
-        return image
-
+    image = image.resize((int(image.width * scaling_factor), 
+        int(image.height * scaling_factor)))
+    return image
 
 def resize_to_square(image, size=1024):
-    """Resize image to a square canvas while maintaining aspect ratio"""
-    try:
-        # Convert to RGBA if needed
-        img = image.convert("RGBA") if isinstance(image, Image.Image) else Image.open(image).convert("RGBA")
 
-        # Resize while maintaining aspect ratio
-        img.thumbnail((size, size), Image.LANCZOS)
+    # Load image if a file path is provided
+    if isinstance(image, str):
+        img = Image.open(image).convert("RGBA")
+    else:
+        img = image.convert("RGBA")  # If already an Image object
+
+    # Resize while maintaining aspect ratio
+    img.thumbnail((size, size), Image.LANCZOS)
 
-        # Create a transparent square canvas
-        square_img = Image.new("RGBA", (size, size), (0, 0, 0, 0))
+    # Create a transparent square canvas
+    square_img = Image.new("RGBA", (size, size), (0, 0, 0, 0))
 
-        # Calculate the position to paste the resized image (centered)
-        x_offset = (size - img.width) // 2
-        y_offset = (size - img.height) // 2
+    # Calculate the position to paste the resized image (centered)
+    x_offset = (size - img.width) // 2
+    y_offset = (size - img.height) // 2
 
-        # Extract the alpha channel as a mask
-        mask = img.split()[3] if img.mode == "RGBA" else None
+    # Extract the alpha channel as a mask
+    mask = img.split()[3] if img.mode == "RGBA" else None
 
-        # Paste the resized image onto the square canvas with the correct transparency mask
-        square_img.paste(img, (x_offset, y_offset), mask)
+    # Paste the resized image onto the square canvas with the correct transparency mask
+    square_img.paste(img, (x_offset, y_offset), mask)
 
-        return square_img
-    except Exception as e:
-        logger.error(f"Error creating square image: {e}")
-        # Return the original image in case of an error
-        return image
+    return square_img
 
 
 def encode_image(image):
-    """Encode image to base64 for API requests"""
     buffer = BytesIO()
-    # Ensure image is in proper format
-    if image.mode not in ["RGB", "RGBA"]:
-        image = image.convert("RGBA")
     image.save(buffer, format="PNG")
     encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
     return f"data:image/png;base64,{encoded_image}"
 
-
 def generate_ai_bg(input_img, prompt):
-    """Generate AI background using external service"""
-    try:
-        # Make sure the prompt is not empty
-        if not prompt or prompt.strip() == "":
-            prompt = "realistic automotive photography, professional lighting"
-            logger.info("Using default prompt for AI background generation")
-        
-        # Use gradio_client for the API call
-        from gradio_client import Client, handle_file
-        
-        try:
-            client = Client("lllyasviel/iclight-v2-vary")
-            result = client.predict(
-                input_fg=handle_file(input_img),
-                bg_source="None",
-                prompt=prompt,
-                image_width=1024,
-                image_height=1024,
-                num_samples=1,
-                seed=12345,
-                steps=25,
-                n_prompt="lowres, bad anatomy, bad hands, cropped, worst quality",
-                cfg=2,
-                gs=5,
-                enable_hr_fix=True,
-                hr_downscale=0.5,
-                lowres_denoise=0.8,
-                highres_denoise=0.99,
-                api_name="/process"
-            )
-            
-            logger.info(f"AI background generation result: {result}")
-            relight_img_path = result[1]
-            
-            # Load the generated image
-            relight_img = Image.open(relight_img_path).convert("RGBA")
-            
-            # Ensure sizes match
-            if relight_img.size != input_img.size:
-                logger.info(f"Resizing generated image from {relight_img.size} to {input_img.size}")
-                relight_img = relight_img.resize(input_img.size, Image.LANCZOS)
-                
-            return relight_img
-            
-        except Exception as e:
-            logger.error(f"Error using gradio_client API: {e}")
-            # Fall back to a simpler method - just use the input image with a color filter
-            relight_img = input_img.copy()
-            color_overlay = Image.new("RGBA", input_img.size, (100, 150, 200, 128))
-            relight_img = Image.alpha_composite(relight_img.convert("RGBA"), color_overlay)
-            logger.info("Using fallback method for image generation")
-            return relight_img
-            
-    except Exception as e:
-        logger.error(f"Error in AI background generation: {e}")
-        # Return the original image in case of an error
-        return input_img.copy()
+    # input_img = resize_image(input_img, 0.01)
+    # hf_input_img = encode_image(input_img)
+
+    # handler = fal_client.submit(
+    #     "fal-ai/iclight-v2",
+    #     arguments={
+    #         "prompt": prompt,
+    #         "image_url": hf_input_img
+    #     },
+    #     webhook_url="https://optional.webhook.url/for/results",
+    # )
+
+    # request_id = handler.request_id
+
+    # status = fal_client.status("fal-ai/iclight-v2", request_id, with_logs=True)
+
+    # result = fal_client.result("fal-ai/iclight-v2", request_id)
+
+    # relight_img_path = result['images'][0]['url']
+
+    # response = requests.get(relight_img_path, stream=True)
 
+    # relight_img = Image.open(BytesIO(response.content)).convert("RGBA")
+
+    from gradio_client import Client, handle_file
+
+    client = Client("lllyasviel/iclight-v2-vary")
+    result = client.predict(
+            input_fg=handle_file(input_img),
+            bg_source="None",
+            prompt=prompt,
+            image_width=1024,
+            image_height=1024,
+            num_samples=1,
+            seed=12345,
+            steps=25,
+            n_prompt="lowres, bad anatomy, bad hands, cropped, worst quality",
+            cfg=2,
+            gs=5,
+            enable_hr_fix=True,
+            hr_downscale=0.5,
+            lowres_denoise=0.8,
+            highres_denoise=0.99,
+            api_name="/process"
+    )
+    print(result)
+
+    relight_img_path = result[1]
+
+    # response = requests.get(relight_img_path, stream=True)
+    relight_img = Image.open(relight_img_path).convert("RGBA")
+
+    # relight_img = Image.open(BytesIO(response.content)).convert("RGBA")
+
+    return relight_img
 
 def blend_details(input_image, relit_image, masked_image, scaling_factor=1):
-    """Blend original and relit images using mask with detailed processing"""
-    try:
-        # Ensure all inputs are proper images
-        if input_image is None or relit_image is None or masked_image is None:
-            logger.error("Missing input for blend_details")
-            return input_image if input_image is not None else Image.new("RGB", (800, 600), (0, 0, 0))
-            
-        # Ensure all images have the same size
-        if input_image.size != relit_image.size:
-            logger.warning(f"Relit image size ({relit_image.size}) doesn't match input image size ({input_image.size})")
-            relit_image = relit_image.resize(input_image.size, Image.LANCZOS)
-            
-        if input_image.size != masked_image.size:
-            logger.warning(f"Mask image size ({masked_image.size}) doesn't match input image size ({input_image.size})")
-            masked_image = masked_image.resize(input_image.size, Image.LANCZOS)
-            
-        # Process masked image
-        masked_image_rgb = split_image_with_alpha(masked_image)
-        masked_image_blurred = gaussian_blur(masked_image_rgb, radius=10)
-        
-        # Fix: Add error handling for mask expansion
-        try:
-            grow_mask = expand_mask(masked_image_blurred, -15, True)
-        except Exception as e:
-            logger.error(f"Error expanding mask: {e}")
-            grow_mask = masked_image_blurred.convert("L")
 
-        # Process input image
-        input_image_rgb = split_image_with_alpha(input_image)
-        input_blurred = gaussian_blur(input_image_rgb, radius=10)
-        input_inverted = invert_image(input_image_rgb)
-        
-        # Add blurred and inverted images
-        input_blend_1 = blend_images(input_inverted, input_blurred, blend_percentage=0.5)
-        input_blend_1_inverted = invert_image(input_blend_1)
-        input_blend_2 = blend_images(input_blurred, input_blend_1_inverted, blend_percentage=1.0)
-
-        # Process relit image
-        relit_image_rgb = split_image_with_alpha(relit_image)
-        relit_blurred = gaussian_blur(relit_image_rgb, radius=10)
-        relit_inverted = invert_image(relit_image_rgb)
-        
-        # Add blurred and inverted relit images
-        relit_blend_1 = blend_images(relit_inverted, relit_blurred, blend_percentage=0.5)
-        relit_blend_1_inverted = invert_image(relit_blend_1)
-        relit_blend_2 = blend_images(relit_blurred, relit_blend_1_inverted, blend_percentage=1.0)
+    # input_image = resize_image(input_image)
 
-        # Blend high frequency components
-        high_freq_comp = image_blend_by_mask(relit_blend_2, input_blend_2, grow_mask, blend_percentage=1.0)
-        
-        # Final compositing
-        comped_image = blend_images(relit_blurred, high_freq_comp, blend_percentage=0.65)
-        
-        # Apply levels adjustment
-        final_image = apply_image_levels(comped_image, black_level=83, mid_level=128, white_level=172)
+    # relit_image = resize_image(relit_image)
+
+    # masked_image = resize_image(masked_image)
+
+    masked_image_rgb = split_image_with_alpha(masked_image)
+    masked_image_blurred = gaussian_blur(masked_image_rgb, radius=10)
+    grow_mask = expand_mask(masked_image_blurred, -15, True)
+
+    # grow_mask.save("output/grow_mask.png")
 
-        return final_image
+    # Split images and get RGB channels
+    input_image_rgb = split_image_with_alpha(input_image)
+    input_blurred = gaussian_blur(input_image_rgb, radius=10)
+    input_inverted = invert_image(input_image_rgb)
+
+    # input_blurred.save("output/input_blurred.png")
+    # input_inverted.save("output/input_inverted.png")
         
-    except Exception as e:
-        logger.error(f"Error in blend_details: {e}")
-        # Return the relit image on error as a fallback
-        return relit_image if relit_image is not None else input_image
+    # Add blurred and inverted images
+    input_blend_1 = blend_images(input_inverted, input_blurred, blend_percentage=0.5)
+    input_blend_1_inverted = invert_image(input_blend_1)
+    input_blend_2 = blend_images(input_blurred, input_blend_1_inverted, blend_percentage=1.0)
+
+    # input_blend_2.save("output/input_blend_2.png")
+
+    # Process relit image
+    relit_image_rgb = split_image_with_alpha(relit_image)
+    relit_blurred = gaussian_blur(relit_image_rgb, radius=10)
+    relit_inverted = invert_image(relit_image_rgb)
+
+    # relit_blurred.save("output/relit_blurred.png")
+    # relit_inverted.save("output/relit_inverted.png")
+    
+    # Add blurred and inverted relit images
+    relit_blend_1 = blend_images(relit_inverted, relit_blurred, blend_percentage=0.5)
+    relit_blend_1_inverted = invert_image(relit_blend_1)
+    relit_blend_2 = blend_images(relit_blurred, relit_blend_1_inverted, blend_percentage=1.0)
+
+    # relit_blend_2.save("output/relit_blend_2.png")
 
+    high_freq_comp = image_blend_by_mask(relit_blend_2, input_blend_2, grow_mask, blend_percentage=1.0)
+
+    # high_freq_comp.save("output/high_freq_comp.png")
+
+    comped_image = blend_images(relit_blurred, high_freq_comp, blend_percentage=0.65)
+
+    # comped_image.save("output/comped_image.png")
+
+    final_image = apply_image_levels(comped_image, black_level=83, mid_level=128, white_level=172)
+
+    # final_image.save("output/final_image.png")
+
+    return final_image
 
 @spaces.GPU
 def generate_image(input_img, prompt):