Upload 4 files

.gitignore

@@ -0,0 +1,67 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Model cache (if any)
+models/
+
+# Environment variables
+.env
+.env.local
+.env.*.local
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Model cache
+models/
+generated_images/
+
+# Logs
+*.log
+logs/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Virtual environment
+venv/
+env/
+ENV/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# pytest
+.pytest_cache/
+
+# Coverage
+.coverage
+htmlcov/
+
+# Docker
+.dockerignore

image_generator.py

@@ -0,0 +1,100 @@
+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()

models.py

@@ -0,0 +1,74 @@
+from pydantic import BaseModel, Field
+from typing import List, Optional, Literal
+from enum import Enum
+
+
+class ImageSize(str, Enum):
+    """Supported image sizes (OpenAI compatible)"""
+    SMALL = "256x256"
+    MEDIUM = "512x512" 
+    LARGE = "1024x1024"
+    WIDE = "1792x1024"
+    TALL = "1024x1792"
+
+
+class ImageQuality(str, Enum):
+    """Image quality options"""
+    STANDARD = "standard"
+    HD = "hd"
+
+
+class ImageStyle(str, Enum):
+    """Image style options"""
+    VIVID = "vivid"
+    NATURAL = "natural"
+
+
+class ResponseFormat(str, Enum):
+    """Response format options"""
+    URL = "url"
+    B64_JSON = "b64_json"
+
+
+class ImageGenerationRequest(BaseModel):
+    """OpenAI compatible image generation request"""
+    prompt: str = Field(..., description="A text description of the desired image(s)")
+    model: str = Field(default="dall-e-3", description="The model to use for image generation")
+    n: int = Field(default=1, ge=1, le=10, description="Number of images to generate")
+    quality: ImageQuality = Field(default=ImageQuality.STANDARD, description="Quality of the image")
+    response_format: ResponseFormat = Field(default=ResponseFormat.URL, description="Response format")
+    size: ImageSize = Field(default=ImageSize.LARGE, description="Size of the generated images")
+    style: ImageStyle = Field(default=ImageStyle.VIVID, description="Style of the generated images")
+    user: Optional[str] = Field(default=None, description="A unique identifier representing your end-user")
+
+
+class ImageData(BaseModel):
+    """Individual image data in response"""
+    url: Optional[str] = Field(default=None, description="URL of the generated image")
+    b64_json: Optional[str] = Field(default=None, description="Base64 encoded image data")
+    revised_prompt: Optional[str] = Field(default=None, description="The revised prompt used for generation")
+
+
+class ImageGenerationResponse(BaseModel):
+    """OpenAI compatible image generation response"""
+    created: int = Field(..., description="Unix timestamp of when the image was created")
+    data: List[ImageData] = Field(..., description="List of generated images")
+
+
+class ErrorResponse(BaseModel):
+    """Error response format"""
+    error: dict = Field(..., description="Error details")
+
+
+class ModelInfo(BaseModel):
+    """Model information"""
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str
+
+
+class ModelsResponse(BaseModel):
+    """Models list response"""
+    object: str = "list"
+    data: List[ModelInfo]

requirements.txt

@@ -0,0 +1,7 @@
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+pydantic==2.5.0
+pillow==10.1.0
+huggingface_hub
+requests==2.31.0
+python-multipart==0.0.6