diff --git a/config.py b/config.py index 1cd2228..22ecdde 100644 --- a/config.py +++ b/config.py @@ -296,6 +296,33 @@ TRUSTED_DEFAULT_URLS = list(map(map_legacy_config, TRUSTED_DEFAULT_URLS)) BLUESKY_IDENTIFIER = os.environ.get("BLUESKY_IDENTIFIER", None) BLUESKY_APP_PASSWORD = os.environ.get("BLUESKY_APP_PASSWORD", None) +# Celery Configuration +# Try Redis first, fallback to memory broker for development +try: + import redis + + redis.Redis(host="localhost", port=6379, db=0).ping() + CELERY_BROKER_URL = "redis://localhost:6379/0" +except Exception: # pylint: disable=broad-except + # Fallback to memory broker for development + CELERY_BROKER_URL = "memory://" + print("Warning: Redis not available, using memory broker for development") + +CELERY_RESULT_BACKEND = "django-db" +CELERY_ACCEPT_CONTENT = ["json"] +CELERY_TASK_SERIALIZER = "json" +CELERY_RESULT_SERIALIZER = "json" +CELERY_TIMEZONE = "UTC" +CELERY_TASK_TRACK_STARTED = True +CELERY_TASK_TIME_LIMIT = 300 # 5 minutes +CELERY_TASK_SOFT_TIME_LIMIT = 240 # 4 minutes +CELERY_WORKER_PREFETCH_MULTIPLIER = 1 +CELERY_TASK_ACKS_LATE = True +CELERY_RESULT_EXPIRES = 3600 # 1 hour +CELERY_WORKER_CONCURRENCY = ( + 1 # Max 1 parallel avatar generation task for local development +) + # This MUST BE THE LAST! if os.path.isfile(os.path.join(BASE_DIR, "config_local.py")): from config_local import * # noqa # flake8: noqa # NOQA # pragma: no cover diff --git a/ivatar/__init__.py b/ivatar/__init__.py index 0649992..13dc53f 100644 --- a/ivatar/__init__.py +++ b/ivatar/__init__.py @@ -3,4 +3,10 @@ Module init """ +# This will make sure the app is always imported when +# Django starts so that shared_task will use this app. +from .celery import app as celery_app + +__all__ = ("celery_app",) + app_label = __name__ # pylint: disable=invalid-name diff --git a/ivatar/ai_service.py b/ivatar/ai_service.py new file mode 100644 index 0000000..bdb0acc --- /dev/null +++ b/ivatar/ai_service.py @@ -0,0 +1,461 @@ +# -*- coding: utf-8 -*- +""" +AI service module for text-to-image avatar generation +Supports Stable Diffusion (local and API) for professional avatar generation +""" + +import logging +import requests +import base64 +from io import BytesIO +from PIL import Image +from django.conf import settings + +logger = logging.getLogger(__name__) + + +class AIServiceError(Exception): + """Custom exception for AI service errors""" + + pass + + +class StableDiffusionService: + """ + Service for generating images using Stable Diffusion (local or API) + """ + + # Model-specific token limits + TOKEN_LIMITS = { + "stable_diffusion": 77, # CLIP tokenizer limit + "stable_diffusion_v2": 77, + "stable_diffusion_xl": 77, + } + + def __init__(self): + self.api_url = getattr(settings, "STABLE_DIFFUSION_API_URL", None) + self.api_key = getattr(settings, "STABLE_DIFFUSION_API_KEY", None) + self.timeout = getattr(settings, "STABLE_DIFFUSION_TIMEOUT", 60) + self._pipe = None # Cache for local model + self._tokenizer = None # Cache for tokenizer + + def generate_image( + self, prompt, size=(512, 512), quality="medium", allow_nsfw=False + ): + """ + Generate an image from text prompt using Stable Diffusion + + Args: + prompt (str): Text description of the desired image + size (tuple): Image dimensions (width, height) + quality (str): Generation quality ('low', 'medium', 'high') + allow_nsfw (bool): Whether to allow potentially NSFW content + + Returns: + PIL.Image: Generated image + + Raises: + AIServiceError: If generation fails or prompt is too long + """ + # Validate prompt length first + validation = self.validate_prompt(prompt) + if not validation["valid"]: + raise AIServiceError(validation["warning"]) + + try: + if self.api_url and self.api_key: + return self._generate_via_api(prompt, size, quality, allow_nsfw) + else: + return self._generate_locally(prompt, size, quality, allow_nsfw) + except Exception as e: + logger.error(f"Failed to generate image: {e}") + raise AIServiceError(f"Image generation failed: {str(e)}") + + def validate_prompt(self, prompt, model="stable_diffusion"): + """ + Validate prompt length against model token limits + + Args: + prompt (str): Text prompt to validate + model (str): Model name to check limits for + + Returns: + dict: Validation result with 'valid', 'token_count', 'limit', 'warning' + """ + try: + token_count = self._count_tokens(prompt) + limit = self.TOKEN_LIMITS.get(model, 77) + + is_valid = token_count <= limit + warning = None + + if not is_valid: + warning = f"Prompt too long: {token_count} tokens (limit: {limit}). Please shorten your prompt." + + return { + "valid": is_valid, + "token_count": token_count, + "limit": limit, + "warning": warning, + } + + except Exception as e: + logger.warning(f"Token counting failed: {e}") + return { + "valid": True, # Allow generation if counting fails + "token_count": 0, + "limit": 77, + "warning": None, + } + + def _count_tokens(self, prompt): + """ + Count tokens in a prompt using CLIP tokenizer + """ + try: + if self._tokenizer is None: + from transformers import CLIPTokenizer + + self._tokenizer = CLIPTokenizer.from_pretrained( + "openai/clip-vit-base-patch32" + ) + + tokens = self._tokenizer(prompt, return_tensors="pt", truncation=False)[ + "input_ids" + ] + return tokens.shape[1] + + except ImportError: + # Fallback: more accurate estimation + # CLIP tokenizer typically produces ~1.3 tokens per word for English + words = len(prompt.split()) + return int(words * 1.3) + except Exception as e: + logger.warning(f"Token counting error: {e}") + # Fallback: more accurate estimation + words = len(prompt.split()) + return int(words * 1.3) + + def _is_black_image(self, image): + """ + Check if an image is completely black (common NSFW response from APIs) + + Args: + image (PIL.Image): Image to check + + Returns: + bool: True if image is completely black + """ + # Convert to RGB if necessary + if image.mode != "RGB": + image = image.convert("RGB") + + # Get image data + pixels = list(image.getdata()) + + # Check if all pixels are black (0, 0, 0) + black_pixels = sum(1 for r, g, b in pixels if r == 0 and g == 0 and b == 0) + total_pixels = len(pixels) + + # Consider it a black image if more than 95% of pixels are black + return (black_pixels / total_pixels) > 0.95 + + def _generate_via_api(self, prompt, size, quality, allow_nsfw=False): + """ + Generate image via Stable Diffusion API (Replicate, Hugging Face, etc.) + """ + # Enhanced prompt for avatar generation + enhanced_prompt = f"""professional avatar portrait, {prompt}, high quality, detailed, clean background, centered composition, profile picture style, photorealistic""" + + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + payload = { + "prompt": enhanced_prompt, + "width": size[0], + "height": size[1], + "num_inference_steps": 25 + if quality == "high" + else (20 if quality == "medium" else 15), + "guidance_scale": 7.5, # Balanced for quality and speed + "negative_prompt": "blurry, low quality, distorted, ugly, deformed, bad anatomy", + } + + # Add NSFW safety setting if supported by the API + if allow_nsfw: + payload["safety_tolerance"] = 2 # Some APIs support this + payload["nsfw"] = True # Some APIs support this + + response = requests.post( + self.api_url, json=payload, headers=headers, timeout=self.timeout + ) + + if response.status_code != 200: + error_msg = f"Stable Diffusion API request failed: {response.status_code}" + try: + error_detail = response.json() + error_msg += f" - {error_detail}" + + # Check for NSFW content detection + if isinstance(error_detail, dict): + error_text = str(error_detail).lower() + if ( + "nsfw" in error_text + or "inappropriate" in error_text + or "black image" in error_text + ): + if allow_nsfw: + # If user allowed NSFW but still got blocked, provide a different message + raise AIServiceError( + "Content warning: The AI service still detected inappropriate content even with relaxed settings. Please try a different prompt or contact support if you believe this is an error." + ) + else: + raise AIServiceError( + "Content warning: The AI service detected potentially inappropriate content in your prompt. Please modify your description to be more appropriate for all ages and try again." + ) + elif isinstance(error_detail, str): + if ( + "nsfw" in error_detail.lower() + or "inappropriate" in error_detail.lower() + or "black image" in error_detail.lower() + ): + if allow_nsfw: + raise AIServiceError( + "Content warning: The AI service still detected inappropriate content even with relaxed settings. Please try a different prompt or contact support if you believe this is an error." + ) + else: + raise AIServiceError( + "Content warning: The AI service detected potentially inappropriate content in your prompt. Please modify your description to be more appropriate for all ages and try again." + ) + + except AIServiceError: + # Re-raise our custom NSFW error + raise + except Exception: # pylint: disable=broad-except + error_msg += f" - {response.text}" + + # Also check response text for NSFW warnings + if ( + "nsfw" in response.text.lower() + or "inappropriate" in response.text.lower() + or "black image" in response.text.lower() + ): + if allow_nsfw: + raise AIServiceError( + "Content warning: The AI service still detected inappropriate content even with relaxed settings. Please try a different prompt or contact support if you believe this is an error." + ) + else: + raise AIServiceError( + "Content warning: The AI service detected potentially inappropriate content in your prompt. Please modify your description to be more appropriate for all ages and try again." + ) + + raise AIServiceError(error_msg) + + result = response.json() + + if "image" in result: + # Decode base64 image + image_data = base64.b64decode(result["image"]) + image = Image.open(BytesIO(image_data)) + + # Check if the image is completely black (common NSFW response) + if not allow_nsfw and self._is_black_image(image): + raise AIServiceError( + "Content warning: The AI service detected potentially inappropriate content in your prompt and returned a black image. Please modify your description to be more appropriate for all ages and try again." + ) + + return image + else: + raise AIServiceError("No image data in API response") + + def _generate_locally(self, prompt, size, quality, allow_nsfw=False): + """ + Generate image using local Stable Diffusion installation + This requires diffusers library and a local model + """ + try: + from diffusers import StableDiffusionPipeline + import torch + + # Enhanced prompt for avatar generation + enhanced_prompt = f"""professional avatar portrait, {prompt}, high quality, detailed, clean background, centered composition, profile picture style, photorealistic""" + + # Use cached model if available, otherwise load it + if self._pipe is None: + logger.info("Loading Stable Diffusion model (first time or cache miss)") + self._pipe = StableDiffusionPipeline.from_pretrained( + "runwayml/stable-diffusion-v1-5", + torch_dtype=torch.float16 + if torch.cuda.is_available() + else torch.float32, + ) + + if torch.cuda.is_available(): + self._pipe = self._pipe.to("cuda") + else: + logger.info("Using cached Stable Diffusion model") + + pipe = self._pipe + + # Disable safety checker if NSFW override is enabled + if allow_nsfw: + pipe.safety_checker = None + pipe.requires_safety_checker = False + + # Generate image with optimized settings for speed + image = pipe( + enhanced_prompt, + height=size[1], + width=size[0], + num_inference_steps=25 + if quality == "high" + else (20 if quality == "medium" else 15), + guidance_scale=7.5, # Balanced for quality and speed + negative_prompt="blurry, low quality, distorted, ugly, deformed, bad anatomy", + ).images[0] + + return image + + except ImportError: + logger.warning( + "diffusers library not installed, falling back to placeholder" + ) + return self._generate_placeholder(prompt, size) + except Exception as e: + logger.error(f"Local Stable Diffusion generation failed: {e}") + return self._generate_placeholder(prompt, size) + + def _generate_placeholder(self, prompt, size): + """ + Generate a placeholder image when Stable Diffusion is not available + """ + logger.info("Generating placeholder image") + + # Create a more sophisticated placeholder + img = Image.new("RGBA", size, color=(240, 248, 255, 255)) + + from PIL import ImageDraw, ImageFont + + draw = ImageDraw.Draw(img) + + try: + font = ImageFont.load_default() + except Exception: # pylint: disable=broad-except + font = None + + # Add title + title = "AI Avatar (Stable Diffusion)" + draw.text((10, 10), title, fill="darkblue", font=font) + + # Add prompt + prompt_text = f"Prompt: {prompt[:50]}..." + draw.text((10, 40), prompt_text, fill="black", font=font) + + # Add note + note = "Install Stable Diffusion for real generation" + draw.text((10, 70), note, fill="darkgreen", font=font) + + # Create a more sophisticated avatar placeholder + center_x, center_y = size[0] // 2, size[1] // 2 + 20 + radius = min(size) // 4 + + # Face circle + draw.ellipse( + [ + center_x - radius, + center_y - radius, + center_x + radius, + center_y + radius, + ], + outline="purple", + width=3, + fill=(255, 240, 245), + ) + + # Eyes + eye_radius = radius // 4 + draw.ellipse( + [ + center_x - radius // 2 - eye_radius, + center_y - radius // 2 - eye_radius, + center_x - radius // 2 + eye_radius, + center_y - radius // 2 + eye_radius, + ], + fill="blue", + ) + draw.ellipse( + [ + center_x + radius // 2 - eye_radius, + center_y - radius // 2 - eye_radius, + center_x + radius // 2 + eye_radius, + center_y - radius // 2 + eye_radius, + ], + fill="blue", + ) + + # Smile + smile_y = center_y + radius // 3 + draw.arc( + [ + center_x - radius // 2, + smile_y - radius // 4, + center_x + radius // 2, + smile_y + radius // 4, + ], + 0, + 180, + fill="red", + width=3, + ) + + return img + + +def validate_avatar_prompt(prompt, model="stable_diffusion"): + """ + Convenience function to validate avatar prompts + + Args: + prompt (str): Text description of the avatar + model (str): AI model to use + + Returns: + dict: Validation result with 'valid', 'token_count', 'limit', 'warning' + """ + if model == "stable_diffusion": + service = StableDiffusionService() + return service.validate_prompt(prompt, model) + else: + # For other models, assume they're valid + return {"valid": True, "token_count": 0, "limit": 0, "warning": None} + + +def generate_avatar_image( + prompt, + model="stable_diffusion", + size=(512, 512), + quality="medium", + allow_nsfw=False, +): + """ + Convenience function to generate avatar images + + Args: + prompt (str): Text description of the avatar + model (str): AI model to use (currently only 'stable_diffusion') + size (tuple): Image dimensions + quality (str): Generation quality ('low', 'medium', 'high') + allow_nsfw (bool): Whether to allow potentially NSFW content + + Returns: + PIL.Image: Generated avatar image + """ + if model == "stable_diffusion": + service = StableDiffusionService() + return service.generate_image(prompt, size, quality, allow_nsfw) + else: + raise AIServiceError( + f"Unsupported model: {model}. Only 'stable_diffusion' is currently supported." + ) diff --git a/ivatar/celery.py b/ivatar/celery.py new file mode 100644 index 0000000..0a6ea92 --- /dev/null +++ b/ivatar/celery.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- +""" +Celery configuration for ivatar +""" + +import os +from celery import Celery +from django.conf import settings + +# Set the default Django settings module for the 'celery' program. +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ivatar.settings") + +app = Celery("ivatar") + +# Using a string here means the worker doesn't have to serialize +# the configuration object to child processes. +app.config_from_object("django.conf:settings", namespace="CELERY") + +# Load task modules from all registered Django apps. +app.autodiscover_tasks() + +# Celery configuration +app.conf.update( + # Task routing - use default queue for simplicity + task_default_queue="default", + task_routes={ + "ivatar.tasks.generate_avatar_task": {"queue": "default"}, + "ivatar.tasks.update_queue_positions": {"queue": "default"}, + "ivatar.tasks.cleanup_old_tasks": {"queue": "default"}, + }, + # Worker configuration + worker_prefetch_multiplier=1, + task_acks_late=True, + # Result backend + result_backend="django-db", + result_expires=3600, # 1 hour + # Task time limits + task_time_limit=300, # 5 minutes + task_soft_time_limit=240, # 4 minutes + # Task serialization + task_serializer="json", + accept_content=["json"], + result_serializer="json", + # Timezone + timezone="UTC", + enable_utc=True, +) + +# Set worker concurrency from Django settings +if hasattr(settings, "CELERY_WORKER_CONCURRENCY"): + app.conf.worker_concurrency = settings.CELERY_WORKER_CONCURRENCY + + +@app.task(bind=True) +def debug_task(self): + print(f"Request: {self.request!r}") diff --git a/ivatar/ivataraccount/forms.py b/ivatar/ivataraccount/forms.py index a2f4dcd..f06e9be 100644 --- a/ivatar/ivataraccount/forms.py +++ b/ivatar/ivataraccount/forms.py @@ -217,6 +217,89 @@ class UploadLibravatarExportForm(forms.Form): ) +class GenerateAvatarForm(forms.Form): + """ + Form for generating avatars using AI text-to-image + """ + + MODEL_CHOICES = [ + ("stable_diffusion", "Stable Diffusion"), + # Future models can be added here + ] + + prompt = forms.CharField( + label=_("Avatar Description"), + max_length=500, + widget=forms.Textarea( + attrs={ + "rows": 3, + "placeholder": _( + 'Describe the avatar you want to create, e.g., "A friendly robot with blue eyes"' + ), + "id": "id_prompt", + "data-token-limit": "77", + "data-model": "stable_diffusion", + } + ), + help_text=_( + "Describe the avatar you want to generate. Be specific about appearance, style, and mood.
Stable Diffusion has a 77-token limit. Keep your description concise for best results." + ), + ) + + model = forms.ChoiceField( + label=_("AI Model"), + choices=MODEL_CHOICES, + initial="stable_diffusion", + help_text=_("Select the AI model to use for generation."), + ) + + quality = forms.ChoiceField( + label=_("Generation Quality"), + choices=[ + ("low", _("Low (faster, lower quality)")), + ("medium", _("Medium (balanced)")), + ("high", _("High (slower, better quality)")), + ], + initial="medium", + help_text=_("Higher quality takes longer but produces better results."), + ) + + not_porn = forms.BooleanField( + label=_("Suitable for all ages (no offensive content)"), + required=True, + error_messages={ + "required": _( + 'We only host "G-rated" images and so this field must be checked.' + ) + }, + ) + + can_distribute = forms.BooleanField( + label=_("Can be freely copied"), + required=True, + error_messages={ + "required": _( + "This field must be checked since we need to be able to distribute photos to third parties." + ) + }, + ) + + def clean_prompt(self): + """Validate prompt length against token limits""" + prompt = self.cleaned_data.get("prompt", "") + model = self.cleaned_data.get("model", "stable_diffusion") + + if prompt: + from ivatar.ai_service import validate_avatar_prompt + + validation = validate_avatar_prompt(prompt, model) + + if not validation["valid"]: + raise forms.ValidationError(validation["warning"]) + + return prompt + + class DeleteAccountForm(forms.Form): password = forms.CharField( label=_("Password"), required=False, widget=forms.PasswordInput() diff --git a/ivatar/ivataraccount/migrations/0021_add_ai_generation_fields.py b/ivatar/ivataraccount/migrations/0021_add_ai_generation_fields.py new file mode 100644 index 0000000..6559c15 --- /dev/null +++ b/ivatar/ivataraccount/migrations/0021_add_ai_generation_fields.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- +# Generated by Django 5.2.1 on 2025-09-17 10:06 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("ivataraccount", "0020_confirmedopenid_bluesky_handle"), + ] + + operations = [ + migrations.AddField( + model_name="photo", + name="ai_generated", + field=models.BooleanField( + default=False, help_text="Whether this photo was generated by AI" + ), + ), + migrations.AddField( + model_name="photo", + name="ai_model", + field=models.CharField( + blank=True, + help_text="The AI model used for generation", + max_length=50, + null=True, + ), + ), + migrations.AddField( + model_name="photo", + name="ai_prompt", + field=models.TextField( + blank=True, + help_text="The prompt used to generate this image", + null=True, + ), + ), + migrations.AddField( + model_name="photo", + name="ai_quality", + field=models.CharField( + blank=True, + help_text="The quality setting used", + max_length=20, + null=True, + ), + ), + ] diff --git a/ivatar/ivataraccount/migrations/0022_add_generation_task.py b/ivatar/ivataraccount/migrations/0022_add_generation_task.py new file mode 100644 index 0000000..dfc67f7 --- /dev/null +++ b/ivatar/ivataraccount/migrations/0022_add_generation_task.py @@ -0,0 +1,80 @@ +# -*- coding: utf-8 -*- +# Generated by Django 5.2.1 on 2025-09-17 10:25 + +import django.db.models.deletion +import django.utils.timezone +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("ivataraccount", "0021_add_ai_generation_fields"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.CreateModel( + name="GenerationTask", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ( + "ip_address", + models.GenericIPAddressField(null=True, unpack_ipv4=True), + ), + ("add_date", models.DateTimeField(default=django.utils.timezone.now)), + ("prompt", models.TextField()), + ("model", models.CharField(max_length=50)), + ("quality", models.CharField(max_length=20)), + ("allow_nsfw", models.BooleanField(default=False)), + ( + "status", + models.CharField( + choices=[ + ("pending", "Pending"), + ("processing", "Processing"), + ("completed", "Completed"), + ("failed", "Failed"), + ("cancelled", "Cancelled"), + ], + default="pending", + max_length=20, + ), + ), + ("progress", models.IntegerField(default=0)), + ("queue_position", models.IntegerField(default=0)), + ("task_id", models.CharField(blank=True, max_length=255, null=True)), + ("error_message", models.TextField(blank=True, null=True)), + ( + "generated_photo", + models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.SET_NULL, + to="ivataraccount.photo", + ), + ), + ( + "user", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + to=settings.AUTH_USER_MODEL, + ), + ), + ], + options={ + "verbose_name": "Generation Task", + "verbose_name_plural": "Generation Tasks", + "ordering": ["-add_date"], + }, + ), + ] diff --git a/ivatar/ivataraccount/migrations/0023_add_ai_invalid_field.py b/ivatar/ivataraccount/migrations/0023_add_ai_invalid_field.py new file mode 100644 index 0000000..ddfe845 --- /dev/null +++ b/ivatar/ivataraccount/migrations/0023_add_ai_invalid_field.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +# Generated by Django 5.2.1 on 2025-09-17 10:41 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("ivataraccount", "0022_add_generation_task"), + ] + + operations = [ + migrations.AddField( + model_name="photo", + name="ai_invalid", + field=models.BooleanField( + default=False, + help_text="Whether this AI-generated image is invalid (black, etc.)", + ), + ), + ] diff --git a/ivatar/ivataraccount/models.py b/ivatar/ivataraccount/models.py index 25a4d86..9f6ebbb 100644 --- a/ivatar/ivataraccount/models.py +++ b/ivatar/ivataraccount/models.py @@ -118,6 +118,42 @@ class BaseAccountModel(models.Model): abstract = True +class GenerationTask(BaseAccountModel): + """ + Model to track avatar generation tasks in the queue + """ + + STATUS_CHOICES = [ + ("pending", _("Pending")), + ("processing", _("Processing")), + ("completed", _("Completed")), + ("failed", _("Failed")), + ("cancelled", _("Cancelled")), + ] + + user = models.ForeignKey(User, on_delete=models.CASCADE) + prompt = models.TextField() + model = models.CharField(max_length=50) + quality = models.CharField(max_length=20) + allow_nsfw = models.BooleanField(default=False) + status = models.CharField(max_length=20, choices=STATUS_CHOICES, default="pending") + progress = models.IntegerField(default=0) # 0-100 + queue_position = models.IntegerField(default=0) + task_id = models.CharField(max_length=255, blank=True, null=True) # Celery task ID + error_message = models.TextField(blank=True, null=True) + generated_photo = models.ForeignKey( + "Photo", on_delete=models.SET_NULL, blank=True, null=True + ) + + class Meta: + ordering = ["-add_date"] + verbose_name = _("Generation Task") + verbose_name_plural = _("Generation Tasks") + + def __str__(self): + return f"Task {self.pk}: {self.prompt[:50]}... ({self.status})" + + class Photo(BaseAccountModel): """ Model holding the photos and information about them @@ -128,6 +164,27 @@ class Photo(BaseAccountModel): format = models.CharField(max_length=4) access_count = models.BigIntegerField(default=0, editable=False) + # AI Generation metadata + ai_generated = models.BooleanField( + default=False, help_text=_("Whether this photo was generated by AI") + ) + ai_prompt = models.TextField( + blank=True, null=True, help_text=_("The prompt used to generate this image") + ) + ai_model = models.CharField( + max_length=50, + blank=True, + null=True, + help_text=_("The AI model used for generation"), + ) + ai_quality = models.CharField( + max_length=20, blank=True, null=True, help_text=_("The quality setting used") + ) + ai_invalid = models.BooleanField( + default=False, + help_text=_("Whether this AI-generated image is invalid (black, etc.)"), + ) + class Meta: # pylint: disable=too-few-public-methods """ Class attributes @@ -136,6 +193,49 @@ class Photo(BaseAccountModel): verbose_name = _("photo") verbose_name_plural = _("photos") + def is_valid_avatar(self): + """ + Check if this photo is a valid avatar (not black/invalid) + """ + if not self.ai_generated: + return True # Non-AI photos are assumed valid + + # If we've already marked it as invalid, return False + if self.ai_invalid: + return False + + try: + from PIL import Image + import io + + # Load the image data + image_data = io.BytesIO(self.data) + image = Image.open(image_data) + + # Convert to RGB if needed + if image.mode != "RGB": + image = image.convert("RGB") + + # Check if image is predominantly black (common NSFW response) + pixels = list(image.getdata()) + black_pixels = sum(1 for r, g, b in pixels if r == 0 and g == 0 and b == 0) + total_pixels = len(pixels) + + # If more than 95% black pixels, consider it invalid + black_ratio = black_pixels / total_pixels + is_valid = black_ratio < 0.95 + + # Cache the result + if not is_valid: + self.ai_invalid = True + self.save(update_fields=["ai_invalid"]) + + return is_valid + + except Exception: + # If we can't analyze the image, assume it's valid + return True + def import_image(self, service_name, email_address): """ Allow to import image from other (eg. Gravatar) service diff --git a/ivatar/ivataraccount/templates/assign_photo_email.html b/ivatar/ivataraccount/templates/assign_photo_email.html index 32c13e8..6cdc2ff 100644 --- a/ivatar/ivataraccount/templates/assign_photo_email.html +++ b/ivatar/ivataraccount/templates/assign_photo_email.html @@ -78,7 +78,7 @@ {% endif %} - + {% endblock content %} diff --git a/ivatar/ivataraccount/templates/assign_photo_openid.html b/ivatar/ivataraccount/templates/assign_photo_openid.html index a062cac..094d10d 100644 --- a/ivatar/ivataraccount/templates/assign_photo_openid.html +++ b/ivatar/ivataraccount/templates/assign_photo_openid.html @@ -75,7 +75,7 @@ {% endif %} - + {% endblock content %} diff --git a/ivatar/ivataraccount/templates/avatar_gallery.html b/ivatar/ivataraccount/templates/avatar_gallery.html new file mode 100644 index 0000000..dd4bce4 --- /dev/null +++ b/ivatar/ivataraccount/templates/avatar_gallery.html @@ -0,0 +1,318 @@ +{% extends 'base.html' %} +{% load i18n %} +{% load static %} +{% load bootstrap4 %} + +{% block title %}{% trans 'Avatar Gallery' %}{% endblock %} + +{% block extra_css %} + +{% endblock %} + +{% block content %} + +
+
+
+

{% trans 'Avatar Gallery' %}

+

+ {% trans 'Browse recently generated avatars for inspiration. Click on any avatar to reuse its prompt.' %} +

+ + + {% if user_avatars %} +
+

{% trans 'Your Recent Avatars' %}

+
+ {% for avatar in user_avatars %} +
+
+
+
+ Avatar +
+
+ + {{ avatar.add_date|date:"M d, Y" }} + +

+ {{ avatar.ai_prompt|truncatechars:60 }} +

+
+ + {% trans 'Reuse' %} + + + {% trans 'View' %} + +
+
+
+
+
+ {% endfor %} +
+
+ {% endif %} + + +
+

{% trans 'Community Gallery' %}

+

+ {% trans 'Recently generated avatars from all users (last 30)' %} +

+ + {% if avatars %} +
+ {% for avatar in avatars %} +
+
+
+
+ Avatar +
+
+ + {{ avatar.add_date|date:"M d, Y" }} + {% if avatar.user == user %} + {% trans 'Yours' %} + {% endif %} + +

+ {{ avatar.ai_prompt|truncatechars:60 }} +

+
+ + {% trans 'Reuse' %} + + + {% trans 'View' %} + +
+
+
+
+
+ {% endfor %} +
+ + + {% if is_paginated %} + + {% endif %} + + {% else %} +
+ +
{% trans 'No Avatars Yet' %}
+

+ {% trans 'No AI-generated avatars have been created yet. Be the first to generate one!' %} +

+ + {% trans 'Generate Avatar' %} + +
+ {% endif %} +
+ + +
+ + {% trans 'Generate New Avatar' %} + + + {% trans 'Back to Profile' %} + +
+ + +
+
{% trans 'Gallery Tips' %}
+
    +
  • + {% trans 'Click "Reuse" to copy a prompt and modify it for your own avatar' %} +
    • +
  • + {% trans 'Click "View" to see the full avatar and assign it to your emails' %} +
    • +
  • + {% trans 'Use the gallery to find inspiration for your own avatar descriptions' %} +
    • +
  • + {% trans 'Your own avatars are shown at the top for quick access' %} +
    • +
+
+
+
+
+ +{% endblock %} + +{% block extra_js %} + +{% endblock %} diff --git a/ivatar/ivataraccount/templates/avatar_preview.html b/ivatar/ivataraccount/templates/avatar_preview.html new file mode 100644 index 0000000..b1f59ba --- /dev/null +++ b/ivatar/ivataraccount/templates/avatar_preview.html @@ -0,0 +1,198 @@ +{% extends 'base.html' %} +{% load i18n %} +{% load static %} +{% load bootstrap4 %} + +{% block title %}{% trans 'Avatar Preview' %}{% endblock %} + +{% block extra_css %} +{% endblock %} + +{% block content %} + +
+
+
+

{% trans 'Avatar Preview' %}

+

+ {% trans 'Here\'s your generated avatar. You can refine it or assign it to your email addresses.' %} +

+ + +
+
+

{% trans 'Generated Avatar' %}

+
+ Generated Avatar +
+

+ {% trans 'Generated on' %} {{ photo.add_date|date:"F d, Y \a\t H:i" }} +

+
+
+ + +
+
+
{% trans 'Refine Avatar' %}
+
+
+

+ {% trans 'Not satisfied with the result? Modify your description and generate a new avatar.' %} +

+ +
+ {% csrf_token %} + + +
+ + + + {{ form.prompt.help_text|safe }} + +
+ +
+ + + {{ form.model.help_text }} +
+ +
+ + + {{ form.quality.help_text }} +
+ +
+ + +
+ +
+ + +
+ + +
+ +
+
+
+
+ + + {% if confirmed_emails %} +
+
+
{% trans 'Assign to Email Addresses' %}
+
+
+

+ {% trans 'Assign this avatar to one or more of your confirmed email addresses.' %} +

+ +
+ {% for email in confirmed_emails %} +
+
+ {{ email.email }} + {% if email.photo_id == photo.pk %} + {% trans 'Currently assigned' %} + {% endif %} +
+
+ {% if email.photo_id != photo.pk %} + + {% trans 'Assign' %} + + {% else %} + + {% trans 'Assigned' %} + + {% endif %} +
+
+ {% endfor %} +
+
+
+ {% else %} +
+
+
{% trans 'No Email Addresses' %}
+

+ {% trans 'You need to add and confirm email addresses before you can assign avatars to them.' %} +

+ + {% trans 'Add Email Address' %} + +
+
+ {% endif %} + + +
+
+
+
+ + {% trans 'Generate New Avatar' %} + +
+
+ + {% trans 'Back to Profile' %} + +
+
+
+
+ + +
+
{% trans 'Tips for Better Results' %}
+
    +
  • + {% trans 'Be more specific: "A friendly robot with blue LED eyes and silver metallic body"' %} +
    • +
  • + {% trans 'Add style keywords: "cartoon style", "realistic", "anime", "pixel art"' %} +
    • +
  • + {% trans 'Include mood: "cheerful", "serious", "mysterious", "professional"' %} +
    • +
  • + {% trans 'Specify lighting: "soft lighting", "dramatic shadows", "bright and clear"' %} +
    • +
+
+
+
+
+ +{% endblock %} + +{% block extra_js %} +{% endblock %} diff --git a/ivatar/ivataraccount/templates/generate_avatar.html b/ivatar/ivataraccount/templates/generate_avatar.html new file mode 100644 index 0000000..0fa17c7 --- /dev/null +++ b/ivatar/ivataraccount/templates/generate_avatar.html @@ -0,0 +1,145 @@ +{% extends 'base.html' %} +{% load i18n %} +{% load static %} +{% load bootstrap4 %} + +{% block title %}{% trans 'Generate AI Avatar' %}{% endblock %} + +{% block extra_css %} +{% endblock %} + +{% block content %} + +
+
+
+

{% trans 'Generate AI Avatar' %}

+

+ {% trans 'Create a unique avatar using artificial intelligence. Describe what you want and our AI will generate it for you.' %} +

+ +
+
+
+ {% csrf_token %} + + +
+ + + + {{ form.prompt.help_text|safe }} + +
+ +
+ + + {{ form.model.help_text }} +
+ +
+ + + {{ form.quality.help_text }} +
+ +
+ + +
+ +
+ + +
+ +
+ + + {% trans 'Cancel' %} + +
+
+
+
+ +
+

{% trans 'Tips for Better Results' %}

+
    +
  • + {% trans 'Be specific about appearance, style, and mood' %} +
    • +
  • + {% trans 'Include details like hair color, clothing, or facial expressions' %} +
    • +
  • + {% trans 'Try different art styles: "cartoon", "realistic", "anime", "pixel art"' %} +
    • +
  • + {% trans 'Keep descriptions appropriate for all ages' %} +
    • +
+
+ +
+

{% trans 'Example Prompts' %}

+
+
+
+
+
{% trans 'Character Avatar' %}
+

+ "A friendly robot with blue eyes and a silver body, cartoon style" +

+
+
+
+
+
+
+
{% trans 'Abstract Avatar' %}
+

+ "Colorful geometric shapes forming a face, modern art style" +

+
+
+
+
+
+ +
+
{% trans 'Important Notes' %}
+
    +
  • {% trans 'Avatar generation may take 30-60 seconds depending on server load' %}
    • +
  • {% trans 'Generated avatars are automatically saved to your account' %}
    • +
  • {% trans 'You can assign the generated avatar to any of your email addresses' %}
    • +
  • {% trans 'All generated content must be appropriate for all ages' %}
    • +
+
+
+
+
+ +{% endblock %} + +{% block extra_js %} +{% endblock %} diff --git a/ivatar/ivataraccount/templates/generation_status.html b/ivatar/ivataraccount/templates/generation_status.html new file mode 100644 index 0000000..2e65d67 --- /dev/null +++ b/ivatar/ivataraccount/templates/generation_status.html @@ -0,0 +1,405 @@ +{% extends 'base.html' %} +{% load i18n %} +{% load static %} + +{% block title %}{% trans 'Avatar Generation Status' %}{% endblock %} + +{% block content %} +
+
+
+

{% trans 'Avatar Generation Status' %}

+ +
+
+
{% trans 'Generation Task' %}
+
+
+
+
+

{% trans 'Prompt:' %}

+

{{ task.prompt }}

+
+
+

{% trans 'Model:' %} {{ task.model|title }}

+

{% trans 'Quality:' %} {{ task.quality|title }}

+

{% trans 'Status:' %} + + {{ task.get_status_display }} + {% if task.status == 'processing' %} + + {% elif task.status == 'pending' %} + + {% endif %} + + + + Live + +

+
+
+ + {% if task.status == 'processing' %} +
+
+
+
+ {{ task.progress }}% +
+
+
+
+
+ + {% trans 'Generating your avatar...' %} + +
+
+ + + {% trans 'Updated just now' %} + +
+
+
+
+
+ {% elif task.status == 'pending' %} +
+
+ + {% trans 'Your avatar is in the queue. Position:' %} {{ queue_position }} + {% if queue_length > 1 %} + {% trans 'out of' %} {{ queue_length }} {% trans 'tasks' %} + {% endif %} +
+
+
+ 0% +
+
+ {% trans 'Waiting in queue...' %} +
+ {% elif task.status == 'completed' %} +
+
+ + {% trans 'Avatar generated successfully!' %} +
+ {% if task.generated_photo %} +
+ + {% trans 'View Avatar' %} + +
+ {% endif %} +
+ {% elif task.status == 'failed' %} +
+
+ + {% trans 'Avatar generation failed.' %} + {% if task.error_message %} +
{{ task.error_message }} + {% endif %} +
+
+ + {% trans 'Try Again' %} + +
+
+ {% endif %} +
+
+ + +
+
+
{% trans 'Queue Information' %}
+
+
+
+
+
+

{{ processing_count }}

+

{% trans 'Currently Processing' %}

+
+
+
+
+

{{ queue_length }}

+

{% trans 'Pending Tasks' %}

+
+
+
+
+

2

+

{% trans 'Max Parallel Jobs' %}

+
+
+
+
+
+ + + {% if user_tasks %} +
+
+
{% trans 'Your Recent Tasks' %}
+
+
+
+ + + + + + + + + + + {% for user_task in user_tasks %} + + + + + + + {% endfor %} + +
{% trans 'Prompt' %}{% trans 'Status' %}{% trans 'Created' %}{% trans 'Actions' %}
+ + {{ user_task.prompt }} + + + + {{ user_task.get_status_display }} + + {{ user_task.add_date|date:"M d, H:i" }} + {% if user_task.status == 'completed' and user_task.generated_photo %} + + + + {% elif user_task.status == 'failed' %} + + + + {% elif user_task.status == 'pending' or user_task.status == 'processing' %} + + + + {% endif %} +
+
+
+
+ {% endif %} + +
+ + {% trans 'Generate New Avatar' %} + + + {% trans 'Back to Profile' %} + +
+
+
+
+{% endblock %} + +{% block extra_js %} + + + +{% endblock %} diff --git a/ivatar/ivataraccount/urls.py b/ivatar/ivataraccount/urls.py index ae6a1c4..0e85d53 100644 --- a/ivatar/ivataraccount/urls.py +++ b/ivatar/ivataraccount/urls.py @@ -27,6 +27,14 @@ from .views import ResendConfirmationMailView from .views import IvatarLoginView from .views import DeleteAccountView from .views import ExportView +from .views import ( + GenerateAvatarView, + AvatarPreviewView, + AvatarGalleryView, + ReusePromptView, + GenerationStatusView, + task_status_api, +) # Define URL patterns, self documenting # To see the fancy, colorful evaluation of these use: @@ -77,6 +85,26 @@ urlpatterns = [ # pylint: disable=invalid-name ProfileView.as_view(), name="profile_with_profile_username", ), + path("generate_avatar/", GenerateAvatarView.as_view(), name="generate_avatar"), + re_path( + r"generation_status/(?P\d+)/", + GenerationStatusView.as_view(), + name="generation_status", + ), + re_path( + r"api/task_status/(?P\d+)/", task_status_api, name="task_status_api" + ), + re_path( + r"avatar_preview/(?P\d+)/", + AvatarPreviewView.as_view(), + name="avatar_preview", + ), + path("avatar_gallery/", AvatarGalleryView.as_view(), name="avatar_gallery"), + re_path( + r"reuse_prompt/(?P\d+)/", + ReusePromptView.as_view(), + name="reuse_prompt", + ), path("add_email/", AddEmailView.as_view(), name="add_email"), path("add_openid/", AddOpenIDView.as_view(), name="add_openid"), path("upload_photo/", UploadPhotoView.as_view(), name="upload_photo"), diff --git a/ivatar/ivataraccount/views.py b/ivatar/ivataraccount/views.py index 1cca502..f435596 100644 --- a/ivatar/ivataraccount/views.py +++ b/ivatar/ivataraccount/views.py @@ -21,8 +21,9 @@ from django.utils.decorators import method_decorator from django.contrib.messages.views import SuccessMessageMixin from django.contrib import messages from django.views.generic.edit import FormView, UpdateView -from django.views.generic.base import View, TemplateView +from django.views.generic.base import View, TemplateView, RedirectView from django.views.generic.detail import DetailView +from django.views.generic.list import ListView from django.contrib.auth import authenticate, login from django.contrib.auth.forms import UserCreationForm, SetPasswordForm from django.contrib.auth.views import LoginView @@ -30,9 +31,11 @@ from django.contrib.auth.views import ( PasswordResetView as PasswordResetViewOriginal, ) from django.utils.translation import gettext_lazy as _ -from django.http import HttpResponseRedirect, HttpResponse +from django.http import HttpResponseRedirect, HttpResponse, JsonResponse from django.urls import reverse_lazy, reverse from django.shortcuts import render +from django.views.decorators.http import require_http_methods +import logging from django_openid_auth.models import UserOpenID from openid import oidutil @@ -54,12 +57,17 @@ from .gravatar import get_photo as get_gravatar_photo from .forms import AddEmailForm, UploadPhotoForm, AddOpenIDForm from .forms import UpdatePreferenceForm, UploadLibravatarExportForm -from .forms import DeleteAccountForm +from .forms import DeleteAccountForm, GenerateAvatarForm from .models import UnconfirmedEmail, ConfirmedEmail, Photo from .models import UnconfirmedOpenId, ConfirmedOpenId, DjangoOpenIDStore from .models import UserPreference from .models import file_format from .read_libravatar_export import read_gzdata as libravatar_read_gzdata +from ivatar.ai_service import generate_avatar_image, AIServiceError +from ivatar.tasks import generate_avatar_task, update_queue_positions +from ivatar.ivataraccount.models import GenerationTask + +logger = logging.getLogger(__name__) def openid_logging(message, level=0): @@ -1367,3 +1375,358 @@ class ExportView(SuccessMessageMixin, TemplateView): ] = f'attachment; filename="libravatar-export_{user.username}.xml.gz"' response.write(bytesobj.read()) return response + + +class GenerateAvatarView(SuccessMessageMixin, FormView): + """ + View for generating avatars using AI text-to-image + """ + + template_name = "generate_avatar.html" + form_class = GenerateAvatarForm + success_message = _("Avatar generated successfully!") + + @method_decorator(login_required) + def dispatch(self, *args, **kwargs): + return super().dispatch(*args, **kwargs) + + def get_initial(self): + """Pre-populate form with reused prompt if available""" + initial = super().get_initial() + + # Check for reused prompt from gallery + reuse_prompt = self.request.session.get("reuse_prompt", "") + if reuse_prompt: + initial["prompt"] = reuse_prompt + # Clear the reused prompt from session + del self.request.session["reuse_prompt"] + + return initial + + def form_valid(self, form): + """ + Handle form submission and queue avatar generation + """ + try: + # Get form data + prompt = form.cleaned_data["prompt"] + model = form.cleaned_data["model"] + quality = form.cleaned_data["quality"] + + # Create generation task + task = GenerationTask.objects.create( + user=self.request.user, + prompt=prompt, + model=model, + quality=quality, + allow_nsfw=False, # Always false - no NSFW override allowed + status="pending", + ) + + # Update queue positions + update_queue_positions.delay() + + # Queue the generation task + celery_task = generate_avatar_task.delay( + task_id=task.pk, + user_id=self.request.user.pk, + prompt=prompt, + model=model, + quality=quality, + allow_nsfw=False, # Always false - no NSFW override allowed + ) + + # Store task ID + task.task_id = celery_task.id + task.save() + + # Store prompt in session for refinement + self.request.session["last_avatar_prompt"] = prompt + self.request.session["user_consent_given"] = True + + messages.success( + self.request, + _("Avatar generation queued! You'll be notified when it's ready."), + ) + + # Redirect to task status page + return HttpResponseRedirect( + reverse_lazy("generation_status", kwargs={"task_id": task.pk}) + ) + + except Exception as e: + logger.error(f"Unexpected error in avatar generation: {e}") + messages.error( + self.request, _("An unexpected error occurred. Please try again.") + ) + return self.form_invalid(form) + + +class GenerationStatusView(TemplateView): + """ + View for showing avatar generation status and progress + """ + + template_name = "generation_status.html" + + @method_decorator(login_required) + def dispatch(self, *args, **kwargs): + return super().dispatch(*args, **kwargs) + + def get_context_data(self, **kwargs): + context = super().get_context_data(**kwargs) + task_id = kwargs.get("task_id") + + try: + task = GenerationTask.objects.get(pk=task_id, user=self.request.user) + context["task"] = task + + # Get queue information + pending_tasks = GenerationTask.objects.filter(status="pending").order_by( + "add_date" + ) + processing_tasks = GenerationTask.objects.filter(status="processing") + + context["queue_length"] = pending_tasks.count() + context["processing_count"] = processing_tasks.count() + context["queue_position"] = task.queue_position + + # Get user's other tasks + user_tasks = GenerationTask.objects.filter(user=self.request.user).order_by( + "-add_date" + )[:5] + context["user_tasks"] = user_tasks + + except GenerationTask.DoesNotExist: + messages.error(self.request, _("Generation task not found.")) + return HttpResponseRedirect(reverse_lazy("profile")) + + return context + + +class AvatarPreviewView(SuccessMessageMixin, FormView): + """ + View for previewing generated avatars and allowing refinements + """ + + template_name = "avatar_preview.html" + form_class = GenerateAvatarForm + success_message = _("Avatar regenerated successfully!") + + @method_decorator(login_required) + def dispatch(self, *args, **kwargs): + return super().dispatch(*args, **kwargs) + + def get_context_data(self, **kwargs): + """Add photo and related data to context""" + context = super().get_context_data(**kwargs) + + try: + photo = Photo.objects.get( + pk=self.kwargs["photo_id"], user=self.request.user + ) + context["photo"] = photo + context["photo_url"] = reverse("raw_image", kwargs={"pk": photo.pk}) + + # Get user's confirmed emails for assignment + context["confirmed_emails"] = self.request.user.confirmedemail_set.all() + + except Photo.DoesNotExist: + messages.error(self.request, _("Avatar not found.")) + return HttpResponseRedirect(reverse_lazy("profile")) + + return context + + def get_initial(self): + """Pre-populate form with current prompt if available""" + initial = super().get_initial() + initial["model"] = "stable_diffusion" + initial["quality"] = "medium" + + # Try to get the prompt from the session or URL parameters + prompt = self.request.session.get("last_avatar_prompt", "") + if prompt: + initial["prompt"] = prompt + + # Pre-check consent checkboxes since user already gave consent + if self.request.session.get("user_consent_given", False): + initial["not_porn"] = True + initial["can_distribute"] = True + + return initial + + def form_valid(self, form): + """ + Handle refinement - generate new avatar with modified prompt + """ + try: + # Generate new avatar with refined prompt + prompt = form.cleaned_data["prompt"] + model = form.cleaned_data["model"] + quality = form.cleaned_data["quality"] + + generated_image = generate_avatar_image( + prompt=prompt, + model=model, + size=(512, 512), + quality=quality, + allow_nsfw=False, # Always false - no NSFW override allowed + ) + + # Convert PIL image to bytes + img_buffer = BytesIO() + generated_image.save(img_buffer, format="PNG") + img_data = img_buffer.getvalue() + + # Create new Photo object + new_photo = Photo() + new_photo.user = self.request.user + new_photo.ip_address = get_client_ip(self.request)[0] + new_photo.data = img_data + new_photo.format = "png" + + # Store AI generation metadata + new_photo.ai_generated = True + new_photo.ai_prompt = prompt + new_photo.ai_model = model + new_photo.ai_quality = "medium" # Default quality + + new_photo.save() + + # Store the new prompt and preserve consent in session for further refinement + self.request.session["last_avatar_prompt"] = prompt + self.request.session["user_consent_given"] = True + + messages.success( + self.request, + _( + "Avatar regenerated successfully! You can refine it further or assign it to your email addresses." + ), + ) + + # Redirect to preview the new avatar + return HttpResponseRedirect( + reverse_lazy("avatar_preview", kwargs={"photo_id": new_photo.pk}) + ) + + except Photo.DoesNotExist: + messages.error(self.request, _("Original avatar not found.")) + return HttpResponseRedirect(reverse_lazy("profile")) + except AIServiceError as e: + messages.error( + self.request, + _("Failed to regenerate avatar: %(error)s") % {"error": str(e)}, + ) + return self.form_invalid(form) + except Exception as e: + messages.error( + self.request, + _("An unexpected error occurred: %(error)s") % {"error": str(e)}, + ) + return self.form_invalid(form) + + +class AvatarGalleryView(ListView): + """ + View for displaying a gallery of recent AI-generated avatars + """ + + template_name = "avatar_gallery.html" + context_object_name = "avatars" + paginate_by = 12 # Show 12 avatars per page + + @method_decorator(login_required) + def dispatch(self, *args, **kwargs): + return super().dispatch(*args, **kwargs) + + def get_queryset(self): + """Get the last 30 AI-generated avatars from all users, excluding invalid ones""" + return ( + Photo.objects.filter( + ai_generated=True, + ai_prompt__isnull=False, + ai_invalid=False, # Exclude invalid images + ) + .exclude(ai_prompt="") + .order_by("-add_date")[:30] + ) + + def get_context_data(self, **kwargs): + """Add additional context data""" + context = super().get_context_data(**kwargs) + + # Add user's own avatars for quick access (filtered for validity) + context["user_avatars"] = ( + Photo.objects.filter( + user=self.request.user, + ai_generated=True, + ai_prompt__isnull=False, + ai_invalid=False, # Exclude invalid images + ) + .exclude(ai_prompt="") + .order_by("-add_date")[:10] + ) + + return context + + +class ReusePromptView(RedirectView): + """ + View to reuse a prompt from the gallery + """ + + permanent = False + + def get_redirect_url(self, *args, **kwargs): + """Redirect to generate avatar page with pre-filled prompt""" + try: + photo = Photo.objects.get( + pk=kwargs["photo_id"], ai_generated=True, ai_prompt__isnull=False + ) + + # Store the prompt in session for the generate form + self.request.session["reuse_prompt"] = photo.ai_prompt + + # Redirect to generate avatar page + return reverse_lazy("generate_avatar") + + except Photo.DoesNotExist: + messages.error(self.request, _("Avatar not found.")) + return reverse_lazy("avatar_gallery") + + +@method_decorator(login_required, name="dispatch") +@require_http_methods(["GET"]) +def task_status_api(request, task_id): + """ + API endpoint to get task status for AJAX requests + """ + try: + task = GenerationTask.objects.get(pk=task_id, user=request.user) + + # Get queue information + pending_tasks = GenerationTask.objects.filter(status="pending").order_by( + "add_date" + ) + processing_tasks = GenerationTask.objects.filter(status="processing") + + data = { + "status": task.status, + "progress": task.progress, + "queue_position": task.queue_position, + "queue_length": pending_tasks.count(), + "processing_count": processing_tasks.count(), + "error_message": task.error_message, + "generated_photo_id": task.generated_photo.pk + if task.generated_photo + else None, + "status_display": task.get_status_display(), + } + + return JsonResponse(data) + + except GenerationTask.DoesNotExist: + return JsonResponse({"error": "Task not found"}, status=404) + except Exception as e: + logger.error(f"Error in task status API: {e}") + return JsonResponse({"error": "Internal server error"}, status=500) diff --git a/ivatar/settings.py b/ivatar/settings.py index 30e0232..e0757c5 100644 --- a/ivatar/settings.py +++ b/ivatar/settings.py @@ -33,6 +33,7 @@ INSTALLED_APPS = [ "django.contrib.messages", "django.contrib.staticfiles", "social_django", + "django_celery_results", ] MIDDLEWARE = [ diff --git a/ivatar/tasks.py b/ivatar/tasks.py new file mode 100644 index 0000000..a7287e4 --- /dev/null +++ b/ivatar/tasks.py @@ -0,0 +1,182 @@ +# -*- coding: utf-8 -*- +""" +Celery tasks for avatar generation +""" + +import logging +from celery import shared_task +from django.contrib.auth.models import User +from io import BytesIO + +from ivatar.ai_service import generate_avatar_image, AIServiceError +from ivatar.ivataraccount.models import GenerationTask, Photo + +logger = logging.getLogger(__name__) + + +@shared_task(bind=True, name="ivatar.tasks.generate_avatar_task") +def generate_avatar_task(self, task_id, user_id, prompt, model, quality, allow_nsfw): + """ + Background task to generate avatar images + + Args: + task_id: GenerationTask ID + user_id: User ID + prompt: Avatar description + model: AI model to use + quality: Generation quality + allow_nsfw: Whether to allow NSFW content + + Returns: + dict: Task result with photo_id or error + """ + try: + # Get the task object + task = GenerationTask.objects.get(pk=task_id) + + # Update task status + task.status = "processing" + task.task_id = self.request.id + task.progress = 10 + task.save() + + logger.info(f"Starting avatar generation task {task_id} for user {user_id}") + + # Update progress + self.update_state( + state="PROGRESS", + meta={"progress": 20, "status": "Initializing generation..."}, + ) + task.progress = 20 + task.save() + + # Generate the avatar image + self.update_state( + state="PROGRESS", meta={"progress": 30, "status": "Generating image..."} + ) + task.progress = 30 + task.save() + + generated_image = generate_avatar_image( + prompt=prompt, + model=model, + size=(512, 512), + quality=quality, + allow_nsfw=allow_nsfw, + ) + + # Update progress + self.update_state( + state="PROGRESS", meta={"progress": 70, "status": "Processing image..."} + ) + task.progress = 70 + task.save() + + # Convert PIL image to bytes + img_buffer = BytesIO() + generated_image.save(img_buffer, format="PNG") + img_data = img_buffer.getvalue() + + # Get user + user = User.objects.get(pk=user_id) + + # Create Photo object + photo = Photo() + photo.user = user + photo.ip_address = "127.0.0.1" # Default IP for background tasks + photo.data = img_data + photo.format = "png" + + # Store AI generation metadata + photo.ai_generated = True + photo.ai_prompt = prompt + photo.ai_model = model + photo.ai_quality = quality + + photo.save() + + # Update progress + self.update_state( + state="PROGRESS", meta={"progress": 90, "status": "Saving avatar..."} + ) + task.progress = 90 + task.save() + + # Update task with completed status + task.status = "completed" + task.progress = 100 + task.generated_photo = photo + task.save() + + logger.info( + f"Completed avatar generation task {task_id}, created photo {photo.pk}" + ) + + return {"status": "completed", "photo_id": photo.pk, "task_id": task_id} + + except AIServiceError as e: + logger.error(f"AI service error in task {task_id}: {e}") + task.status = "failed" + task.error_message = str(e) + task.progress = 0 + task.save() + + return {"status": "failed", "error": str(e), "task_id": task_id} + + except Exception as e: + logger.error(f"Unexpected error in task {task_id}: {e}") + task.status = "failed" + task.error_message = str(e) + task.progress = 0 + task.save() + + return { + "status": "failed", + "error": f"Unexpected error: {str(e)}", + "task_id": task_id, + } + + +@shared_task(name="ivatar.tasks.update_queue_positions") +def update_queue_positions(): + """ + Update queue positions for pending tasks + """ + try: + pending_tasks = GenerationTask.objects.filter(status="pending").order_by( + "add_date" + ) + + for index, task in enumerate(pending_tasks): + task.queue_position = index + 1 + task.save() + + logger.info(f"Updated queue positions for {len(pending_tasks)} pending tasks") + + except Exception as e: + logger.error(f"Error updating queue positions: {e}") + + +@shared_task(name="ivatar.tasks.cleanup_old_tasks") +def cleanup_old_tasks(): + """ + Clean up old completed/failed tasks + """ + try: + from django.utils import timezone + from datetime import timedelta + + # Delete tasks older than 7 days + cutoff_date = timezone.now() - timedelta(days=7) + + old_tasks = GenerationTask.objects.filter( + add_date__lt=cutoff_date, status__in=["completed", "failed", "cancelled"] + ) + + count = old_tasks.count() + old_tasks.delete() + + logger.info(f"Cleaned up {count} old tasks") + + except Exception as e: + logger.error(f"Error cleaning up old tasks: {e}") diff --git a/templates/_account_bar.html b/templates/_account_bar.html index fa1ebd2..88f4eba 100644 --- a/templates/_account_bar.html +++ b/templates/_account_bar.html @@ -12,6 +12,8 @@