change simple

app.py

@@ -2,11 +2,20 @@ import spaces
 import gradio as gr
 import re
 from PIL import Image
-import flux1_img2img
+
 import os
 import numpy as np
 import shutil
 #shutil.rmtree("/home/user/app/.gradio/cached_examples/23")
+import torch
+from diffusers import FluxImg2ImgPipeline
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = FluxImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
+
+
 
 def sanitize_prompt(prompt):
   # Allow only alphanumeric characters, spaces, and basic punctuation
@@ -14,7 +23,26 @@ def sanitize_prompt(prompt):
   sanitized_prompt = allowed_chars.sub("", prompt)
   return sanitized_prompt
 
-#@spaces.GPU(duration=180)
+#@spaces.GPU
+def process_img2img(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
+
+    generators = []
+    generator = torch.Generator(device).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]
+
+
+@spaces.GPU(duration=180)
 def process_images(image, image2=None,prompt="a girl",inpaint_model="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0):
     print("start process_images")
     try:
@@ -37,11 +65,11 @@ def process_images(image, image2=None,prompt="a girl",inpaint_model="black-fores
             mask = image['layers'][0]
 
 
-        output = flux1_img2img.process_image(image["background"],mask,prompt,inpaint_model,strength,seed)
+        output = process_img2img(image["background"],mask,prompt,inpaint_model,strength,seed)
     except Exception as e:
         print(f"An error occurred: {e}")
         gr.Error(e)
-        
+    print("end process_images")
     return output