Update app.py

app.py

@@ -9,7 +9,7 @@ import spaces
 import torch
 from PIL import Image
 from diffusers import FluxInpaintPipeline
-
+from diffusers import FluxImg2ImgPipeline
 
 MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024
@@ -33,7 +33,8 @@ def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
 
 pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
-
+pipe2 = FluxImg2ImgPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
 def resize_image_dimensions(
     original_resolution_wh: Tuple[int, int],
@@ -79,21 +80,34 @@ def process(
 
     if not image:
         gr.Info("Please upload an image.")
-        return None, None
-
+        return result, None
+        
+    width, height = resize_image_dimensions(original_resolution_wh=image.size)
+    resized_image = image.resize((width, height), Image.LANCZOS)
+    if randomize_seed_checkbox:
+       seed_slicer = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed_slicer)
+ 
     if not mask:
         gr.Info("Please draw a mask on the image.")
-        return None, None
+        pipe2.load_lora_weights("SIGMitch/KIT")
+        result = pipe2(
+        prompt=input_text,
+        image=resized_image,
+        width=width,
+        height=height,
+        strength=strength_slider,
+        generator=generator,
+        joint_attention_kwargs={"scale": 1.2},
+        num_inference_steps=num_inference_steps_slider
+        ).images[0]
+        print('INFERENCE DONE')
+        return result, None
 
-    width, height = resize_image_dimensions(original_resolution_wh=image.size)
-    resized_image = image.resize((width, height), Image.LANCZOS)
     resized_mask = mask.resize((width, height), Image.LANCZOS)
 
-    if randomize_seed_checkbox:
-        seed_slicer = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed_slicer)
-    pipe.load_lora_weights("SIGMitch/KIT")
-    result = pipe(
+     pipe.load_lora_weights("SIGMitch/KIT")
+     result = pipe(
         prompt=input_text,
         image=resized_image,
         mask_image=resized_mask,