Update app.py

app.py

@@ -13,14 +13,20 @@ import multiprocessing
 import io
 import time
 from tqdm import tqdm
+from google.cloud import storage
+import json
 
 hf_token = os.getenv("HF_TOKEN")
 redis_host = os.getenv("REDIS_HOST")
 redis_port = int(os.getenv("REDIS_PORT", 6379))
 redis_password = os.getenv("REDIS_PASSWORD")
+gcs_credentials = json.loads(os.getenv("GCS_CREDENTIALS"))
+gcs_bucket_name = os.getenv("GCS_BUCKET_NAME")
 
 HfFolder.save_token(hf_token)
 
+storage_client = storage.Client.from_service_account_info(gcs_credentials)
+
 def connect_to_redis():
     while True:
         try:
@@ -57,6 +63,16 @@ def save_object_to_redis(key, obj):
     except redis.exceptions.RedisError as e:
         print(f"Failed to save object to Redis: {e}")
 
+def upload_to_gcs(bucket_name, blob_name, data):
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(blob_name)
+    blob.upload_from_string(data)
+
+def download_from_gcs(bucket_name, blob_name):
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(blob_name)
+    return blob.download_as_bytes()
+
 def get_model_or_download(model_id, redis_key, loader_func):
     model = load_object_from_redis(redis_key)
     if model:
@@ -66,6 +82,8 @@ def get_model_or_download(model_id, redis_key, loader_func):
             model = loader_func(model_id, torch_dtype=torch.float16)
             pbar.update(1)
         save_object_to_redis(redis_key, model)
+        model_bytes = pickle.dumps(model)
+        upload_to_gcs(gcs_bucket_name, redis_key, model_bytes)
     except Exception as e:
         print(f"Failed to load or save model: {e}")
         return None
@@ -82,6 +100,7 @@ def generate_image(prompt):
             image.save(buffered, format="JPEG")
             image_bytes = buffered.getvalue()
             save_object_to_redis(redis_key, image_bytes)
+            upload_to_gcs(gcs_bucket_name, redis_key, image_bytes)
         except Exception as e:
             print(f"Failed to generate image: {e}")
             return None
@@ -100,6 +119,7 @@ def edit_image_with_prompt(image_bytes, prompt, strength=0.75):
             edited_image.save(buffered, format="JPEG")
             edited_image_bytes = buffered.getvalue()
             save_object_to_redis(redis_key, edited_image_bytes)
+            upload_to_gcs(gcs_bucket_name, redis_key, edited_image_bytes)
         except Exception as e:
             print(f"Failed to edit image: {e}")
             return None
@@ -115,6 +135,7 @@ def generate_song(prompt, duration=10):
                 pbar.update(1)
             song_bytes = song[0].getvalue()
             save_object_to_redis(redis_key, song_bytes)
+            upload_to_gcs(gcs_bucket_name, redis_key, song_bytes)
         except Exception as e:
             print(f"Failed to generate song: {e}")
             return None
@@ -129,6 +150,7 @@ def generate_text(prompt):
                 text = text_gen_pipeline(prompt, max_new_tokens=256)[0]["generated_text"].strip()
                 pbar.update(1)
             save_object_to_redis(redis_key, text)
+            upload_to_gcs(gcs_bucket_name, redis_key, text.encode())
         except Exception as e:
             print(f"Failed to generate text: {e}")
             return None
@@ -152,6 +174,7 @@ def generate_flux_image(prompt):
             flux_image.save(buffered, format="JPEG")
             flux_image_bytes = buffered.getvalue()
             save_object_to_redis(redis_key, flux_image_bytes)
+            upload_to_gcs(gcs_bucket_name, redis_key, flux_image_bytes)
         except Exception as e:
             print(f"Failed to generate flux image: {e}")
             return None
@@ -168,6 +191,7 @@ def generate_code(prompt):
                 code = starcoder_tokenizer.decode(outputs[0])
                 pbar.update(1)
             save_object_to_redis(redis_key, code)
+            upload_to_gcs(gcs_bucket_name, redis_key, code.encode())
         except Exception as e:
             print(f"Failed to generate code: {e}")
             return None
@@ -185,6 +209,7 @@ def generate_video(prompt):
                 video = export_to_video(pipe(prompt, num_inference_steps=25).frames)
                 pbar.update(1)
             save_object_to_redis(redis_key, video)
+            upload_to_gcs(gcs_bucket_name, redis_key, video.encode())
         except Exception as e:
             print(f"Failed to generate video: {e}")
             return None
@@ -203,6 +228,7 @@ def test_model_meta_llama():
                 response = meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"].strip()
                 pbar.update(1)
             save_object_to_redis(redis_key, response)
+            upload_to_gcs(gcs_bucket_name, redis_key, response.encode())
         except Exception as e:
             print(f"Failed to test Meta-Llama: {e}")
             return None
@@ -223,6 +249,8 @@ def train_model(model, dataset, epochs, batch_size, learning_rate):
             pbar.update(epochs)
         save_object_to_redis("trained_model", model)
         save_object_to_redis("training_results", output_dir.getvalue())
+        upload_to_gcs(gcs_bucket_name, "trained_model", pickle.dumps(model))
+        upload_to_gcs(gcs_bucket_name, "training_results", output_dir.getvalue())
     except Exception as e:
         print(f"Failed to train model: {e}")