Update app.py

app.py

@@ -1,22 +1,26 @@
 import gradio as gr
-from gradio_client import Client
-
-# Initialize the client with the model endpoint
-client = Client("black-forest-labs/FLUX.1-dev")
-
-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
+import torch
+from transformers import DalleMini, DalleMiniProcessor
+from PIL import Image
+
+# Load model and processor
+model_id = "dalle-mini/dalle-mega"
+model = DalleMini.from_pretrained(model_id)
+processor = DalleMiniProcessor.from_pretrained(model_id)
+
+# Function to generate image
+def generate_image(prompt, num_inference_steps=50):
+    inputs = processor(prompt, return_tensors="pt")
+    
+    # Generate images
+    with torch.no_grad():
+        outputs = model.generate(**inputs, num_inference_steps=num_inference_steps)
+    
+    # Convert to PIL image
+    image = processor.decode(outputs[0], skip_special_tokens=True)
+    image = Image.open(io.BytesIO(image))
+    
+    return image
 
 # Define the Gradio interface
 with gr.Blocks() as demo:
@@ -24,11 +28,6 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
-        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():
@@ -36,10 +35,10 @@ with gr.Blocks() as demo:
 
     result = gr.Image(label="Generated Image")
 
-    # Define the button click action
+    # Connect the function to the button
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs=[prompt, num_inference_steps],
         outputs=result
     )