Update flux1_img2img.py

flux1_img2img.py

@@ -1,66 +1,29 @@
-import os
 import torch
 from diffusers import FluxImg2ImgPipeline
+
 from PIL import Image
 import sys
 import spaces
 
-# Set memory optimization flags
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
-
-# Global pipe variable for lazy loading
-pipe = None
-
-def resize_image(image: Image.Image, max_dim: int = 512) -> Image.Image:
-    """Resizes image to fit within max_dim while preserving aspect ratio"""
-    w, h = image.size
-    ratio = min(max_dim / w, max_dim / h)
-    if ratio < 1.0:
-        new_w = int(w * ratio)
-        new_h = int(h * ratio)
-        image = image.resize((new_w, new_h), Image.LANCZOS)
-    return image
-
-def get_pipe(model_id="black-forest-labs/FLUX.1-schnell"):
-    global pipe
-    if pipe is None:
-        pipe = FluxImg2ImgPipeline.from_pretrained(
-            model_id, 
-            torch_dtype=torch.float16
-        ).to("cuda")
-    return pipe
+# I only test with FLUX.1-schnell
 
 @spaces.GPU
-def process_image(image, mask_image, prompt="a person", model_id="black-forest-labs/FLUX.1-schnell", strength=0.75, seed=0, num_inference_steps=4):
+def process_image(image,mask_image,prompt="a person",model_id="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,num_inference_steps=4):
     print("start process image process_image")
-    if image is None:
+    if image == None:
         print("empty input image returned")
         return None
     
-    # Resize image to reduce memory usage
-    image = resize_image(image, max_dim=512)
-    
-    # Get model using lazy loading
-    model = get_pipe(model_id)
+    pipe = FluxImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    pipe.to("cuda")
 
     generators = []
     generator = torch.Generator("cuda").manual_seed(seed)
     generators.append(generator)
-    
-    # Use autocast for better memory efficiency
-    with torch.cuda.amp.autocast(dtype=torch.float16):
-        with torch.no_grad():
-            # more parameter see https://huggingface.co/docs/diffusers/api/pipelines/flux#diffusers.FluxInpaintPipeline
-            print(prompt)
-            output = model(
-                prompt=prompt, 
-                image=image,
-                generator=generator,
-                strength=strength,
-                guidance_scale=0,
-                num_inference_steps=num_inference_steps,
-                max_sequence_length=256
-            )
+    # more parameter see https://huggingface.co/docs/diffusers/api/pipelines/flux#diffusers.FluxInpaintPipeline
+    print(prompt)
+    output = pipe(prompt=prompt, image=image,generator=generator,strength=strength
+                  ,guidance_scale=0,num_inference_steps=num_inference_steps,max_sequence_length=256)
 
     # TODO support mask
     return output.images[0]
@@ -68,6 +31,6 @@ def process_image(image, mask_image, prompt="a person", model_id="black-forest-l
 if __name__ == "__main__":
     #args input-image input-mask output
     image = Image.open(sys.argv[1])
-    mask = Image.open(sys.argv[2])
-    output = process_image(image, mask)
+    mask  = Image.open(sys.argv[2])
+    output = process_image(image,mask)
     output.save(sys.argv[3])