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)
+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.float16)
+    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])