mirror of
https://github.com/khoaliber/khoj.git
synced 2026-03-02 13:18:18 +00:00
Drop image generation support for Stability AI models
Reduces maintenance burden by dropping support for old ai model providers
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
# Generated by Django 5.1.14 on 2025-11-30 00:47
|
||||
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
def delete_stability_configs(apps, schema_editor):
|
||||
TextToImageModelConfig = apps.get_model("database", "TextToImageModelConfig")
|
||||
TextToImageModelConfig.objects.filter(model_type="stability-ai").delete()
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
dependencies = [
|
||||
("database", "0097_serverchatsettings_priority"),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.RunPython(delete_stability_configs, migrations.RunPython.noop),
|
||||
migrations.AlterField(
|
||||
model_name="texttoimagemodelconfig",
|
||||
name="model_type",
|
||||
field=models.CharField(
|
||||
choices=[
|
||||
("openai", "Openai"),
|
||||
("replicate", "Replicate"),
|
||||
("google", "Google"),
|
||||
],
|
||||
default="openai",
|
||||
max_length=200,
|
||||
),
|
||||
),
|
||||
]
|
||||
@@ -575,7 +575,6 @@ class SearchModelConfig(DbBaseModel):
|
||||
class TextToImageModelConfig(DbBaseModel):
|
||||
class ModelType(models.TextChoices):
|
||||
OPENAI = "openai"
|
||||
STABILITYAI = "stability-ai"
|
||||
REPLICATE = "replicate"
|
||||
GOOGLE = "google"
|
||||
|
||||
|
||||
@@ -112,8 +112,6 @@ async def text_to_image(
|
||||
webp_image_bytes = generate_image_with_openai(
|
||||
image_prompt, text_to_image_config, text2image_model, image_shape
|
||||
)
|
||||
elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.STABILITYAI:
|
||||
webp_image_bytes = generate_image_with_stability(image_prompt, text_to_image_config, text2image_model)
|
||||
elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.REPLICATE:
|
||||
webp_image_bytes = generate_image_with_replicate(
|
||||
image_prompt, text_to_image_config, text2image_model, image_shape
|
||||
@@ -220,35 +218,6 @@ def generate_image_with_openai(
|
||||
return convert_image_to_webp(base64.b64decode(image))
|
||||
|
||||
|
||||
@retry(
|
||||
retry=retry_if_exception_type(requests.RequestException),
|
||||
wait=wait_random_exponential(min=1, max=10),
|
||||
stop=stop_after_attempt(3),
|
||||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||
reraise=True,
|
||||
)
|
||||
def generate_image_with_stability(
|
||||
improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
|
||||
):
|
||||
"Generate image using Stability AI"
|
||||
|
||||
# Call Stability AI API to generate image
|
||||
response = requests.post(
|
||||
"https://api.stability.ai/v2beta/stable-image/generate/sd3",
|
||||
headers={"authorization": f"Bearer {text_to_image_config.api_key}", "accept": "image/*"},
|
||||
files={"none": ""},
|
||||
data={
|
||||
"prompt": improved_image_prompt,
|
||||
"model": text2image_model,
|
||||
"mode": "text-to-image",
|
||||
"output_format": "png",
|
||||
"aspect_ratio": "1:1",
|
||||
},
|
||||
)
|
||||
# Convert png to webp for faster loading
|
||||
return convert_image_to_webp(response.content)
|
||||
|
||||
|
||||
@retry(
|
||||
retry=retry_if_exception_type(requests.RequestException),
|
||||
wait=wait_random_exponential(min=1, max=10),
|
||||
|
||||
Reference in New Issue
Block a user