Update tk_app.py

tk_app.py

@@ -0,0 +1,198 @@
+import tkinter as tk
+from tkinter import ttk, scrolledtext, filedialog
+from PIL import Image, ImageTk
+import asyncio
+import random
+from threading import RLock
+import requests
+import io
+import base64
+from all_models import models
+
+class ImageGeneratorApp:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("Image Generator")
+        self.lock = RLock()
+        
+        self.models = models  # Use the imported models
+        self.num_models = 6  # Number of models to display
+        
+        self.create_widgets()
+        
+    def create_widgets(self):
+        self.notebook = ttk.Notebook(self.root)
+        self.notebook.pack(fill=tk.BOTH, expand=True)
+        
+        self.create_local_tab()
+        self.create_api_tab()
+    
+    def create_local_tab(self):
+        local_frame = ttk.Frame(self.notebook)
+        self.notebook.add(local_frame, text="Local Generation")
+        
+        # Prompt input
+        prompt_label = ttk.Label(local_frame, text="Your prompt:")
+        prompt_label.pack(pady=5)
+        self.local_prompt_input = scrolledtext.ScrolledText(local_frame, height=4)
+        self.local_prompt_input.pack(pady=5, padx=10, fill=tk.X)
+        
+        # Generate button
+        generate_button = ttk.Button(local_frame, text="Generate Images", command=self.generate_images)
+        generate_button.pack(pady=10)
+        
+        # Image display area
+        image_frame = ttk.Frame(local_frame)
+        image_frame.pack(pady=10, padx=10)
+        
+        self.image_labels = []
+        for i in range(self.num_models):
+            label = ttk.Label(image_frame)
+            label.grid(row=i//3, column=i%3, padx=5, pady=5)
+            self.image_labels.append(label)
+        
+        # Model selection
+        model_frame = ttk.LabelFrame(local_frame, text="Model Selection")
+        model_frame.pack(pady=10, padx=10, fill=tk.X)
+        
+        self.model_vars = []
+        for model in self.models[:self.num_models]:
+            var = tk.BooleanVar(value=True)
+            cb = ttk.Checkbutton(model_frame, text=model, variable=var)
+            cb.pack(anchor=tk.W)
+            self.model_vars.append(var)
+    
+    def create_api_tab(self):
+        api_frame = ttk.Frame(self.notebook)
+        self.notebook.add(api_frame, text="API Generation")
+        
+        # Model selection
+        model_label = ttk.Label(api_frame, text="Model:")
+        model_label.pack(pady=5)
+        self.api_model_var = tk.StringVar(value=self.models[0])
+        model_combobox = ttk.Combobox(api_frame, textvariable=self.api_model_var, values=self.models)
+        model_combobox.pack(pady=5, padx=10, fill=tk.X)
+        
+        # Prompt input
+        prompt_label = ttk.Label(api_frame, text="Your prompt:")
+        prompt_label.pack(pady=5)
+        self.api_prompt_input = scrolledtext.ScrolledText(api_frame, height=4)
+        self.api_prompt_input.pack(pady=5, padx=10, fill=tk.X)
+        
+        # Generate button
+        generate_button = ttk.Button(api_frame, text="Generate Image", command=self.generate_api_image)
+        generate_button.pack(pady=10)
+        
+        # Image display
+        self.api_image_label = ttk.Label(api_frame)
+        self.api_image_label.pack(pady=10)
+        
+        # Add to gallery button
+        add_gallery_button = ttk.Button(api_frame, text="Add to Gallery", command=self.add_to_gallery)
+        add_gallery_button.pack(pady=10)
+        
+        # Gallery display
+        self.gallery_frame = ttk.Frame(api_frame)
+        self.gallery_frame.pack(pady=10, padx=10, fill=tk.BOTH, expand=True)
+    
+    async def generate_image(self, model, prompt):
+        # This is a placeholder for actual model inference
+        # In a real application, you would call the actual model here
+        await asyncio.sleep(random.uniform(1, 3))  # Random delay to simulate processing time
+        return Image.new('RGB', (256, 256), color=random.choice(['red', 'green', 'blue']))
+    
+    async def generate_all_images(self, prompt):
+        tasks = []
+        for model, var in zip(self.models[:self.num_models], self.model_vars):
+            if var.get():
+                task = asyncio.create_task(self.generate_image(model, prompt))
+                tasks.append(task)
+        
+        results = await asyncio.gather(*tasks)
+        
+        for label, image in zip(self.image_labels, results):
+            if image:
+                photo = ImageTk.PhotoImage(image)
+                label.configure(image=photo)
+                label.image = photo
+            else:
+                label.configure(image='')
+    
+    def generate_images(self):
+        prompt = self.local_prompt_input.get("1.0", tk.END).strip()
+        if not prompt:
+            return
+        
+        for label in self.image_labels:
+            label.configure(image='')
+        
+        asyncio.run(self.generate_all_images(prompt))
+    
+    def generate_api_image(self):
+        model_str = self.api_model_var.get()
+        prompt = self.api_prompt_input.get("1.0", tk.END).strip()
+        
+        if not prompt:
+            return
+        
+        # Make API call
+        url = "https://k00b404-huggingfacediffusion-custom.hf.space/run/gen_fn_4"
+        payload = {
+            "data": [
+                model_str,
+                prompt
+            ]
+        }
+        response = requests.post(url, json=payload)
+        
+        if response.status_code == 200:
+            result = response.json()
+            image_data = base64.b64decode(result['data'][0].split(',')[1])
+            image = Image.open(io.BytesIO(image_data))
+            photo = ImageTk.PhotoImage(image)
+            self.api_image_label.configure(image=photo)
+            self.api_image_label.image = photo
+            self.current_api_image = image
+        else:
+            print(f"Error: {response.status_code}")
+    
+    def add_to_gallery(self):
+        if not hasattr(self, 'current_api_image'):
+            return
+        
+        # Save the image to a temporary file
+        temp_file = 'temp_image.png'
+        self.current_api_image.save(temp_file)
+        
+        # Make API call to add to gallery
+        url = "https://k00b404-huggingfacediffusion-custom.hf.space/run/add_gallery_4"
+        files = {
+            'data': ('image.png', open(temp_file, 'rb'), 'image/png')
+        }
+        data = {
+            'data': ['', self.api_model_var.get(), '[]']
+        }
+        response = requests.post(url, files=files, data=data)
+        
+        if response.status_code == 200:
+            result = response.json()
+            self.update_gallery(result['data'][0])
+        else:
+            print(f"Error: {response.status_code}")
+    
+    def update_gallery(self, gallery_data):
+        for widget in self.gallery_frame.winfo_children():
+            widget.destroy()
+        
+        for i, item in enumerate(gallery_data):
+            image_data = base64.b64decode(item['image'].split(',')[1])
+            image = Image.open(io.BytesIO(image_data))
+            photo = ImageTk.PhotoImage(image)
+            label = ttk.Label(self.gallery_frame, image=photo)
+            label.image = photo
+            label.grid(row=i//3, column=i%3, padx=5, pady=5)
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = ImageGeneratorApp(root)
+    root.mainloop()