repository

package
v0.0.0-...-bfbcaf2 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Apr 25, 2024 License: MIT Imports: 47 Imported by: 15

Documentation

Overview

Description: Processes realtime messages from cog and updates the database

Index

Constants

View Source 
const FREE_CREDIT_TYPE_ID = "3b12b23e-478b-4c18-8e34-70b3f0af1ee6"
View Source 
const MOCK_ADMIN_UUID = "00000000-0000-0000-0000-000000000000"

Mock user IDs

View Source 
const MOCK_ALT_UUID = "00000000-0000-0000-0000-000000000003"
View Source 
const MOCK_GENERATION_MODEL_ID = "b972a2b8-f39e-4ee3-a670-05e3acdd821c"

Mock generation model IDs and scheduler IDs

View Source 
const MOCK_NORMAL_UUID = "00000000-0000-0000-0000-000000000001"
View Source 
const MOCK_NO_CREDITS_UUID = "00000000-0000-0000-0000-000000000002"
View Source 
const MOCK_SCHEDULER_ID = "b4dff6e9-91a7-449b-b1a7-c25000e3ccd0"
View Source 
const MOCK_UPSCALE_MODEL_ID = "b972a2b8-f39e-4ee3-a670-05e3acdd821e"

Mock upscale

View Source 
const MOCK_VOICEOVER_MODEL_ID = "b972a2b8-f39e-4ee3-a670-05e3acdd821f"

Mock voiceover

View Source 
const MOCK_VOICEOVER_SPEAKER_ID = "b4dff6e9-91a7-449b-b1a7-c25000e3ccd1"
View Source 
const REFUND_CREDIT_TYPE_ID = "7ca94fd6-c201-4ca6-a9bf-4473c83e30b4"
View Source 
const TIPPABLE_CREDIT_TYPE_ID = "74cc89d3-fb00-4cca-9573-95c6786cf186"
View Source 
const TIPPED_CREDIT_TYPE_ID = "d7e2e70e-aaa4-475c-a522-1f9c1c28f25f"

Variables

View Source 
var NEVER_EXPIRE = time.Date(2100, 1, 1, 5, 0, 0, 0, time.UTC)

Expiration date for manual invoices (non-recurring)

View Source 
var Now = time.Now

For mocking

View Source 
var UsernameExistsErr = fmt.Errorf("username_exists")

Functions

This section is empty.

Types

type CachedCount 

type CachedCount struct {
	Count    int       `json:"count"`
	CachedAt time.Time `json:"cached_at"`
}

type GalleryData 

type GalleryData struct {
	ID                 uuid.UUID  `json:"id,omitempty" sql:"id"`
	ImageURL           string     `json:"image_url"`
	UpscaledImageURL   string     `json:"upscaled_image_url,omitempty"`
	CreatedAt          time.Time  `json:"created_at" sql:"created_at"`
	UpdatedAt          time.Time  `json:"updated_at" sql:"updated_at"`
	Width              int32      `json:"width" sql:"generation_width"`
	Height             int32      `json:"height" sql:"generation_height"`
	InferenceSteps     int32      `json:"inference_steps" sql:"generation_inference_steps"`
	GuidanceScale      float32    `json:"guidance_scale" sql:"generation_guidance_scale"`
	Seed               int        `json:"seed,omitempty" sql:"generation_seed"`
	ModelID            uuid.UUID  `json:"model_id" sql:"model_id"`
	SchedulerID        uuid.UUID  `json:"scheduler_id" sql:"scheduler_id"`
	PromptText         string     `json:"prompt_text" sql:"prompt_text"`
	PromptID           uuid.UUID  `json:"prompt_id" sql:"prompt_id"`
	NegativePromptText string     `json:"negative_prompt_text,omitempty" sql:"negative_prompt_text"`
	NegativePromptID   *uuid.UUID `json:"negative_prompt_id,omitempty" sql:"negative_prompt_id"`
	UserID             *uuid.UUID `json:"user_id,omitempty" sql:"user_id"`
	Score              *float32   `json:"score,omitempty" sql:"score"`
	Username           *string    `json:"username,omitempty" sql:"username"`
	User               *UserType  `json:"user,omitempty" sql:"user"`
	PromptStrength     *float32   `json:"prompt_strength,omitempty" sql:"prompt_strength"`
	WasAutoSubmitted   bool       `json:"was_auto_submitted" sql:"was_auto_submitted"`
	IsPublic           bool       `json:"is_public" sql:"is_public"`
	LikeCount          int        `json:"like_count" sql:"like_count"`
	LikeCountTrending  *int       `json:"like_count_trending,omitempty" sql:"like_count_trending"`
	IsLiked            *bool      `json:"is_liked,omitempty" sql:"liked_by_user"`
}

type GenerationQueryWithOutputsData 

type GenerationQueryWithOutputsData struct {
	ID                     uuid.UUID                 `json:"id" sql:"id"`
	Height                 int32                     `json:"height" sql:"height"`
	Width                  int32                     `json:"width" sql:"width"`
	InferenceSteps         int32                     `json:"inference_steps" sql:"inference_steps"`
	Seed                   int                       `json:"seed" sql:"seed"`
	Status                 string                    `json:"status" sql:"status"`
	GuidanceScale          float32                   `json:"guidance_scale" sql:"guidance_scale"`
	SchedulerID            uuid.UUID                 `json:"scheduler_id" sql:"scheduler_id"`
	ModelID                uuid.UUID                 `json:"model_id" sql:"model_id"`
	PromptID               *uuid.UUID                `json:"prompt_id,omitempty" sql:"prompt_id"`
	NegativePromptID       *uuid.UUID                `json:"negative_prompt_id,omitempty" sql:"negative_prompt_id"`
	CreatedAt              time.Time                 `json:"created_at" sql:"created_at"`
	UpdatedAt              time.Time                 `json:"updated_at" sql:"updated_at"`
	StartedAt              *time.Time                `json:"started_at,omitempty" sql:"started_at"`
	CompletedAt            *time.Time                `json:"completed_at,omitempty" sql:"completed_at"`
	Username               string                    `json:"username,omitempty" sql:"username"`
	UserID                 uuid.UUID                 `json:"user_id,omitempty" sql:"user_id"`
	IsFavorited            bool                      `json:"is_favorited" sql:"is_favorited"`
	IsPublic               bool                      `json:"is_public" sql:"is_public"`
	Outputs                []GenerationUpscaleOutput `json:"outputs"`
	Prompt                 PromptType                `json:"prompt"`
	NegativePrompt         *PromptType               `json:"negative_prompt,omitempty"`
	WasAutoSubmitted       bool                      `json:"was_auto_submitted" sql:"was_auto_submitted"`
	InitImageURL           string                    `json:"init_image_url,omitempty" sql:"init_image_url"`
	InitImageURLSigned     string                    `json:"init_image_url_signed,omitempty"`
	User                   *UserType                 `json:"user,omitempty"`
	PromptStrength         *float32                  `json:"prompt_strength,omitempty" sql:"prompt_strength"`
	AestheticRatingScore   *float32                  `json:"aesthetic_rating_score,omitempty"`
	AestheticArtifactScore *float32                  `json:"aesthetic_artifact_score,omitempty"`
}

type GenerationQueryWithOutputsMeta 

type GenerationQueryWithOutputsMeta[T GenerationQueryWithOutputsMetaCursor] struct {
	Total   *int                                        `json:"total_count,omitempty"`
	Outputs []GenerationQueryWithOutputsResultFormatted `json:"outputs"`
	Next    T                                           `json:"next,omitempty"`
}

Paginated meta for querying generations

type GenerationQueryWithOutputsMetaCursor 

type GenerationQueryWithOutputsMetaCursor interface {
	*uint | *time.Time
}

type GenerationQueryWithOutputsResult 

type GenerationQueryWithOutputsResult struct {
	OutputID               *uuid.UUID                     `json:"output_id,omitempty" sql:"output_id"`
	IsLiked                *bool                          `json:"is_liked,omitempty" sql:"is_liked"`
	LikedAt                *time.Time                     `json:"liked_at,omitempty" sql:"liked_at"`
	AestheticArtifactScore *float32                       `json:"aesthetic_artifact_score,omitempty"`
	AestheticRatingScore   *float32                       `json:"aesthetic_rating_score,omitempty"`
	LikeCount              int                            `json:"like_count" sql:"like_count"`
	LikeCountTrending      *int                           `json:"like_count_trending,omitempty"`
	ImageUrl               string                         `json:"image_url,omitempty" sql:"image_path"`
	UpscaledImageUrl       string                         `json:"upscaled_image_url,omitempty" sql:"upscaled_image_path"`
	GalleryStatus          generationoutput.GalleryStatus `json:"gallery_status,omitempty" sql:"output_gallery_status"`
	DeletedAt              *time.Time                     `json:"deleted_at,omitempty" sql:"deleted_at"`
	GenerationQueryWithOutputsData
}

type GenerationQueryWithOutputsResultFormatted 

type GenerationQueryWithOutputsResultFormatted struct {
	GenerationUpscaleOutput
	Generation GenerationQueryWithOutputsData `json:"generation"`
}

type GenerationUpscaleOutput 

type GenerationUpscaleOutput struct {
	ID                     uuid.UUID                      `json:"id"`
	ImageUrl               string                         `json:"image_url,omitempty"`
	AudioFileUrl           string                         `json:"audio_file_url,omitempty"`
	VideoFileUrl           string                         `json:"video_file_url,omitempty"`
	AudioArray             []float64                      `json:"audio_array,omitempty"`
	AudioDuration          *float32                       `json:"audio_duration,omitempty"`
	UpscaledImageUrl       string                         `json:"upscaled_image_url,omitempty"`
	GalleryStatus          generationoutput.GalleryStatus `json:"gallery_status,omitempty"`
	OutputID               *uuid.UUID                     `json:"output_id,omitempty"`
	CreatedAt              *time.Time                     `json:"created_at,omitempty"`
	IsFavorited            bool                           `json:"is_favorited"`
	InitImageUrl           string                         `json:"init_image_url,omitempty"`
	WasAutoSubmitted       bool                           `json:"was_auto_submitted"`
	IsPublic               bool                           `json:"is_public"`
	DenoiseAudio           *bool                          `json:"denoise_audio,omitempty"`
	RemoveSilence          *bool                          `json:"remove_silence,omitempty"`
	LikeCount              *int                           `json:"like_count,omitempty"`
	IsLiked                *bool                          `json:"is_liked,omitempty"`
	AestheticRatingScore   *float32                       `json:"aesthetic_rating_score,omitempty"`
	AestheticArtifactScore *float32                       `json:"aesthetic_artifact_score,omitempty"`
}

type MatchBannedPrompts 

type MatchBannedPrompts struct {
	ID            uuid.UUID `json:"id" sql:"id"`
	Prompt        string    `json:"prompt" sql:"prompt"`
	Similarity    float32   `json:"similarity" sql:"similarity"`
	ShouldBanUser bool      `json:"should_ban_user" sql:"should_ban_user"`
}

type OperationQueryResult 

type OperationQueryResult struct {
	ID            uuid.UUID           `json:"id"`
	OperationType shared.ProcessType  `json:"operation_type"`
	CreatedAt     time.Time           `json:"created_at"`
	StartedAt     time.Time           `json:"started_at"`
	CompletedAt   time.Time           `json:"completed_at"`
	NumOutputs    int                 `json:"num_outputs"`
	Cost          int32               `json:"cost"`
	Source        enttypes.SourceType `json:"source"`
}

type OperationQueryResultMeta 

type OperationQueryResultMeta struct {
	Operations []OperationQueryResult `json:"operations"`
	Next       *time.Time             `json:"next,omitempty"`
}

type OperationType 

type OperationType string
const (
	OperationTypeGeneration OperationType = "generation"
	OperationTypeUpscale    OperationType = "upscale"
)

type PromptType 

type PromptType struct {
	ID   uuid.UUID `json:"id"`
	Text string    `json:"text"`
}

type Repository 

type Repository struct {
	DB             *ent.Client
	ConnInfo       database.SqlDBConn
	Redis          *database.RedisWrapper
	Ctx            context.Context
	Qdrant         *qdrant.QdrantClient
	QueueThrottler *shared.UserQueueThrottlerMap
}

func (*Repository) AddAdhocCreditsIfEligible 

func (r *Repository) AddAdhocCreditsIfEligible(creditType *ent.CreditType, userID uuid.UUID, lineItemID string) (added bool, err error)

Adds credits of creditType to user if they do not already have any belonging to stripe invoice line item

func (*Repository) AddCreditsIfEligible 

func (r *Repository) AddCreditsIfEligible(creditType *ent.CreditType, userID uuid.UUID, expiresAt time.Time, lineItemId string, DB *ent.Client) (added bool, err error)

Add credits of creditType to user if they do not have any un-expired credits of this type

func (*Repository) AddCreditsToUser 

func (r *Repository) AddCreditsToUser(creditType *ent.CreditType, userID uuid.UUID) error

Give credits to user

func (*Repository) AddToQueueLog 

func (r *Repository) AddToQueueLog(messageId string, priority int, DB *ent.Client) (*ent.MqLog, error)

Add to queue log

func (*Repository) ApplyUserGenerationsFilters 

func (r *Repository) ApplyUserGenerationsFilters(query *ent.GenerationQuery, filters *requests.QueryGenerationFilters, omitEdges bool) *ent.GenerationQuery

Apply all filters to root ent query

func (*Repository) ApplyUserVoiceoverFilters 

func (r *Repository) ApplyUserVoiceoverFilters(query *ent.VoiceoverQuery, filters *requests.QueryVoiceoverFilters, omitEdges bool) *ent.VoiceoverQuery

Apply all filters to root ent query

func (*Repository) BanDomains 

func (r *Repository) BanDomains(domains []string, deleteData bool) (int, error)

Ban domains

func (*Repository) BanUsers 

func (r *Repository) BanUsers(userIDs []uuid.UUID, deleteData bool) (int, error)

Ban users

func (*Repository) BulkUpdateGalleryStatusForOutputs 

func (r *Repository) BulkUpdateGalleryStatusForOutputs(outputIDs []uuid.UUID, galleryStatus generationoutput.GalleryStatus) (int, error)

Bulk change gallery status

func (*Repository) CheckIfEmailExists 

func (r *Repository) CheckIfEmailExists(email string) (string, bool, error)

Check if email already exists

func (*Repository) CollectUnclaimedTips 

func (r *Repository) CollectUnclaimedTips(toUserId uuid.UUID, toUserDiscordId string, DB *ent.Client) (claimted int32, err error)

Collect unclaimed tips

func (*Repository) CreateCreditType 

func (r *Repository) CreateCreditType(name string, amount int32, description *string, stripeProductID *string, ctype credittype.Type) (*ent.CreditType, error)

func (*Repository) CreateGeneration 

func (r *Repository) CreateGeneration(userID uuid.UUID, deviceType, deviceOs, deviceBrowser, countryCode string, req requests.CreateGenerationRequest, productId *string, apiTokenId *uuid.UUID, sourceType enttypes.SourceType, DB *ent.Client) (*ent.Generation, error)

CreateGeneration creates the initial generation in the database Takes in a userID (creator), device info, countryCode, and a request body

func (*Repository) CreateMockData 

func (repo *Repository) CreateMockData(ctx context.Context) error

Just creates some mock data for our tests

func (*Repository) CreateMockGenerationForDeletion 

func (repo *Repository) CreateMockGenerationForDeletion(ctx context.Context) (*ent.Generation, error)

func (*Repository) CreateMockUpscaleForDeletion 

func (repo *Repository) CreateMockUpscaleForDeletion(ctx context.Context) (*ent.Upscale, error)

func (*Repository) CreateUpscale 

func (r *Repository) CreateUpscale(userID uuid.UUID, width, height int32, deviceType, deviceOs, deviceBrowser, countryCode string, req requests.CreateUpscaleRequest, productId *string, systemGenerated bool, apiTokenId *uuid.UUID, sourceType enttypes.SourceType, DB *ent.Client) (*ent.Upscale, error)

CreateUpscale creates the initial generation in the database Takes in a userID (creator), device info, countryCode, and a request body

func (*Repository) CreateUser 

func (r *Repository) CreateUser(id uuid.UUID, email string, stripeCustomerId string, lastSignIn *time.Time, db *ent.Client) (*ent.User, error)

func (*Repository) CreateVoiceover 

func (r *Repository) CreateVoiceover(userID uuid.UUID, deviceType, deviceOs, deviceBrowser, countryCode string, req requests.CreateVoiceoverRequest, productId *string, apiTokenId *uuid.UUID, sourceType enttypes.SourceType, DB *ent.Client) (*ent.Voiceover, error)

CreateVoiceover creates the initial voiceover in the database Takes in a userID (creator), device info, countryCode, and a request body

func (*Repository) DeactivateTokenForUser 

func (r *Repository) DeactivateTokenForUser(id uuid.UUID, userId uuid.UUID) (int, error)

func (*Repository) DeductCreditsFromUser 

func (r *Repository) DeductCreditsFromUser(userID uuid.UUID, amount int32, forTip bool, DB *ent.Client) (success bool, err error)

Deduct credits from user, starting with credits that expire soonest. Return true if deduction was successful

func (*Repository) DeleteCreditsWithLineItemID 

func (r *Repository) DeleteCreditsWithLineItemID(lineItemID string) error

func (*Repository) DeleteFromQueueLog 

func (r *Repository) DeleteFromQueueLog(messageId string, DB *ent.Client) (int, error)

Delete from queue log

func (*Repository) FailCogMessageDueToTimeoutIfTimedOut 

func (r *Repository) FailCogMessageDueToTimeoutIfTimedOut(msg requests.CogWebhookMessage)

Consider a generation/upscale a failure due to timeout

func (*Repository) GetAllGenerationModels 

func (r *Repository) GetAllGenerationModels() ([]*ent.GenerationModel, error)

func (*Repository) GetAllSchedulers 

func (r *Repository) GetAllSchedulers() ([]*ent.Scheduler, error)

func (*Repository) GetAllUpscaleModels 

func (r *Repository) GetAllUpscaleModels() ([]*ent.UpscaleModel, error)

func (*Repository) GetAllVoiceoverModels 

func (r *Repository) GetAllVoiceoverModels() ([]*ent.VoiceoverModel, error)

func (*Repository) GetAllVoiceoverSpeakers 

func (r *Repository) GetAllVoiceoverSpeakers() ([]*ent.VoiceoverSpeaker, error)

func (*Repository) GetAvgGenerationQueueTime 

func (r *Repository) GetAvgGenerationQueueTime(since time.Time, limit int) (float64, error)

Get avg(started_at - created_at) over a time period with limit

func (*Repository) GetBannedUsersToDelete 

func (r *Repository) GetBannedUsersToDelete() ([]*ent.User, error)

Get banned users to delete

func (*Repository) GetCreditTypeByID 

func (r *Repository) GetCreditTypeByID(id uuid.UUID) (*ent.CreditType, error)

func (*Repository) GetCreditTypeByStripeProductID 

func (r *Repository) GetCreditTypeByStripeProductID(stripeProductID string) (*ent.CreditType, error)

func (*Repository) GetCreditTypeList 

func (r *Repository) GetCreditTypeList() ([]*ent.CreditType, error)

func (*Repository) GetCreditsForUser 

func (r *Repository) GetCreditsForUser(userID uuid.UUID) ([]*UserCreditsQueryResult, error)

Get credits for user that are not expired

func (*Repository) GetFreeCreditReplenishesAtForUser 

func (r *Repository) GetFreeCreditReplenishesAtForUser(userID uuid.UUID) (*time.Time, *ent.Credit, *ent.CreditType, error)

func (*Repository) GetFreeUsers 

func (r *Repository) GetFreeUsers() ([]*ent.User, error)

For credit replenishment

func (*Repository) GetGeneration 

func (r *Repository) GetGeneration(id uuid.UUID) (*ent.Generation, error)

Get generation by ID

func (*Repository) GetGenerationCount 

func (r *Repository) GetGenerationCount(filters *requests.QueryGenerationFilters) (int, error)

Gets the count of generations with outputs user has with filters

func (*Repository) GetGenerationCountAdmin 

func (r *Repository) GetGenerationCountAdmin(filters *requests.QueryGenerationFilters) (int, error)

func (*Repository) GetGenerationOutput 

func (r *Repository) GetGenerationOutput(id uuid.UUID) (*ent.GenerationOutput, error)

Get generation output by ID

func (*Repository) GetGenerationOutputForUser 

func (r *Repository) GetGenerationOutputForUser(id uuid.UUID, userID uuid.UUID) (*ent.GenerationOutput, error)

Get generation output for user

func (*Repository) GetGenerationOutputIDsLikedByUser 

func (r *Repository) GetGenerationOutputIDsLikedByUser(userID uuid.UUID, limit int) ([]uuid.UUID, error)

Take user ID and limit, return array of output_ids they have liked in descending order of created_at

func (*Repository) GetGenerationOutputLikeCountForUser 

func (r *Repository) GetGenerationOutputLikeCountForUser(userID uuid.UUID) (int, error)

Get how many likes a user has received for their generation outputs

func (*Repository) GetGenerationOutputLikeCountForUserExcludingSelfLikes 

func (r *Repository) GetGenerationOutputLikeCountForUserExcludingSelfLikes(userID uuid.UUID) (int, error)

Get how many likes a user has received for their generation outputs excluding self

func (*Repository) GetGenerationOutputWidthHeight 

func (r *Repository) GetGenerationOutputWidthHeight(outputID uuid.UUID) (width, height int32, err error)

Get width/height for generation output

func (*Repository) GetGenerationOutputsLikedByUser 

func (r *Repository) GetGenerationOutputsLikedByUser(userID uuid.UUID, generationOutputIDs []uuid.UUID) (map[uuid.UUID]struct{}, error)

Take user id and generation output ID array, return a map containing uuids for faster lookup.

func (*Repository) GetGenerations 

func (r *Repository) GetGenerations(limit int) ([]*ent.Generation, error)

Get last N generations, a basic view

func (*Repository) GetGenerationsByIDList 

func (r *Repository) GetGenerationsByIDList(ids []uuid.UUID) ([]*ent.Generation, error)

Get generations by IDs

func (*Repository) GetGenerationsQueuedOrStarted 

func (r *Repository) GetGenerationsQueuedOrStarted() ([]*ent.Generation, error)

func (*Repository) GetNSubscribers 

func (r *Repository) GetNSubscribers() (int, error)

Get N subscribers

func (*Repository) GetNonExpiredCreditTotalForUser 

func (r *Repository) GetNonExpiredCreditTotalForUser(userID uuid.UUID, DB *ent.Client) (int, error)

func (*Repository) GetNonFreeCreditSum 

func (r *Repository) GetNonFreeCreditSum(userID uuid.UUID) (int, error)

Determine if a user has non-free credits or not

func (*Repository) GetNonUpscaledGalleryItems 

func (r *Repository) GetNonUpscaledGalleryItems(limit int) ([]*ent.GenerationOutput, error)

func (*Repository) GetOrCreateDeviceInfo 

func (r *Repository) GetOrCreateDeviceInfo(deviceType, deviceOs, deviceBrowser string, DB *ent.Client) (deviceInfoId uuid.UUID, err error)

Get a device_info ID given inputs

func (*Repository) GetOrCreateFreeCreditType 

func (r *Repository) GetOrCreateFreeCreditType(DB *ent.Client) (*ent.CreditType, error)

func (*Repository) GetOrCreatePrompts 

func (r *Repository) GetOrCreatePrompts(promptText, negativePromptText string, promptType prompt.Type, DB *ent.Client) (promptId *uuid.UUID, negativePromptId *uuid.UUID, err error)

Create device info and prompts if they don't exist, otherwise get existing

func (*Repository) GetOrCreateRefundCreditType 

func (r *Repository) GetOrCreateRefundCreditType(db *ent.Client) (*ent.CreditType, error)

Refund credit types are used to refund credits in case of generation failure, they do not expire

func (*Repository) GetOrCreateTippableCreditType 

func (r *Repository) GetOrCreateTippableCreditType(DB *ent.Client) (*ent.CreditType, error)

Tippable credit types can't be spent, but can be gifted to other users in social channels

func (*Repository) GetOrCreateTippedCreditType 

func (r *Repository) GetOrCreateTippedCreditType(DB *ent.Client) (*ent.CreditType, error)

Tipped credit types are used to track credits that have been gifted to a user

func (*Repository) GetPaidCreditSum 

func (r *Repository) GetPaidCreditSum(userID uuid.UUID) (int, error)

Determine if user has paid credits, with a non null stripe line item ID

func (*Repository) GetPromptFromOutputID 

func (r *Repository) GetPromptFromOutputID(outputID uuid.UUID) (string, error)

func (*Repository) GetPromptTextById 

func (r *Repository) GetPromptTextById(id uuid.UUID) (text string, err error)

Get prompt text by prompt ID

func (*Repository) GetPublicGenerationOutput 

func (r *Repository) GetPublicGenerationOutput(id uuid.UUID) (*ent.GenerationOutput, error)

Get public generation output by ID Get generation output by ID

func (*Repository) GetQueuePosition 

func (r *Repository) GetQueuePosition(messageId string) (position int, total int, err error)

func (*Repository) GetQueuedItems 

func (r *Repository) GetQueuedItems(DB *ent.Client) ([]*responses.QueuedItem, error)

Get entire queuee ordered by created_at and priority

func (*Repository) GetRoles 

func (r *Repository) GetRoles(userID uuid.UUID) ([]string, error)

func (*Repository) GetSuperAdminUserIDs 

func (r *Repository) GetSuperAdminUserIDs() ([]uuid.UUID, error)

func (*Repository) GetTippableSumForUser 

func (r *Repository) GetTippableSumForUser(userID uuid.UUID) (int, error)

func (*Repository) GetToken 

func (r *Repository) GetToken(id uuid.UUID) (*ent.ApiToken, error)

func (*Repository) GetTokenByHashedToken 

func (r *Repository) GetTokenByHashedToken(hashedToken string) (*ent.ApiToken, error)

func (*Repository) GetTokenCountByUserID 

func (r *Repository) GetTokenCountByUserID(userID uuid.UUID) (int, error)

func (*Repository) GetTokensByUserID 

func (r *Repository) GetTokensByUserID(userID uuid.UUID, filters *requests.ApiTokenQueryFilters) ([]*ent.ApiToken, error)

func (*Repository) GetUpscale 

func (r *Repository) GetUpscale(id uuid.UUID) (*ent.Upscale, error)

Get upscale by ID

func (*Repository) GetUpscaleOutputWithPath 

func (r *Repository) GetUpscaleOutputWithPath(path string) (*ent.UpscaleOutput, error)

func (*Repository) GetUpscalesQueuedOrStarted 

func (r *Repository) GetUpscalesQueuedOrStarted() ([]*ent.Upscale, error)

func (*Repository) GetUser 

func (r *Repository) GetUser(id uuid.UUID) (*ent.User, error)

func (*Repository) GetUserByDiscordID 

func (r *Repository) GetUserByDiscordID(discordId string) (*ent.User, error)

Get user by discord ID

func (*Repository) GetUserByEmail 

func (r *Repository) GetUserByEmail(email string) (*ent.User, error)

Get user by email

func (*Repository) GetUserByStripeCustomerId 

func (r *Repository) GetUserByStripeCustomerId(customerId string) (*ent.User, error)

func (*Repository) GetUserByUsername 

func (r *Repository) GetUserByUsername(username string) (*ent.User, error)

Get user by username

func (*Repository) GetUserGenerationOutputs 

func (r *Repository) GetUserGenerationOutputs(userId uuid.UUID) ([]*ent.GenerationOutput, error)

func (*Repository) GetUserIDsByUsernames 

func (r *Repository) GetUserIDsByUsernames(usernames []string) ([]uuid.UUID, error)

Get user ids by usernames

func (*Repository) GetUserWithRoles 

func (r *Repository) GetUserWithRoles(id uuid.UUID) (*ent.User, error)

func (*Repository) GetUsersToDelete 

func (r *Repository) GetUsersToDelete() ([]*ent.User, error)

Get non-banned users to delete

func (*Repository) GetUsersUniqueNegativePromptIds 

func (r *Repository) GetUsersUniqueNegativePromptIds(negativePromptIds []uuid.UUID, userId uuid.UUID) ([]uuid.UUID, error)

Returns UUIDs of prompts that are unique to this user

func (*Repository) GetUsersUniquePromptIds 

func (r *Repository) GetUsersUniquePromptIds(promptIds []uuid.UUID, userId uuid.UUID) ([]uuid.UUID, error)

Returns UUIDs of prompts that are unique to this user

func (*Repository) GetVoiceover 

func (r *Repository) GetVoiceover(id uuid.UUID) (*ent.Voiceover, error)

func (*Repository) GetVoiceoverCount 

func (r *Repository) GetVoiceoverCount(filters *requests.QueryVoiceoverFilters) (int, error)

Gets the count of voiceovers with outputs user has with filters

func (*Repository) GetVoiceoversQueuedOrStarted 

func (r *Repository) GetVoiceoversQueuedOrStarted() ([]*ent.Voiceover, error)

func (*Repository) GetVoiceverSpeakersWithName 

func (r *Repository) GetVoiceverSpeakersWithName(limit int) ([]*ent.VoiceoverSpeaker, error)

func (*Repository) GiveFreeCredits 

func (r *Repository) GiveFreeCredits(userID uuid.UUID, DB *ent.Client) (added bool, err error)

Give free credits if eligible

func (*Repository) GiveFreeTippableCredits 

func (r *Repository) GiveFreeTippableCredits(userID uuid.UUID, DB *ent.Client) (added bool, err error)

Give free credits if eligible

func (*Repository) IsBanned 

func (r *Repository) IsBanned(userId uuid.UUID) (bool, error)

Get is banned

func (*Repository) IsBannedPromptEmbedding 

func (r *Repository) IsBannedPromptEmbedding(embedding []float32, DB *ent.Client) ([]MatchBannedPrompts, error)

func (*Repository) IsSuperAdmin 

func (r *Repository) IsSuperAdmin(userID uuid.UUID) (bool, error)

func (*Repository) MakeGenerationOutputsPrivateForUser 

func (r *Repository) MakeGenerationOutputsPrivateForUser(outputIDs []uuid.UUID, userID uuid.UUID, isSubscriber bool) (int, error)

Sets is_public to false on the output IDs given by the user If gallery status is not rejected or approved, sets gallery status to not_submitted Only applies for outputs where was_auto_submitted is false

func (*Repository) MakeGenerationOutputsPublicForUser 

func (r *Repository) MakeGenerationOutputsPublicForUser(outputIDs []uuid.UUID, userID uuid.UUID) (int, error)

Sets is_public to true on the output IDs given by the user If gallery status is not rejected or approved, sets gallery status to submitted

func (*Repository) MarkGenerationOutputsForDeletion 

func (r *Repository) MarkGenerationOutputsForDeletion(generationOutputIDs []uuid.UUID) (int, error)

Marks generations for deletions by setting deleted_at

func (*Repository) MarkGenerationOutputsForDeletionForUser 

func (r *Repository) MarkGenerationOutputsForDeletionForUser(generationOutputIDs []uuid.UUID, userID uuid.UUID) (int, error)

Marks generations for deletions by setting deleted_at, only if they belong to the user with ID userID

func (*Repository) MarkUpscaleOutputForDeletion 

func (r *Repository) MarkUpscaleOutputForDeletion(upscaleOutputIDs []uuid.UUID) (int, error)

Marks upscales for deletions by setting deleted_at

func (*Repository) MarkUpscaleOutputForDeletionBasedOnGenerationOutputIDs 

func (r *Repository) MarkUpscaleOutputForDeletionBasedOnGenerationOutputIDs(generationOutputIDs []uuid.UUID, DB *ent.Client) (int, error)

Marks upscales for deletion if they are associated with the given generation outputs

func (*Repository) MarkUpscaleOutputForDeletionForUser 

func (r *Repository) MarkUpscaleOutputForDeletionForUser(upscaleOutputIDs []uuid.UUID, userID uuid.UUID) (int, error)

Marks upscales for deletions by setting deleted_at, only if they belong to the user with ID userID

func (*Repository) MarkVoiceoverOutputsForDeletion 

func (r *Repository) MarkVoiceoverOutputsForDeletion(outputIDs []uuid.UUID) (int, error)

Marks voiceovers for deletions by setting deleted_at

func (*Repository) MarkVoiceoverOutputsForDeletionForUser 

func (r *Repository) MarkVoiceoverOutputsForDeletionForUser(outputIDs []uuid.UUID, userID uuid.UUID) (int, error)

Marks voiceovers for deletions by setting deleted_at, only if they belong to the user with ID userID

func (*Repository) NewAPIToken 

func (r *Repository) NewAPIToken(userId uuid.UUID, name string) (dbToken *ent.ApiToken, token string, err error)

func (*Repository) NewAPITokenForAuthClient 

func (r *Repository) NewAPITokenForAuthClient(userId uuid.UUID, authClient *ent.AuthClient) (dbToken *ent.ApiToken, token string, err error)

func (*Repository) ProcessCogMessage 

func (r *Repository) ProcessCogMessage(msg requests.CogWebhookMessage) error

Process a cog message into database

func (*Repository) QueryGenerations 

func (r *Repository) QueryGenerations(per_page int, cursor *time.Time, filters *requests.QueryGenerationFilters) (*GenerationQueryWithOutputsMeta[*time.Time], error)

Get user generations from the database using page options Cursor actually represents created_at, we paginate using this for performance reasons If present, we will get results after the cursor (anything before, represents previous pages) ! using ent .With... doesn't use joins, so we construct our own query to make it more efficient

func (*Repository) QueryGenerationsAdmin 

func (r *Repository) QueryGenerationsAdmin(per_page int, cursor *time.Time, callingUserId *uuid.UUID, filters *requests.QueryGenerationFilters) (*GenerationQueryWithOutputsMeta[*time.Time], error)

Alternate version for performance when we can't index by user_id

func (*Repository) QueryUserOperations 

func (r *Repository) QueryUserOperations(userId uuid.UUID, limit int, cursor *time.Time) (*OperationQueryResultMeta, error)

Query generations and upscales combined for user

func (*Repository) QueryUsers 

func (r *Repository) QueryUsers(
	emailSearch string,
	per_page int,
	cursor *time.Time,
	productIds []string,
	banned *bool,
) (*UserQueryMeta, error)

Query all users with filters per_page is how many rows to return cursor is created_at on users, will return items with created_at less than cursor

func (*Repository) QueryUsersCount 

func (r *Repository) QueryUsersCount(emailSearch string) (totalCount int, totalCountByProduct map[string]int, err error)

Get count for QueryUsers

func (*Repository) QueryVoiceovers 

func (r *Repository) QueryVoiceovers(per_page int, cursor *time.Time, filters *requests.QueryVoiceoverFilters) (*VoiceoverQueryWithOutputsMeta, error)

Get user voiceovers from the database using page options Cursor actually represents created_at, we paginate using this for performance reasons If present, we will get results after the cursor (anything before, represents previous pages) ! using ent .With... doesn't use joins, so we construct our own query to make it more efficient

func (*Repository) RefundCreditsToUser 

func (r *Repository) RefundCreditsToUser(userID uuid.UUID, amount int32, db *ent.Client) (success bool, err error)

Refund credits for user, starting with credits that expire soonest. Return true if refund was successful

func (*Repository) ReplenishFreeCreditsToEligibleUsers 

func (r *Repository) ReplenishFreeCreditsToEligibleUsers() (int, error)

Replenish free credits where eligible

func (*Repository) RetrieveGalleryDataByID 

func (r *Repository) RetrieveGalleryDataByID(id uuid.UUID, userId *uuid.UUID, callingUserId *uuid.UUID, all bool) (*GalleryData, error)

Retrieved a single generation output by ID, in GalleryData format

func (*Repository) RetrieveGalleryDataWithOutputIDs 

func (r *Repository) RetrieveGalleryDataWithOutputIDs(outputIDs []uuid.UUID, callingUserId *uuid.UUID, allIsPublic bool) ([]GalleryData, error)

Retrieves data in gallery format given output IDs

func (*Repository) RetrieveGenerationsWithOutputIDs 

func (r *Repository) RetrieveGenerationsWithOutputIDs(outputIDs []uuid.UUID, callingUserId *uuid.UUID, admin bool) (*GenerationQueryWithOutputsMeta[*uint], error)

Get user generations in the *GenerationQueryWithOutputsMeta format Using a list of generation_output ids For when we semantic search against our vector db

func (*Repository) RetrieveMostRecentGalleryData 

func (r *Repository) RetrieveMostRecentGalleryData(filters *requests.QueryGenerationFilters, callingUserId *uuid.UUID, per_page int, cursor *time.Time) ([]GalleryData, *time.Time, error)

Retrieves data in gallery format given output IDs Returns data, next cursor, error

func (*Repository) RetrieveMostRecentGalleryDataV2 

func (r *Repository) RetrieveMostRecentGalleryDataV2(filters *requests.QueryGenerationFilters, callingUserId *uuid.UUID, per_page int, cursor *time.Time, offset *int) ([]GalleryData, *time.Time, *int, error)

func (*Repository) SetActiveProductID 

func (r *Repository) SetActiveProductID(id uuid.UUID, stripeProductID string, db *ent.Client) error

func (*Repository) SetDiscordID 

func (r *Repository) SetDiscordID(userId uuid.UUID, discordId string, DB *ent.Client) error

Set discord ID on user

func (*Repository) SetFavoriteGenerationOutputsForUser 

func (r *Repository) SetFavoriteGenerationOutputsForUser(generationOutputIDs []uuid.UUID, userID uuid.UUID, action requests.FavoriteAction) (int, error)

Marks generations for deletions by setting deleted_at, only if they belong to the user with ID userID

func (*Repository) SetGenerationFailed 

func (r *Repository) SetGenerationFailed(generationID string, reason string, nsfwCount int32, db *ent.Client) error

func (*Repository) SetGenerationStarted 

func (r *Repository) SetGenerationStarted(generationID string) error

func (*Repository) SetGenerationSucceeded 

func (r *Repository) SetGenerationSucceeded(generationID string, promptStr string, negativePrompt string, submitToGallery bool, whOutput requests.CogWebhookOutput, nsfwCount int32) ([]*ent.GenerationOutput, error)

func (*Repository) SetIsProcessingInQueueLog 

func (r *Repository) SetIsProcessingInQueueLog(messageId string, isProcessing bool, DB *ent.Client) (int, error)

Set is_processing

func (*Repository) SetOutputsLikedForUser 

func (r *Repository) SetOutputsLikedForUser(generationOutputIDs []uuid.UUID, userID uuid.UUID, action requests.LikeUnlikeAction) error

Inserts like

func (*Repository) SetTokenUsedAndIncrementCreditsSpent 

func (r *Repository) SetTokenUsedAndIncrementCreditsSpent(creditsSpent int, tokenId uuid.UUID) error

func (*Repository) SetUpscaleFailed 

func (r *Repository) SetUpscaleFailed(upscaleID string, reason string, db *ent.Client) error

func (*Repository) SetUpscaleStarted 

func (r *Repository) SetUpscaleStarted(upscaleID string) error

func (*Repository) SetUpscaleSucceeded 

func (r *Repository) SetUpscaleSucceeded(upscaleID, generationOutputID, inputImageUrl string, output requests.CogWebhookOutput) (*ent.UpscaleOutput, error)

! Currently supports 1 output

func (*Repository) SetUsername 

func (r *Repository) SetUsername(userId uuid.UUID, username string) error

func (*Repository) SetVoiceoverFailed 

func (r *Repository) SetVoiceoverFailed(voiceoverId string, reason string, db *ent.Client) error

func (*Repository) SetVoiceoverStarted 

func (r *Repository) SetVoiceoverStarted(voiceoverID string) error

func (*Repository) SetVoiceoverSucceeded 

func (r *Repository) SetVoiceoverSucceeded(voiceoverId, promptStr string, submitToGallery bool, output requests.CogWebhookOutput) (*ent.VoiceoverOutput, error)

! Currently supports 1 output

func (*Repository) SetWantsEmail 

func (r *Repository) SetWantsEmail(userId uuid.UUID, wantsEmail bool) error

Set wants email

func (*Repository) SubmitGenerationOutputsToGalleryForUser 

func (r *Repository) SubmitGenerationOutputsToGalleryForUser(outputIDs []uuid.UUID, userID uuid.UUID) (int, error)

! Deprecated, see makeGenerationOutputsPublicForUser Submits all generation outputs to gallery for review, from user Verifies that all outputs belong to the user Only submits generation outputs that are not already submitted, accepted, or rejected Returns count of how many were submitted

func (*Repository) SyncStripeProductIDs 

func (r *Repository) SyncStripeProductIDs(productCustomerIDMap map[string][]string) error

Sync stripe product IDs

func (*Repository) TipCreditsToUser 

func (r *Repository) TipCreditsToUser(fromUser uuid.UUID, toUser *uuid.UUID, toUserDiscordId string, amount int32) (success bool, err error)

Handle tips from one user to another

func (*Repository) UnbanDomains 

func (r *Repository) UnbanDomains(domains []string) (int, error)

Ban domains

func (*Repository) UnbanUsers 

func (r *Repository) UnbanUsers(userIDs []uuid.UUID) (int, error)

Unban users

func (*Repository) UnsetActiveProductID 

func (r *Repository) UnsetActiveProductID(id uuid.UUID, stripeProductId string, db *ent.Client) (int, error)

Only unset if the active product ID matches the stripe product ID given

func (*Repository) UpdateCache 

func (r *Repository) UpdateCache() error

Update the cache from the database

func (*Repository) UpdateGenerationCountCacheAdmin 

func (r *Repository) UpdateGenerationCountCacheAdmin(filters *requests.QueryGenerationFilters) (int, error)

func (*Repository) UpdateLastSeenAt 

func (r *Repository) UpdateLastSeenAt(id uuid.UUID) error

Update last_seen_at

func (*Repository) WithTx 

func (r *Repository) WithTx(fn func(tx *ent.Tx) error) error

WithTx runs a function in a transaction Usage example: 

if err := r.WithTx(func(tx *ent.Tx) error {
	 Do stuff with tx
	return nil
}); err != nil {

	 Handle error
}

type TaskStatusUpdateResponse 

type TaskStatusUpdateResponse struct {
	MessageType      string                    `json:"message_type"`
	Status           requests.CogTaskStatus    `json:"status"`
	ProcessType      shared.ProcessType        `json:"process_type"`
	Id               string                    `json:"id"`
	UIId             string                    `json:"ui_id,omitempty"`
	StreamId         string                    `json:"stream_id"`
	Error            string                    `json:"error,omitempty"`
	NSFWCount        int32                     `json:"nsfw_count,omitempty"`
	Outputs          []GenerationUpscaleOutput `json:"outputs,omitempty"`
	ForLivePage      bool                      `json:"for_live_page,omitempty"`
	LivePageMessage  *shared.LivePageMessage   `json:"live_page_message,omitempty"`
	RemainingCredits int                       `json:"total_remaining_credits,omitempty"`
}

Represents an update to a generation/upscale in our database

type UserCreditGroupByType 

type UserCreditGroupByType struct {
	ActiveProductID string `json:"active_product_id"`
	Count           int    `json:"count"`
}

type UserCreditsQueryResult 

type UserCreditsQueryResult struct {
	ID                    uuid.UUID `json:"id" sql:"id"`
	RemainingAmount       int32     `json:"remaining_amount" sql:"remaining_amount"`
	ExpiresAt             time.Time `json:"expires_at" sql:"expires_at"`
	CreditTypeID          uuid.UUID `json:"credit_type_id" sql:"credit_type_id"`
	CreditTypeName        string    `json:"credit_type_name" sql:"credit_type_name"`
	CreditTypeDescription string    `json:"credit_type_description" sql:"credit_type_description"`
	CreditTypeAmount      int32     `json:"credit_type_amount" sql:"credit_type_amount"`
}

type UserGenCount 

type UserGenCount struct {
	Total int `json:"total" sql:"total"`
}

type UserQueryCreditType 

type UserQueryCreditType struct {
	ID              uuid.UUID `json:"id"`
	Name            string    `json:"name"`
	StripeProductId string    `json:"stripe_product_id,omitempty"`
}

type UserQueryCredits 

type UserQueryCredits struct {
	RemainingAmount int32               `json:"remaining_amount"`
	ExpiresAt       time.Time           `json:"expires_at"`
	CreditType      UserQueryCreditType `json:"credit_type"`
	ReplenishedAt   time.Time           `json:"replenished_at,omitempty"`
}

type UserQueryMeta 

type UserQueryMeta struct {
	Total            *int              `json:"total_count,omitempty"`
	TotalByProductID map[string]int    `json:"total_count_by_product_id,omitempty"`
	Next             *time.Time        `json:"next,omitempty"`
	Users            []UserQueryResult `json:"users"`
}

Paginated meta for querying generations

type UserQueryResult 

type UserQueryResult struct {
	ID                     uuid.UUID          `json:"id"`
	Email                  string             `json:"email"`
	Username               string             `json:"username"`
	StripeCustomerID       string             `json:"stripe_customer_id"`
	Roles                  []string           `json:"role,omitempty"`
	CreatedAt              time.Time          `json:"created_at"`
	Credits                []UserQueryCredits `json:"credits,omitempty"`
	LastSignInAt           *time.Time         `json:"last_sign_in_at,omitempty"`
	LastSeenAt             time.Time          `json:"last_seen_at"`
	BannedAt               *time.Time         `json:"banned_at,omitempty"`
	DataDeletedAt          *time.Time         `json:"data_deleted_at,omitempty"`
	ScheduledForDeletionOn *time.Time         `json:"scheduled_for_deletion_on,omitempty"`
	StripeProductID        *string            `json:"product_id,omitempty"`
}

type UserType 

type UserType struct {
	Email      string `json:"email,omitempty"`
	Username   string `json:"username"`
	Identifier string `json:"identifier,omitempty"`
}

type VoiceoverQueryWithOutputsData 

type VoiceoverQueryWithOutputsData struct {
	ID               uuid.UUID                 `json:"id" sql:"id"`
	Seed             int                       `json:"seed" sql:"seed"`
	Temperature      float32                   `json:"temperature" sql:"temperature"`
	Status           string                    `json:"status" sql:"status"`
	ModelID          uuid.UUID                 `json:"model_id" sql:"model_id"`
	PromptID         *uuid.UUID                `json:"prompt_id,omitempty" sql:"prompt_id"`
	CreatedAt        time.Time                 `json:"created_at" sql:"created_at"`
	UpdatedAt        time.Time                 `json:"updated_at" sql:"updated_at"`
	StartedAt        *time.Time                `json:"started_at,omitempty" sql:"started_at"`
	CompletedAt      *time.Time                `json:"completed_at,omitempty" sql:"completed_at"`
	PromptText       string                    `json:"prompt_text,omitempty" sql:"prompt_text"`
	IsFavorited      bool                      `json:"is_favorited" sql:"is_favorited"`
	Outputs          []GenerationUpscaleOutput `json:"outputs"`
	Prompt           PromptType                `json:"prompt"`
	WasAutoSubmitted bool                      `json:"was_auto_submitted" sql:"was_auto_submitted"`
	DenoiseAudio     bool                      `json:"denoise_audio" sql:"denoise_audio"`
	RemoveSilence    bool                      `json:"remove_silence" sql:"remove_silence"`
	Speaker          *VoiceoverSpeaker         `json:"speaker,omitempty"`
	// For speaker object
	SpeakerID    *uuid.UUID `json:"speaker_id,omitempty" sql:"speaker_id"`
	NameInWorker string     `json:"name_in_worker,omitempty" sql:"name_in_worker"`
	Locale       string     `json:"locale,omitempty" sql:"locale"`
}

type VoiceoverQueryWithOutputsMeta 

type VoiceoverQueryWithOutputsMeta struct {
	Total   *int                                       `json:"total_count,omitempty"`
	Outputs []VoiceoverQueryWithOutputsResultFormatted `json:"outputs"`
	Next    *time.Time                                 `json:"next,omitempty"`
}

Paginated meta for querying voiceovers

type VoiceoverQueryWithOutputsResult 

type VoiceoverQueryWithOutputsResult struct {
	OutputID      *uuid.UUID `json:"output_id,omitempty" sql:"output_id"`
	AudioFileUrl  string     `json:"audio_file_url,omitempty" sql:"audio_path"`
	VideoFileUrl  string     `json:"video_file_url,omitempty" sql:"video_path"`
	AudioArray    []float64  `json:"audio_array,omitempty" sql:"audio_array"`
	AudioDuration float32    `json:"audio_duration" sql:"audio_duration"`
	DeletedAt     *time.Time `json:"deleted_at,omitempty" sql:"deleted_at"`
	VoiceoverQueryWithOutputsData
}

type VoiceoverQueryWithOutputsResultFormatted 

type VoiceoverQueryWithOutputsResultFormatted struct {
	GenerationUpscaleOutput
	Voiceover VoiceoverQueryWithOutputsData `json:"voiceover"`
}

type VoiceoverSpeaker 

type VoiceoverSpeaker struct {
	ID     uuid.UUID `json:"id"`
	Name   string    `json:"name"`
	Locale string    `json:"locale"`
}

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.