Update app.py

app.py

@@ -1,18 +1,24 @@
 import gradio as gr
 from PIL import Image
 import torch
+from diffusers.utils import load_image
+from diffusers import FluxControlNetModel
+from diffusers.pipelines import FluxControlNetPipeline
+
 #import torch.nn.functional as F
 #import torchvision
 #import torchvision.transforms as T
+#import cv2
+
 from diffusers import StableDiffusionInpaintPipeline
 import numpy as np
-#import cv2
 import os
 import shutil
 from gradio_client import Client, handle_file
 
 # Load the model once globally to avoid repeated loading
 def load_inpainting_model():
+    # Load pipeline
     model_path = "uberRealisticPornMerge_urpmv13Inpainting.safetensors"
     #model_path = "uberRealisticPornMerge_v23Inpainting.safetensors"
     #model_path = "pornmasterFantasy_v4-inpainting.safetensors"
@@ -25,16 +31,48 @@ def load_inpainting_model():
     ).to(device)
     return pipe
 
+# Load the model once globally to avoid repeated loading
+def load_upscaling_model():
+# Load pipeline
+    device = "cpu"  # Explicitly use CPU
+    controlnet = FluxControlNetModel.from_pretrained(
+        "jasperai/Flux.1-dev-Controlnet-Upscaler",
+        torch_dtype=torch.float32
+    )
+    pipe = FluxControlNetPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        controlnet=controlnet,
+        torch_dtype=torch.float32
+    ).to(device)
+    return pipe
+
 # Preload the model once
 inpaint_pipeline = load_inpainting_model()
+# Preload the model once
+upscale_pipeline = load_upscaling_model()
 
 # Function to resize image (simpler interpolation method for speed)
 def resize_to_match(input_image, output_image):
+
+    w, h = output_image.size
+    control_image = output_image.resize((w * 4, h * 4))
+   
+    scaled_image = pipe(
+        prompt="", 
+        control_image=control_image,
+        controlnet_conditioning_scale=0.6,
+        num_inference_steps=28, 
+        guidance_scale=3.5,
+        height=control_image.size[1],
+        width=control_image.size[0]
+    ).images[0]    
+    
+    return scaled_image
+    
     #torch_img = pil_to_torch(input_image)
     #torch_img_scaled = F.interpolate(torch_img.unsqueeze(0),mode='trilinear').squeeze(0)
     #output_image = torchvision.transforms.functional.to_pil_image(torch_img_scaled, mode=None)
-    #return output_image
-    return output_image.resize(input_image.size, Image.BICUBIC)  # Use BILINEAR for faster resizing
+    #return output_image.resize(input_image.size, Image.BICUBIC)  # Use BILINEAR for faster resizing
 
 # Function to generate the mask using Florence SAM Masking API (Replicate)
 def generate_mask(image_path, text_prompt="clothing"):
@@ -58,7 +96,7 @@ def inpaint_image(input_image, mask_image):
     prompt = "undress, naked"
     result = inpaint_pipeline(prompt=prompt, image=input_image, mask_image=mask_image)
     inpainted_image = result.images[0]
-    #inpainted_image = resize_to_match(input_image, inpainted_image)
+    inpainted_image = resize_to_match(input_image, inpainted_image)
     return inpainted_image
 
 # Function to process input image and mask