Update app.py

app.py

@@ -24,6 +24,32 @@ if torch.cuda.is_available():
     pipe = FluxFillPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")
 
 @spaces.GPU()
+def gen(
+        prompt,
+        image,
+        mask_image,
+        width,
+        height,
+        num_inference_steps,
+        seed,
+        guidance_scale, 
+):
+    generator = torch.Generator("cpu").manual_seed(seed)
+    result = pipe(
+        prompt=prompt,
+        image=image,
+        mask_image=mask_image,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        max_sequence_length=512,
+    ).images[0]
+    
+    return result
+
+
 def inpaintGen(
         imgMask,
         inpaint_prompt: str,
@@ -60,17 +86,16 @@ def inpaintGen(
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator("cpu").manual_seed(seed)
 
-    result = pipe(
-        prompt=inpaint_prompt,
-        image=source_img,
-        mask_image=binary_mask,
-        width=new_width,
-        height=new_height,
-        num_inference_steps=num_steps,
-        generator=generator,
-        guidance_scale=guidance,
-        max_sequence_length=512,
-    ).images[0]
+    result = gen(
+        inpaint_prompt,
+        source_img,
+        binary_mask,
+        new_width,
+        new_height,
+        num_steps,
+        seed,
+        guidance,
+    )
     
     return result, seed
 
@@ -126,7 +151,6 @@ def add_border_and_mask(image, zoom_all=1.0, zoom_left=0, zoom_right=0, zoom_up=
     return bordered_image, mask
 
 
-@spaces.GPU()
 def outpaintGen(
     img,
     outpaint_prompt: str,
@@ -136,10 +160,10 @@ def outpaintGen(
     zoom_right: float,
     zoom_up: float,
     zoom_down: float,
-    op_guidance: float,
-    op_num_steps: int,
-    op_seed: int,
-    op_randomize_seed: bool
+    guidance: float,
+    num_steps: int,
+    seed: int,
+    randomize_seed: bool
 ):
     image = Image.open(img)
 
@@ -155,21 +179,19 @@ def outpaintGen(
 
     width, height = new_image.size
     
-    if op_randomize_seed:
-        op_seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator("cpu").manual_seed(op_seed)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-    result = pipe(
-        prompt=outpaint_prompt,
-        image=new_image,
-        mask_image=mask_image,
-        width=width,
-        height=height,
-        num_inference_steps=op_num_steps,
-        generator=generator,
-        guidance_scale=op_guidance,
-        max_sequence_length=512,
-    ).images[0]
+    result = gen(
+        outpaint_prompt,
+        new_image,
+        mask_image,
+        width,
+        height,
+        num_steps,
+        seed,
+        guidance,
+    )
     
     return result, seed