Update image_generator.py

image_generator.py

@@ -1,100 +1,178 @@
-import base64
-import io
-import os
-import time
-import logging
-from typing import List
-from PIL import Image
-from huggingface_hub import InferenceClient
-from config import config
-from models import ImageGenerationRequest, ImageData, ResponseFormat
-
-logger = logging.getLogger(__name__)
-
-
-class ImageGenerator:
-    """Text-to-image generator using Hugging Face InferenceClient"""
-    
-    def __init__(self):
-        self.client = None
-        self._ensure_output_dir()
-        
-    def _ensure_output_dir(self):
-        """Ensure output directory exists"""
-        os.makedirs(config.OUTPUT_DIR, exist_ok=True)
-        
-    def _get_client(self):
-        """Get or create the InferenceClient"""
-        if self.client is None:
-            self.client = InferenceClient(
-                provider="replicate",
-                api_key=config.HF_TOKEN,
-            )
-        return self.client
-    
-    def _image_to_base64(self, image: Image.Image) -> str:
-        """Convert PIL Image to base64 string"""
-        buffer = io.BytesIO()
-        image.save(buffer, format="PNG")
-        img_str = base64.b64encode(buffer.getvalue()).decode()
-        return img_str
-    
-    def _save_image(self, image: Image.Image, filename: str) -> str:
-        """Save image and return URL"""
-        filepath = os.path.join(config.OUTPUT_DIR, filename)
-        image.save(filepath)
-        return f"{config.BASE_URL}/images/{filename}"
-    
-    async def generate_images(self, request: ImageGenerationRequest) -> List[ImageData]:
-        """Generate images based on the request"""
-        client = self._get_client()
-        
-        # Generate images
-        results = []
-        
-        for i in range(request.n):
-            try:
-                logger.info(f"Generating image {i+1}/{request.n} for prompt: {request.prompt[:50]}...")
-                
-                # Generate the image using HuggingFace InferenceClient
-                image = client.text_to_image(
-                    request.prompt,
-                    model=config.DEFAULT_MODEL,
-                )
-                
-                # Create response based on format
-                if request.response_format == ResponseFormat.B64_JSON:
-                    image_data = ImageData(
-                        b64_json=self._image_to_base64(image),
-                        revised_prompt=request.prompt
-                    )
-                else:
-                    # Save image and return URL
-                    timestamp = int(time.time())
-                    filename = f"generated_{timestamp}_{i}.png"
-                    url = self._save_image(image, filename)
-                    image_data = ImageData(
-                        url=url,
-                        revised_prompt=request.prompt
-                    )
-                
-                results.append(image_data)
-                logger.info(f"Successfully generated image {i+1}/{request.n}")
-                
-            except Exception as e:
-                logger.error(f"Failed to generate image {i+1}: {e}")
-                # Continue with other images
-                continue
-        
-        if not results:
-            raise Exception("Failed to generate any images")
-            
-        return results
-    
-    def cleanup(self):
-        """Cleanup resources"""
-        self.client = None
-
-
-# Global instance
-image_generator = ImageGenerator()
+import base64
+import io
+import os
+import sys
+import time
+import logging
+import tempfile
+import subprocess
+from typing import List
+from enum import Enum
+
+# Install required packages
+def install_packages():
+    """Install required packages using pip"""
+    packages = [
+        "pillow",
+        "huggingface_hub",
+        "pydantic"
+    ]
+    
+    for package in packages:
+        try:
+            __import__(package.replace("-", "_"))
+            print(f"{package} already installed")
+        except ImportError:
+            print(f"Installing {package}...")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+# Install packages before importing
+install_packages()
+
+from PIL import Image
+from huggingface_hub import InferenceClient
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+# Define models directly in the file
+class ResponseFormat(str, Enum):
+    URL = "url"
+    B64_JSON = "b64_json"
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str
+    model: str = "black-forest-labs/flux-schnell"
+    n: int = 1
+    size: str = "1024x1024"
+    quality: str = "standard"
+    response_format: ResponseFormat = ResponseFormat.URL
+
+class ImageData(BaseModel):
+    url: str = None
+    b64_json: str = None
+    revised_prompt: str = None
+
+class ImageGenerator:
+    """Text-to-image generator using Hugging Face InferenceClient"""
+    
+    def __init__(self, hf_token: str = None):
+        self.client = None
+        self.hf_token = hf_token or os.getenv("HF_TOKEN")
+        self.output_dir = tempfile.mkdtemp(prefix="image_gen_")
+        self.base_url = "http://localhost:8000"  # Default base URL
+        self.default_model = "black-forest-labs/flux-schnell"
+        self._ensure_output_dir()
+        
+    def _ensure_output_dir(self):
+        """Ensure output directory exists"""
+        os.makedirs(self.output_dir, exist_ok=True)
+        print(f"Using temporary directory: {self.output_dir}")
+        
+    def _get_client(self):
+        """Get or create the InferenceClient"""
+        if self.client is None:
+            if not self.hf_token:
+                raise ValueError("HuggingFace token is required. Set HF_TOKEN environment variable or pass it to constructor.")
+            
+            self.client = InferenceClient(
+                token=self.hf_token,
+            )
+        return self.client
+    
+    def _image_to_base64(self, image: Image.Image) -> str:
+        """Convert PIL Image to base64 string"""
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        img_str = base64.b64encode(buffer.getvalue()).decode()
+        return img_str
+    
+    def _save_image(self, image: Image.Image, filename: str) -> str:
+        """Save image and return URL"""
+        filepath = os.path.join(self.output_dir, filename)
+        image.save(filepath)
+        return f"{self.base_url}/images/{filename}"
+    
+    def set_config(self, hf_token: str = None, base_url: str = None, default_model: str = None):
+        """Set configuration parameters"""
+        if hf_token:
+            self.hf_token = hf_token
+            self.client = None  # Reset client to use new token
+        if base_url:
+            self.base_url = base_url
+        if default_model:
+            self.default_model = default_model
+    
+    async def generate_images(self, request: ImageGenerationRequest) -> List[ImageData]:
+        """Generate images based on the request"""
+        client = self._get_client()
+        
+        # Generate images
+        results = []
+        
+        for i in range(request.n):
+            try:
+                logger.info(f"Generating image {i+1}/{request.n} for prompt: {request.prompt[:50]}...")
+                
+                # Generate the image using HuggingFace InferenceClient
+                image = client.text_to_image(
+                    request.prompt,
+                    model=request.model or self.default_model,
+                )
+                
+                # Create response based on format
+                if request.response_format == ResponseFormat.B64_JSON:
+                    image_data = ImageData(
+                        b64_json=self._image_to_base64(image),
+                        revised_prompt=request.prompt
+                    )
+                else:
+                    # Save image and return URL
+                    timestamp = int(time.time())
+                    filename = f"generated_{timestamp}_{i}.png"
+                    url = self._save_image(image, filename)
+                    image_data = ImageData(
+                        url=url,
+                        revised_prompt=request.prompt
+                    )
+                
+                results.append(image_data)
+                logger.info(f"Successfully generated image {i+1}/{request.n}")
+                
+            except Exception as e:
+                logger.error(f"Failed to generate image {i+1}: {e}")
+                # Continue with other images
+                continue
+        
+        if not results:
+            raise Exception("Failed to generate any images")
+            
+        return results
+    
+    def cleanup(self):
+        """Cleanup resources and temporary directory"""
+        self.client = None
+        # Clean up temporary directory
+        import shutil
+        if os.path.exists(self.output_dir):
+            shutil.rmtree(self.output_dir)
+            print(f"Cleaned up temporary directory: {self.output_dir}")
+
+# Example usage
+if __name__ == "__main__":
+    # Create generator instance
+    generator = ImageGenerator()
+    
+    # Set HuggingFace token (replace with your actual token)
+    generator.set_config(hf_token="your_hf_token_here")
+    
+    # Example request
+    request = ImageGenerationRequest(
+        prompt="A beautiful sunset over mountains",
+        n=1,
+        response_format=ResponseFormat.URL
+    )
+    
+    # Note: This would need to be run in an async context
+    # results = await generator.generate_images(request)
+    print("Image generator setup complete!")