bytedream/hf_api.py

"""
Hugging Face API Client for Byte Dream
Use Byte Dream models directly from Hugging Face Hub
"""

import torch
import requests
import base64
from io import BytesIO
from PIL import Image
from typing import Optional, List, Union
import time


class HuggingFaceAPI:
    """
    Client for Hugging Face Inference API
    Allows using Byte Dream models without downloading them
    """
    
    def __init__(
        self,
        repo_id: str,
        token: Optional[str] = None,
        use_gpu: bool = False,
    ):
        """
        Initialize Hugging Face API client
        
        Args:
            repo_id: Repository ID (e.g., "username/ByteDream")
            token: Hugging Face API token (optional but recommended)
            use_gpu: Request GPU inference (if available)
        """
        self.repo_id = repo_id
        self.token = token
        self.use_gpu = use_gpu
        
        # API endpoints
        self.inference_api_url = f"https://api-inference.huggingface.co/models/{repo_id}"
        self.headers = {}
        
        if token:
            self.headers["Authorization"] = f"Bearer {token}"
        
        print(f"✓ Hugging Face API initialized for: {repo_id}")
    
    def query(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 512,
        height: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        seed: Optional[int] = None,
    ) -> Image.Image:
        """
        Query the model using Inference API
        
        Args:
            prompt: Text prompt
            negative_prompt: Negative prompt
            width: Image width
            height: Image height
            num_inference_steps: Number of denoising steps
            guidance_scale: Guidance scale
            seed: Random seed
            
        Returns:
            Generated PIL Image
        """
        payload = {
            "inputs": prompt,
            "parameters": {
                "negative_prompt": negative_prompt,
                "width": width,
                "height": height,
                "num_inference_steps": num_inference_steps,
                "guidance_scale": guidance_scale,
            }
        }
        
        if seed is not None:
            payload["parameters"]["seed"] = seed
        
        # Make request
        response = requests.post(
            self.inference_api_url,
            headers=self.headers,
            json=payload,
        )
        
        # Handle errors
        if response.status_code == 503:
            # Model is loading
            print("Model is loading on HF servers. Waiting...")
            time.sleep(5)
            return self.query(prompt, negative_prompt, width, height, 
                            num_inference_steps, guidance_scale, seed)
        
        response.raise_for_status()
        
        # Parse image
        image_bytes = response.content
        image = Image.open(BytesIO(image_bytes))
        
        return image
    
    def query_batch(
        self,
        prompts: List[str],
        negative_prompt: str = "",
        width: int = 512,
        height: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        seeds: Optional[List[int]] = None,
    ) -> List[Image.Image]:
        """
        Generate multiple images
        
        Args:
            prompts: List of prompts
            negative_prompt: Negative prompt
            width: Image width
            height: Image height
            num_inference_steps: Number of steps
            guidance_scale: Guidance scale
            seeds: List of seeds
            
        Returns:
            List of PIL Images
        """
        images = []
        
        for i, prompt in enumerate(prompts):
            seed = seeds[i] if seeds and i < len(seeds) else None
            
            print(f"Generating image {i+1}/{len(prompts)}...")
            image = self.query(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                seed=seed,
            )
            
            images.append(image)
        
        return images


class ByteDreamHFClient:
    """
    High-level client for Byte Dream on Hugging Face
    Supports both local inference and API usage
    """
    
    def __init__(
        self,
        repo_id: str,
        token: Optional[str] = None,
        use_api: bool = False,
        device: str = "cpu",
    ):
        """
        Initialize Byte Dream HF client
        
        Args:
            repo_id: Repository ID on Hugging Face
            token: HF API token
            use_api: Use Inference API instead of local inference
            device: Device for local inference
        """
        self.repo_id = repo_id
        self.token = token
        self.use_api = use_api
        self.device = device
        
        if use_api:
            self.api_client = HuggingFaceAPI(repo_id, token)
            print("✓ Using Hugging Face Inference API")
        else:
            # Load model locally
            from bytedream.generator import ByteDreamGenerator
            self.generator = ByteDreamGenerator(
                hf_repo_id=repo_id,
                config_path="config.yaml",
                device=device,
            )
            print("✓ Model loaded locally from Hugging Face")
    
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 512,
        height: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        seed: Optional[int] = None,
    ) -> Image.Image:
        """
        Generate image from prompt
        
        Args:
            prompt: Text description
            negative_prompt: Things to avoid
            width: Image width
            height: Image height
            num_inference_steps: Number of steps
            guidance_scale: Guidance scale
            seed: Random seed
            
        Returns:
            Generated PIL Image
        """
        if self.use_api:
            return self.api_client.query(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                seed=seed,
            )
        else:
            return self.generator.generate(
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt else None,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                seed=seed,
            )
    
    def generate_batch(
        self,
        prompts: List[str],
        negative_prompt: str = "",
        width: int = 512,
        height: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        seeds: Optional[List[int]] = None,
    ) -> List[Image.Image]:
        """
        Generate multiple images
        
        Args:
            prompts: List of text descriptions
            negative_prompt: Things to avoid
            width: Image width
            height: Image height
            num_inference_steps: Number of steps
            guidance_scale: Guidance scale
            seeds: List of random seeds
            
        Returns:
            List of PIL Images
        """
        if self.use_api:
            return self.api_client.query_batch(
                prompts=prompts,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                seeds=seeds,
            )
        else:
            return self.generator.generate_batch(
                prompts=prompts,
                negative_prompt=negative_prompt if negative_prompt else None,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                seeds=seeds,
            )


# Example usage
if __name__ == "__main__":
    # Example 1: Use Inference API
    print("=" * 60)
    print("Example 1: Using Hugging Face Inference API")
    print("=" * 60)
    
    # You need a token for private models or higher rate limits
    # token = "hf_xxxxxxxxxxxxx"
    
    try:
        client = ByteDreamHFClient(
            repo_id="Enzo8930302/ByteDream",  # Replace with your repo
            # token=token,  # Optional but recommended
            use_api=True,  # Set True to use API
        )
        
        image = client.generate(
            prompt="A beautiful sunset over mountains, digital art",
            negative_prompt="ugly, blurry, low quality",
            width=512,
            height=512,
            num_inference_steps=50,
            guidance_scale=7.5,
            seed=42,
        )
        
        image.save("output_api.png")
        print("✓ Image saved to output_api.png")
        
    except Exception as e:
        print(f"Error: {e}")
        print("Make sure the model exists on Hugging Face")
    
    # Example 2: Download and run locally
    print("\n" + "=" * 60)
    print("Example 2: Download and run locally on CPU")
    print("=" * 60)
    
    try:
        client_local = ByteDreamHFClient(
            repo_id="Enzo8930302/ByteDream",
            use_api=False,  # Download and run locally
            device="cpu",
        )
        
        image_local = client_local.generate(
            prompt="A futuristic city at night, cyberpunk style",
            width=512,
            height=512,
            num_inference_steps=30,
        )
        
        image_local.save("output_local.png")
        print("✓ Image saved to output_local.png")
        
    except Exception as e:
        print(f"Error: {e}")