Update app.py

app.py

@@ -1,26 +1,70 @@
 import gradio as gr
-from gtts import gTTS
-import speech_recognition as sr
-import tempfile
-import os
+import numpy as np
+import random
+from diffusers import DiffusionPipeline
+import torch
 from huggingface_hub import InferenceClient
+from PIL import Image
+from io import BytesIO
 
-# Initialize the Hugging Face model client
+# Initialize the Hugging Face client for chat
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-def respond(audio_file, history, system_message, max_tokens, temperature, top_p):
-    # Convert audio file to text
-    recognizer = sr.Recognizer()
-    temp_audio_path = audio_file.name
+# Initialize the DiffusionPipeline for image generation
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if torch.cuda.is_available():
+    torch.cuda.max_memory_allocated(device=device)
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to(device)
+else: 
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
+    pipe = pipe.to(device)
 
-    try:
-        with sr.AudioFile(temp_audio_path) as source:
-            audio_data = recognizer.record(source)
-            message = recognizer.recognize_google(audio_data)
-    except Exception as e:
-        return "Error in recognizing audio", None
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    image = pipe(
+        prompt = prompt, 
+        negative_prompt = negative_prompt,
+        guidance_scale = guidance_scale, 
+        num_inference_steps = num_inference_steps, 
+        width = width, 
+        height = height,
+        generator = generator
+    ).images[0] 
+    return image
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    # Check for image generation request
+    if "generate an image" in message.lower():
+        prompt = message.replace("generate an image", "").strip()
+        image = infer(
+            prompt=prompt,
+            negative_prompt="",
+            seed=0,
+            randomize_seed=True,
+            width=512,
+            height=512,
+            guidance_scale=7.5,
+            num_inference_steps=50
+        )
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = buffered.getvalue()
+        return "Here is your generated image:", img_str
 
-    # Prepare messages for the model
     messages = [{"role": "system", "content": system_message}]
     for val in history:
         if val[0]:
@@ -30,46 +74,59 @@ def respond(audio_file, history, system_message, max_tokens, temperature, top_p)
 
     messages.append({"role": "user", "content": message})
 
-    # Get response from the model
     response = ""
-    for msg in client.chat_completion(
+    for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        token = msg.choices[0].delta.content
+        token = message.choices[0].delta.content
         response += token
+        yield response
 
-    # Convert response to speech
-    tts = gTTS(response, lang='en', slow=False)
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp3') as temp_file:
-        tts.save(temp_file.name)
-        temp_file.seek(0)
-        audio_output = temp_file.read()
-
-    os.remove(temp_file.name)
-
-    return response, audio_output
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=respond,
-    inputs=[
-        gr.Audio(type="filepath", label="Record your audio"),
-        gr.State(value=[]),  # No `label` parameter
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
-    outputs=[
-        gr.Textbox(label="Response"),
-        gr.Audio(label="Response Audio", type="file")
-    ],
-    live=True
-)
+# Define Gradio Blocks interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Chat and Image Generation")
+    
+    with gr.Row():
+        with gr.Column():
+            chat_interface = gr.ChatInterface(
+                respond,
+                additional_inputs=[
+                    gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+                    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+                    gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.95,
+                        step=0.05,
+                        label="Top-p (nucleus sampling)",
+                    ),
+                ],
+            )
+    
+    def process_image_request(prompt):
+        image = infer(
+            prompt=prompt,
+            negative_prompt="",
+            seed=0,
+            randomize_seed=True,
+            width=512,
+            height=512,
+            guidance_scale=7.5,
+            num_inference_steps=50
+        )
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        return buffered.getvalue()
+    
+    gr.Examples(
+        examples=["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "An astronaut riding a green horse", "A delicious ceviche cheesecake slice"],
+        inputs=[gr.Textbox(label="Prompt", placeholder="Enter your prompt")],
+        outputs=[gr.Image()]
+    )
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+demo.queue().launch()