Update app.py

app.py

@@ -9,15 +9,14 @@ import gradio as gr
 import numpy as np
 from lxml.html import fromstring
 from pathos.threading import ThreadPool as Pool
-from diffusers import FlaxDiffusionPipeline
 #from diffusers.pipelines.flux import FluxPipeline
 #from diffusers.utils import export_to_gif
 #from huggingface_hub import hf_hub_download
 #from safetensors.torch import load_file
-#from diffusers import FlaxStableDiffusionXLPipeline
+from diffusers import FlaxStableDiffusionPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = FlaxDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to(device)
+pipe = FlaxStableDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to(device)
 
 def translate(text,lang):
     
@@ -95,10 +94,10 @@ def Piper(_do,neg):
         return None
 
 @spaces.GPU(duration=25)
-def negator(_dont):
-    neg = pipe.prepare_inputs(_dont)
-    print(neg)
-    return neg
+def tok(txt):
+    toks = pipe.prepare_inputs(txt)
+    print(toks)
+    return toks
     
 def infer(p1,p2):
     name = generate_random_string(12)+".png"
@@ -107,7 +106,7 @@ def infer(p1,p2):
         _do.append(f'{p1}')
     if p2 != "":
         _dont = f'{p2} where in {p1}'
-        neg = negator(_dont)
+        neg = tok(_dont)
     else:
         neg = None
     output = Piper('A '+" ".join(_do),neg)