Update app.py

app.py

@@ -1,104 +1,110 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-# Chatbot Model Configuration
-CHATBOT_CLIENT = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-# Image Generation Model Configuration
-IMAGE_CLIENT = InferenceClient("black-forest-labs/FLUX.1-schnell")
-
-def chatbot_respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message="You are a friendly Chatbot.",
-    max_tokens=512,
-    temperature=0.7,
-    top_p=0.95,
-):
+class AICompanion:
     """
-    Generate chatbot responses using Hugging Face's Zephyr model
+    Combined AI Companion with Chatbot and Image Generation
     """
-    messages = [{"role": "system", "content": system_message}]
+    def __init__(self):
+        # Chatbot Model Configuration
+        self.chatbot_client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+        
+        # Image Generation Model Configuration
+        self.image_client = InferenceClient("black-forest-labs/FLUX.1-schnell")
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    def chatbot_respond(
+        self,
+        message,
+        history: list[tuple[str, str]],
+        system_message="You are a friendly Chatbot.",
+        max_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+    ):
+        """
+        Generate chatbot responses using Hugging Face's Zephyr model
+        """
+        messages = [{"role": "system", "content": system_message}]
 
-    messages.append({"role": "user", "content": message})
+        for val in history:
+            if val[0]:
+                messages.append({"role": "user", "content": val[0]})
+            if val[1]:
+                messages.append({"role": "assistant", "content": val[1]})
 
-    response = ""
-    for message in CHATBOT_CLIENT.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
+        messages.append({"role": "user", "content": message})
 
-def generate_image(prompt):
-    """
-    Generate image using Flux.1-schnell model
-    """
-    if not prompt:
-        return None
-    
-    try:
-        image = IMAGE_CLIENT.text_to_image(prompt)
-        return image
-    except Exception as e:
-        return f"Error generating image: {str(e)}"
+        response = ""
+        for message in self.chatbot_client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            stream=True,
+            temperature=temperature,
+            top_p=top_p,
+        ):
+            token = message.choices[0].delta.content or ""
+            response += token
+            yield response
 
-def create_combined_interface():
-    """
-    Create a combined Gradio interface with chatbot and image generation
-    """
-    with gr.Blocks() as demo:
-        gr.Markdown("# AI Companion: Chat and Image Generation")
+    def generate_image(self, prompt):
+        """
+        Generate image using Flux.1-schnell model
+        """
+        if not prompt:
+            return None
         
-        with gr.Tab("Chatbot"):
-            chatbot_interface = gr.ChatInterface(
-                chatbot_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)",
-                    ),
-                ]
-            )
-        
-        with gr.Tab("Image Generation"):
-            with gr.Row():
-                image_prompt = gr.Textbox(label="Enter Image Prompt")
-                generate_btn = gr.Button("Generate Image")
-            
-            image_output = gr.Image(label="Generated Image")
+        try:
+            return self.image_client.text_to_image(prompt)
+        except Exception as e:
+            return f"Error generating image: {str(e)}"
+
+    def create_interface(self):
+        """
+        Create a combined Gradio interface with chatbot and image generation
+        """
+        with gr.Blocks() as demo:
+            gr.Markdown("# AI Companion: Chat and Image Generation")
             
-            generate_btn.click(
-                fn=generate_image, 
-                inputs=image_prompt, 
-                outputs=image_output
-            )
+            with gr.Tab("Chatbot"):
+                chatbot_interface = gr.ChatInterface(
+                    self.chatbot_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)",
+                        ),
+                    ]
+                )
             
-            image_prompt.submit(
-                fn=generate_image, 
-                inputs=image_prompt, 
-                outputs=image_output
-            )
+            with gr.Tab("Image Generation"):
+                with gr.Row():
+                    image_prompt = gr.Textbox(label="Enter Image Prompt")
+                    generate_btn = gr.Button("Generate Image")
+                
+                image_output = gr.Image(label="Generated Image")
+                
+                generate_btn.click(
+                    fn=self.generate_image, 
+                    inputs=image_prompt, 
+                    outputs=image_output
+                )
+                
+                image_prompt.submit(
+                    fn=self.generate_image, 
+                    inputs=image_prompt, 
+                    outputs=image_output
+                )
 
-    return demo
+        return demo
 
 # Launch the combined interface
 if __name__ == "__main__":
-    combined_demo = create_combined_interface()
+    companion = AICompanion()
+    combined_demo = companion.create_interface()
     combined_demo.launch()