Update app.py

app.py

@@ -13,7 +13,8 @@ from pathos.threading import ThreadPool as Pool
 #from diffusers.utils import export_to_gif
 #from huggingface_hub import hf_hub_download
 #from safetensors.torch import load_file
-from diffusers import FlaxStableDiffusionPipeline, DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
+from diffusers.utils import load_image
 #import jax
 #import jax.numpy as jnp
 
@@ -22,6 +23,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 #pipe = FlaxStableDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=jnp.bfloat16, token=os.getenv("hf_token")).to(device)
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1", token=os.getenv("hf_token")).to(device)
 
+pipe2 = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+).to(device)
+pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
+
 def translate(text,lang):
     
     if text == None or lang == None:
@@ -81,7 +87,7 @@ def generate_random_string(length):
     return ''.join(random.choice(characters) for _ in range(length))
 
 @spaces.GPU(duration=35)
-def Piper(_do,neg):
+def Piper(_do):
     try:
         retu = pipe(
             _do,
@@ -97,6 +103,18 @@ def Piper(_do,neg):
         print(e)
         return None
 
+@spaces.GPU(duration=75)
+def Piper2(name,neg):
+    try:
+        retu = pipe2(
+            negative_prompt=neg,
+            image=name
+        )
+        return retu
+    except Exception as e:
+        print(e)
+        return None
+
 @spaces.GPU(duration=25)
 def tok(txt):
     toks = pipe.tokenizer(txt)['input_ids']
@@ -110,15 +128,24 @@ def infer(p1,p2):
         _do.append(f'{p1}')
     if p2 != "":
         _dont = f'{p2} where in {p1}'
-        neg = tok(_dont)
+        neg = _dont
     else:
         neg = None
-    output = Piper('A '+" ".join(_do),neg)
+    output = Piper('A '+" ".join(_do))
     if output == None:
-        return output
+        return None
     else:
         output.images[0].save(name)
-    return name
+    if neg == None:
+        return name
+
+    init_image = load_image(name).convert("RGB")
+    output2 = Piper2(init_image,neg)
+    if output2 == None:
+        return None
+    else:
+        output2.images[0].save("_"+name)
+    return "_"+name
 
 css="""
 input, input::placeholder {