alfrds commited on
Commit
9dbb565
·
1 Parent(s): 7703dd9

adding new changes for generation and permisions

Browse files
Files changed (1) hide show
  1. generate.py +15 -3
generate.py CHANGED
@@ -21,9 +21,13 @@ pipe = StableDiffusionPipeline.from_pretrained(
21
 
22
  pipe = pipe.to(device)
23
 
24
- s3 = boto3.resource('s3')
 
 
25
 
26
- s3_client = boto3.client('s3')
 
 
27
 
28
 
29
  def text_to_image(prompt, save_as, key_id):
@@ -46,13 +50,21 @@ def text_to_image(prompt, save_as, key_id):
46
 
47
  # Full path of the file in the bucket
48
  s3_key = "public/" + image_name
 
49
 
50
  # Upload the image to the S3 bucket
51
  s3.upload_fileobj(image_buffer, S3_BUCKET_NAME, s3_key)
 
52
 
53
  def generator_image(prompt):
54
  prompt = prompt
55
- image = pipe(prompt).images[0]
 
 
 
 
 
 
56
 
57
  # Save the image in S3
58
  save_image_to_s3(image)
 
21
 
22
  pipe = pipe.to(device)
23
 
24
+ s3 = boto3.resource('s3',
25
+ aws_access_key_id=AWS_ACCESS_KEY_ID,
26
+ aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
27
 
28
+ s3_client = boto3.client('s3',
29
+ aws_access_key_id=AWS_ACCESS_KEY_ID,
30
+ aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
31
 
32
 
33
  def text_to_image(prompt, save_as, key_id):
 
50
 
51
  # Full path of the file in the bucket
52
  s3_key = "public/" + image_name
53
+ print('Saving image to s3')
54
 
55
  # Upload the image to the S3 bucket
56
  s3.upload_fileobj(image_buffer, S3_BUCKET_NAME, s3_key)
57
+ print('Image saved to s3')
58
 
59
  def generator_image(prompt):
60
  prompt = prompt
61
+ print('Starting to generate the image ...')
62
+ try:
63
+ image = pipe(prompt).images[0]
64
+ except Exception as e:
65
+ print('Error: ', e)
66
+
67
+ print('Image generation completed')
68
 
69
  # Save the image in S3
70
  save_image_to_s3(image)