Update app.py

app.py

@@ -1,65 +1,47 @@
 import gradio as gr
-import torch
-from diffusers import DiffusionPipeline  # Note: Change `FluxPipeline` to `DiffusionPipeline` if `FluxPipeline` is not correct
-from PIL import Image
+from gradio_client import Client
 
-# Function to determine the device and handle model loading
-def setup_pipeline():
-    # Check for CUDA availability
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    # Load the diffusion model
-    try:
-        pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-        if device == "cpu":
-            # If using CPU, ensure model is offloaded to avoid GPU-specific features
-            pipeline.enable_model_cpu_offload()
-        else:
-            # Move model to GPU
-            pipeline.to(device)
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        raise e
+# Initialize the client with the model endpoint
+client = Client("black-forest-labs/FLUX.1-dev")
 
-    return pipeline, device
-
-pipeline, device = setup_pipeline()
-
-def generate_image(prompt, guidance_scale=7.5, num_inference_steps=50):
-    # Generate an image based on the prompt
-    with torch.no_grad():
-        try:
-            images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
-        except Exception as e:
-            print(f"Error generating image: {e}")
-            raise e
-    
-    # Assuming pipeline returns a list of images, just take the first one
-    img = images[0]
-    
-    # Convert PIL image to format suitable for Gradio
-    return img
+def generate_image(prompt, seed=0, randomize_seed=True, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28):
+    # Make the API request
+    result = client.predict(
+        prompt=prompt,
+        seed=seed,
+        randomize_seed=randomize_seed,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        api_name="/infer"
+    )
+    return result
 
-# Set up Gradio interface
+# Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Text to Image Generation")
-    
+
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
-        guidance_scale = gr.Slider(minimum=1, maximum=15, step=0.1, value=7.5, label="Guidance Scale")
-        num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=50, label="Number of Inference Steps")
-    
+        seed = gr.Slider(minimum=0, maximum=100000, step=1, value=0, label="Seed")
+        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+        width = gr.Slider(minimum=256, maximum=2048, step=32, value=1024, label="Width")
+        height = gr.Slider(minimum=256, maximum=2048, step=32, value=1024, label="Height")
+        guidance_scale = gr.Slider(minimum=1, maximum=15, step=0.1, value=3.5, label="Guidance Scale")
+        num_inference_steps = gr.Slider(minimum=1, maximum=50, step=1, value=28, label="Number of Inference Steps")
+
     with gr.Row():
         generate_button = gr.Button("Generate Image")
-    
+
     result = gr.Image(label="Generated Image")
-    
-    # Connect the function to the button
+
+    # Define the button click action
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, guidance_scale, num_inference_steps],
+        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs=result
     )
 
-# Launch the app
+# Launch the Gradio app
 demo.launch()