flux_avatar_manager.py

# filename: flux_avatar_manager.py

import os
from pathlib import Path
from PIL import Image
from huggingface_hub import HfApi
from flux_client_wrapper import FluxImageGenerator


class AvatarManager:
    def __init__(
        self,
        hf_space_id="K00B404/chatQwenne",
        hf_folder="chatbot_images",
        initial_prompt="a photo realistic fullbody ,head to toe, adult cyborg Chatbuddy avatar ",
        client=None
    ):
        self.hf_space_id = hf_space_id
        self.hf_folder = Path(hf_folder)
        self.hf_folder.mkdir(parents=True, exist_ok=True)
        self.api = HfApi()
        self.client = client or FluxImageGenerator()
        self.initial_prompt = initial_prompt
        self.avatar_image_path = self.hf_folder / "avatar.png"
        self.has_uploaded = False

    def generate_initial_avatar(self):
        if self.avatar_image_path.exists():
            print("Avatar already exists locally.")
            return Image.open(self.avatar_image_path)

        result = self.client.generate_image(prompt=self.initial_prompt)
        # Assuming `result` is a path or PIL.Image object
        if isinstance(result, str) and result.endswith((".png", ".jpg", ".jpeg")):
            image = Image.open(result)
        elif isinstance(result, Image.Image):
            image = result
        else:
            raise ValueError("Unrecognized image format from result.")

        image.save(self.avatar_image_path)
        print(f"Saved avatar to {self.avatar_image_path}")
        return image

    def upload_avatar_to_hf(self):
        if not self.avatar_image_path.exists():
            raise FileNotFoundError("Avatar image not found, generate it first.")

        if self.has_uploaded:
            print("Avatar already uploaded.")
            return

        self.api.upload_folder(
            folder_path=str(self.hf_folder),
            repo_id=self.hf_space_id,
            repo_type="space"
        )
        self.has_uploaded = True
        print(f"Uploaded avatar to {self.hf_space_id}/{self.hf_folder.name}")

    def get_avatar_path(self):
        return self.avatar_image_path if self.avatar_image_path.exists() else None

    def refresh_avatar(self, new_prompt=None):
        prompt = new_prompt or self.initial_prompt
        print(f"Refreshing avatar with prompt: {prompt}")
        self.initial_prompt = prompt
        if self.avatar_image_path.exists():
            os.remove(self.avatar_image_path)
        self.has_uploaded = False
        return self.generate_initial_avatar()