add flux

app.py

@@ -6,13 +6,23 @@ import random
 import time
 from diffusers import AutoPipelineForText2Image
 from diffusers import AutoPipelineForImage2Image
+from diffusers import FluxPipeline
 from diffusers.utils import load_image
 pipeline_text2image = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo").to("cuda")
 pipeline_image2image = AutoPipelineForImage2Image.from_pipe(pipeline_text2image).to("cuda")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cuda")
 
-def text2img(prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe.",guidance_scale=0.0, num_inference_steps=1):
-    
-    image = pipeline_text2image(prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
+def text2img(prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe.",model="flux",guidance_scale=0.0, num_inference_steps=1):
+    if model=="flux":
+        image = pipe(
+            prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            max_sequence_length=256,
+            generator=torch.Generator("cpu").manual_seed(0)
+        ).images[0]
+    else:
+        image = pipeline_text2image(prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
     return image
 
 def img2img(image,prompt="A cinematic shot of a baby racoon wearing an intricate italian priest robe.", guidance_scale=0.0, num_inference_steps=1,strength=0.5):
@@ -26,6 +36,9 @@ gradio_app_text2img = gr.Interface(
     fn=text2img,
     inputs=[
         gr.Text(),
+                gr.Dropdown(
+            ["flux", "stable diffusion"], label="model", info="model to use"
+        ),
         gr.Slider(0.0, 2.0, value=1,step=0.1),
         gr.Slider(2.0, 20.0, value=1,step=1)
     ],