Update flux1_img2img.py

flux1_img2img.py

@@ -1,36 +1,36 @@
-import torch
-from diffusers import FluxImg2ImgPipeline
-
-from PIL import Image
-import sys
-import spaces
-
-# 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):
-    print("start process image process_image")
-    if image == None:
-        print("empty input image returned")
-        return None
-    
-    pipe = FluxImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-    pipe.to("cuda")
-
-    generators = []
-    generator = torch.Generator("cuda").manual_seed(seed)
-    generators.append(generator)
-    # 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]
-
-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)
-    output.save(sys.argv[3])
+import torch
+from diffusers import FluxImg2ImgPipeline
+from PIL import Image
+import sys
+import spaces
+
+# 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):
+    print("start process image process_image")
+    if image is None:
+        print("empty input image returned")
+        return None
+
+    # Ensure image is in RGB mode (helps with WebP and other formats)
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    pipe = FluxImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    pipe.to("cuda")
+
+    generator = torch.Generator("cuda").manual_seed(seed)
+    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]
+
+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)
+    output.save(sys.argv[3])