Last active
September 28, 2021 13:08
-
-
Save mfrashad/2fb66f1c25e40bc046464db736b11a33 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| import clipit | |
| import time | |
| from datetime import datetime | |
| import firebase_admin | |
| from firebase_admin import credentials, firestore, storage | |
| if not firebase_admin._apps: | |
| cred = credentials.Certificate("YOUR_CREDENTIAL_FILE") | |
| firebase_admin.initialize_app(cred, { | |
| 'storageBucket': 'YOUR_BUCKET_URL' | |
| }) | |
| db = firestore.client() | |
| bucket = storage.bucket() | |
| def generate(doc_id, prompt, quality, style, aspect, email): | |
| torch.cuda.empty_cache() | |
| clipit.reset_settings() | |
| use_pixeldraw = (style == 'pixel art') | |
| use_clipdraw = (style == 'painting') | |
| clipit.add_settings(prompts=prompt, | |
| seed=seed, | |
| aspect=aspect, | |
| quality=quality, | |
| use_pixeldraw=use_pixeldraw, | |
| use_clipdraw=use_clipdraw, | |
| make_video=True) | |
| settings = clipit.apply_settings() | |
| clipit.do_init(settings) | |
| clipit.do_run(settings) | |
| data = { | |
| "seed": seed, | |
| "prompt": prompt, | |
| "quality": quality, | |
| "aspect": aspect, | |
| "type": style, | |
| "user": email, | |
| "created_at": datetime.now() | |
| } | |
| db.collection('generated_images').document(doc_id).set(data) | |
| email_results_mailgun(email, prompt) | |
| transaction = db.transaction() | |
| @firestore.transactional | |
| def claim_task(transaction, queue_objects_ref): | |
| # query firestore | |
| queue_objects = queue_objects_ref.stream(transaction=transaction) | |
| # pull the document from the iterable | |
| next_item = None | |
| for doc in queue_objects: | |
| next_item = doc | |
| # if queue is empty return status code of 2 | |
| if not next_item: | |
| return {"status": 2} | |
| # get information from the document | |
| next_item_data = next_item.to_dict() | |
| next_item_data["status"] = 0 | |
| next_item_data['id'] = next_item.id | |
| # delete the document and return the information | |
| transaction.delete(next_item.reference) | |
| return next_item_data | |
| # initialize query | |
| queue_objects_ref = ( | |
| db.collection("queue") | |
| .order_by("created_at", direction="ASCENDING") | |
| .limit(1) | |
| ) | |
| transaction_attempts = 0 | |
| while True: | |
| try: | |
| # apply transaction | |
| next_item_data = claim_task(transaction, queue_objects_ref) | |
| if next_item_data['status'] == 0: | |
| generate(next_item_data['id'], | |
| next_item_data['prompt'], | |
| next_item_data['quality'], | |
| next_item_data['type'], | |
| next_item_data['aspect'], | |
| next_item_data['email']) | |
| print(f"Generated {next_item_data['prompt']} for {next_item_data['email']}") | |
| except Exception as e: | |
| print(f"Could not apply transaction. Error: {e}") | |
| time.sleep(5) | |
| transaction_attempts += 1 | |
| if transaction_attempts > 20: | |
| db.collection("errors").add({ | |
| "exception": f"Could not apply transaction. Error: {e}", | |
| "time": str(datetime.now()) | |
| }) | |
| exit() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment