feat: Integrate FLUX.1 [schnell] model with Gradio for image generation

app.py

@@ -1,32 +1,39 @@
 import gradio as gr
-import numpy as np
-import random
-import spaces
-import torch
-from diffusers import DiffusionPipeline
+from gradio_client import Client, handle_file
 
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1").to(device)
+flux_1_schnell_space = "https://black-forest-labs-flux-1-schnell.hf.space"
+client = None
+job = None
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
 
-@spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps, 
-            generator = generator,
-            guidance_scale=0.0
-    ).images[0] 
-    return image, seed
+    global job
+    global client
+    if client is None:
+        try:
+            client = Client(flux_1_schnell_space)
+            print(f"Loaded custom model from {flux_1_schnell_space}")
+        except ValueError as e:
+            print(f"Failed to load custom model: {e}")
+            client = None
+            raise gr.Error("Failed to load client for " + flux_1_schnell_space)
+        
+    try:
+        job = client.submit(
+            prompt=prompt,
+            seed=seed,
+            randomize_seed=randomize_seed,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            api_name="/infer"
+        )
+        result = job.result()
+    except ValueError as e:
+        raise gr.Error(e)
+    
+    return result
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",