Update app.py

app.py

@@ -2,12 +2,13 @@ import gradio as gr
 import random
 import torch
 import spaces
-from diffusers import DiffusionPipeline
 from langchain_groq import ChatGroq
 from langchain.schema import HumanMessage, SystemMessage
 import os
 from PIL import Image
 import numpy as np
+import base64
+from io import BytesIO
 
 # Set up API keys
 GROQ_API_KEY = os.environ.get('GROQ_API_KEY')
@@ -15,12 +16,6 @@ GROQ_API_KEY = os.environ.get('GROQ_API_KEY')
 # Set up LLM
 llm = ChatGroq(temperature=0, model_name='llama-3.1-8b-instant', groq_api_key=GROQ_API_KEY)
 
-# Set up DiffusionPipeline on CPU
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 512
 
@@ -36,7 +31,6 @@ def generate_detailed_prompt(user_input):
     Given a simple description, create an elaborate and detailed prompt that can be used to generate high-quality images.
     Your response should be concise and no longer than 3 sentences.
     Use the following examples as a guide for the level of detail and creativity expected:
-
     """ + "\n\n".join([f"Input: {input}\nOutput: {output}" for input, output in few_shot_examples]))
 
     human_message = HumanMessage(content=f"Generate a detailed image prompt based on this input, using no more than 3 sentences: {user_input}")
@@ -44,19 +38,27 @@ def generate_detailed_prompt(user_input):
     response = llm([system_message, human_message])
     return response.content
 
+# Initialize the schnell client
+from huggingface_hub import InferenceClient
+client = InferenceClient("black-forest-labs/FLUX.1-schnell")
+
 @spaces.GPU()
 def generate_image(prompt, width=512, height=512, num_inference_steps=4):
     seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    image = pipe(
-        prompt=prompt,
+    
+    # Use the schnell client to generate the image
+    result = client.text_to_image(
+        prompt,
+        negative_prompt="",
         width=width,
         height=height,
         num_inference_steps=num_inference_steps,
-        generator=generator,
-        guidance_scale=0.0
-    ).images[0]
-
+        guidance_scale=0.0,
+        seed=seed
+    )
+    
+    # Convert the image to a PIL Image object
+    image = Image.open(BytesIO(result))
     return image
 
 # Gradio UI setup
@@ -117,4 +119,4 @@ with gr.Blocks(css=css, theme='gradio/soft') as demo:
         outputs=[result]
     )
 
-demo.launch(share=True)
+demo.launch()