update origin

app.py

@@ -1,63 +1,41 @@
-import os
+import gradio as gr
 import boto3
-import io
 from PIL import Image
-import requests
-import gradio as gr
+import io
+
+# Define your AWS S3 setup here
+s3_client = boto3.client(
+    's3',
+    aws_access_key_id='AKIA34AOT6TNS2T5EZOB',
+    aws_secret_access_key='aCXXKADXHcOgAIjtp3Nv9lK6uCZpm9D5CcQjWhb0',
+    region_name='ap-southeast-1'
+)
 
-# AWS S3 configuration
 bucket_name = 'flux-image-generation'
 
-# Retrieve AWS credentials from environment variables
-aws_access_key_id = os.getenv('AKIA34AOT6TNS2T5EZOB')
-aws_secret_access_key = os.getenv('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')
+def upload_to_s3(image, bucket_name, s3_key):
+    s3_client.upload_fileobj(image, bucket_name, s3_key)
 
-# Initialize S3 client
-s3 = boto3.client(
-    's3', 
-    aws_access_key_id=aws_access_key_id, 
-    aws_secret_access_key=aws_secret_access_key
-)
+def process_and_upload_image(image):
+    # Convert image to bytes
+    image_bytes = io.BytesIO()
+    image.save(image_bytes, format='PNG')
+    image_bytes.seek(0)
 
-# Load the Gradio model (replace with the actual Gradio loading function if needed)
-model = gr.load("models/black-forest-labs/FLUX.1-dev")
-
-def generate_and_store_image(prompt):
-    try:
-        # Generate an image URL from the prompt using the FLUX model
-        # Assuming the model returns a URL to the generated image
-        image_url = model(prompt)
-        
-        # Download the image from the URL
-        response = requests.get(image_url)
-        response.raise_for_status()  # Check for HTTP errors
-        
-        # Convert downloaded image to PIL Image
-        img = Image.open(io.BytesIO(response.content))
-        
-        # Save image to in-memory file
-        in_mem_file = io.BytesIO()
-        img.save(in_mem_file, format='PNG')
-        in_mem_file.seek(0)
-        
-        # Upload image to S3
-        image_name = f'{prompt.replace(" ", "_")}.png'  # Unique image name based on prompt
-        s3.upload_fileobj(in_mem_file, bucket_name, image_name)
-        
-        # Return S3 URL of the image
-        s3_url = f'https://{bucket_name}.s3.amazonaws.com/{image_name}'
-        return s3_url
-    except Exception as e:
-        return f"Error generating or uploading image: {e}"
-
-# Set up Gradio interface using the latest API
-iface = gr.Interface(
-    fn=generate_and_store_image,
-    inputs=gr.Textbox(lines=2, placeholder="Enter a prompt for the FLUX model"),
-    outputs=gr.Textbox(),  # Change to `gr.Textbox()` for text output
-    title="Generate and Store Image with FLUX Model",
-    description="Enter a prompt, and the image will be generated and saved to S3."
-)
+    # Generate a unique key for the S3 object
+    s3_key = f'generated_images/{str(uuid.uuid4())}.png'
+
+    # Upload the image to S3
+    upload_to_s3(image_bytes, bucket_name, s3_key)
+
+    return s3_key  # Return the S3 key or URL for reference
+
+# Define your Gradio interface
+def generate_image():
+    # Your code to generate an image
+    image = ...  # This should be the generated image
+    s3_key = process_and_upload_image(image)
+    return s3_key
 
-# Launch the Gradio interface
+iface = gr.Interface(fn=generate_image, inputs=[], outputs="text")
 iface.launch()