update app.py v2

app.py

@@ -1,13 +1,11 @@
-import requests
-import io
 from PIL import Image
 import gradio as gr
-from together import Together
+import torch
+from diffusers import FluxPipeline
 from groq import Groq  # Import the Groq library
 
-# Initialize the Together and Groq clients with your API keys
-client = Together(api_key="d43d9d8b78242356408329a19a9cdce9290a8c3a254db6293fa6293d207e7e38")
-groq_client = Groq(api_key="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")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
 
 def enhance_prompt(user_prompt):
     """Enhances the given prompt using Groq and returns the refined prompt."""
@@ -46,26 +44,19 @@ def enhance_prompt(user_prompt):
 def generate_image(prompt):
     """Generates an image using the refined prompt."""
     try:
-        response = client.images.generate(
-            prompt=prompt,
-            model="black-forest-labs/FLUX.1-dev-lora",
-            width=1024,
-            height=768,
-            steps=28,
-            n=1,
-            response_format="url",
-            image_loras=[
-                {"path": "https://huggingface.co/XLabs-AI/flux-RealismLora", "scale": 1},
-            ],
-        )
-        image_url = response.data[0].url
-        image_response = requests.get(image_url)
-        image_bytes = image_response.content
-        img = Image.open(io.BytesIO(image_bytes))
+        image = pipe(
+        prompt,
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        num_inference_steps=50,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
+).images[0]
     except Exception as e:
         # Optionally, handle errors (you can also return a default error image)
-        img = None
-    return img
+        image = None
+    return image
 
 # Build the Gradio interface with a two-step process
 with gr.Blocks(css=".gradio-container {background-color: #f9f9f9; padding: 20px;}") as demo: