Update app.py via AI Editor

app.py

@@ -1,6 +1,6 @@
 import base64
 import dash
-from dash import dcc, html, Input, Output, State
+from dash import dcc, html, Input, Output, State, callback_context
 import dash_bootstrap_components as dbc
 from dash.exceptions import PreventUpdate
 import google.generativeai as genai
@@ -9,18 +9,18 @@ import logging
 import threading
 import time
 import os
+import flask
+import uuid
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Updated STYLES list
 STYLES = [
     "photographic", "3d-model", "analog-film", "anime", "cinematic", "comic-book", 
     "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly", 
     "modeling-compound", "neon-punk", "origami", "pixel-art", "tile-texture"
 ]
 
-# Default negative prompt (hidden from UI)
 DEFAULT_NEGATIVE_PROMPT = """
 ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, 
 extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, 
@@ -28,65 +28,43 @@ cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, di
 plastic, cartoonish, artificial, fake, unnatural, blurry, smooth, lack of detail, low quality
 """
 
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+external_stylesheets = [dbc.themes.BOOTSTRAP]
+server = flask.Flask(__name__)
+app = dash.Dash(__name__, server=server, external_stylesheets=external_stylesheets)
 
-app.layout = dbc.Container([
-    html.H1("ImaGen", className="my-4"),
-    dbc.Row([
-        # Left column: Form entry
-        dbc.Col([
-            dbc.Card([
-                dbc.CardBody([
-                    dbc.Textarea(id="prompt", placeholder="Enter your prompt", className="mb-3"),
-                    dcc.Dropdown(
-                        id="style",
-                        options=[{"label": s.replace("-", " ").title(), "value": s} for s in STYLES],
-                        value="photographic",
-                        placeholder="Select style",
-                        className="mb-3"
-                    ),
-                    dbc.Button("Generate Image", id="submit-btn", color="primary", className="mb-3"),
-                    dbc.Accordion([
-                        dbc.AccordionItem(
-                            [
-                                dbc.Label("Aspect Ratio"),
-                                dcc.Dropdown(
-                                    id="aspect-ratio",
-                                    options=[
-                                        {"label": ar, "value": ar} for ar in 
-                                        ["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"]
-                                    ],
-                                    value="1:1"
-                                ),
-                                dbc.Label("Steps"),
-                                dcc.Slider(id="steps", min=4, max=50, step=1, value=30, marks={4: '4', 25: '25', 50: '50'}),
-                            ],
-                            title="Advanced Settings",
-                        ),
-                    ], start_collapsed=True, className="mb-3"),
-                ])
-            ], className="mb-4"),
-        ], width=6),
-        # Right column: Image preview
-        dbc.Col([
-            dbc.Card([
-                dbc.CardBody([
-                    dcc.Loading(
-                        id="loading",
-                        type="circle",
-                        children=[
-                            html.Div(id="status-message", className="mb-3"),
-                            html.Img(id="image-output", className="img-fluid mb-3"),
-                            html.Div(id="enhanced-prompt-output", className="mb-3"),
-                            dbc.Button("Download Image", id="download-btn", color="secondary", className="mb-3", disabled=True),
-                            dcc.Download(id="download-image")
-                        ]
-                    ),
-                ])
-            ]),
-        ], width=6),
-    ]),
-], fluid=True)
+app.title = "ImaGen"
+
+# Global in-memory session storage and locks
+SESSION_DATA = {}
+SESSION_LOCKS = {}
+
+def get_session_id():
+    if hasattr(flask.g, "session_id"):
+        return flask.g.session_id
+    session_id = flask.request.cookies.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+    flask.g.session_id = session_id
+    return session_id
+
+@app.server.before_request
+def ensure_session_id():
+    session_id = flask.request.cookies.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        flask.g.set_cookie = session_id
+    flask.g.session_id = session_id or flask.g.get("set_cookie", None)
+    # Ensure session state and lock
+    if session_id not in SESSION_DATA:
+        SESSION_DATA[session_id] = {'image': None, 'enhanced_prompt': None, 'status': None}
+    if session_id not in SESSION_LOCKS:
+        SESSION_LOCKS[session_id] = threading.Lock()
+
+@app.server.after_request
+def set_session_cookie(response):
+    if hasattr(flask.g, "set_cookie"):
+        response.set_cookie("session_id", flask.g.set_cookie, httponly=True, samesite='Lax')
+    return response
 
 def enhance_prompt(google_api_key, prompt, style):
     genai.configure(api_key=google_api_key)
@@ -108,17 +86,13 @@ def enhance_prompt(google_api_key, prompt, style):
 
     Enhanced prompt:
     """
-    
     try:
         response = model.generate_content(enhanced_prompt_request)
-        
         enhanced_prompt = response.text.strip()
-        
         prefixes_to_remove = ["Enhanced prompt:", "Here's the enhanced prompt:", "The enhanced prompt is:"]
         for prefix in prefixes_to_remove:
             if enhanced_prompt.lower().startswith(prefix.lower()):
                 enhanced_prompt = enhanced_prompt[len(prefix):].strip()
-        
         logging.info(f"Enhanced prompt: {enhanced_prompt}")
         return enhanced_prompt
     except Exception as e:
@@ -127,12 +101,10 @@ def enhance_prompt(google_api_key, prompt, style):
 
 def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, steps, aspect_ratio):
     url = "https://api.stability.ai/v2beta/stable-image/generate/core"
-    
     headers = {
         "Accept": "image/*",
         "Authorization": f"Bearer {stability_api_key}"
     }
-    
     data = {
         "prompt": f"{enhanced_prompt}, Style: {style}, highly detailed, high quality, descriptive, sharp focus, intricate details",
         "negative_prompt": negative_prompt,
@@ -143,14 +115,9 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, s
         "style_preset": style,
         "aspect_ratio": aspect_ratio,
     }
-    
     try:
         response = requests.post(url, headers=headers, files={"none": ''}, data=data, timeout=60)
         response.raise_for_status()
-        
-        logging.debug(f"Response headers: {response.headers}")
-        logging.debug(f"Response content type: {response.headers.get('content-type')}")
-        
         if response.headers.get('content-type').startswith('image/'):
             image_data = response.content
             if len(image_data) < 1000:
@@ -160,7 +127,6 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, s
             error_message = response.text
             logging.error(f"Unexpected content type: {response.headers.get('content-type')}. Response: {error_message}")
             raise Exception(f"Unexpected content type: {response.headers.get('content-type')}. Response: {error_message}")
-    
     except requests.exceptions.RequestException as e:
         logging.error(f"Request failed: {str(e)}")
         raise Exception(f"Request failed: {str(e)}")
@@ -169,7 +135,6 @@ def process_and_generate(google_api_key, stability_api_key, prompt, style, steps
     try:
         set_status("Enhancing prompt...")
         enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
-        
         set_status("Generating image...")
         max_attempts = 3
         for attempt in range(max_attempts):
@@ -188,72 +153,161 @@ def process_and_generate(google_api_key, stability_api_key, prompt, style, steps
         set_status(f"Error: {str(e)}")
         return None, str(e)
 
+app.layout = dbc.Container([
+    dbc.Row([
+        dbc.Col([
+            html.H1("ImaGen", className="text-center mb-4")
+        ], width=12)
+    ]),
+    dbc.Row([
+        dbc.Col([
+            dbc.Card([
+                dbc.CardBody([
+                    dbc.Textarea(
+                        id="prompt", 
+                        placeholder="Tell me what image you want me to build.", 
+                        className="mb-3",
+                        style={"height": "120px", "whiteSpace": "pre-wrap", "wordWrap": "break-word"}
+                    ),
+                    dcc.Dropdown(
+                        id="style",
+                        options=[{"label": s.replace("-", " ").title(), "value": s} for s in STYLES],
+                        value="photographic",
+                        placeholder="Select style",
+                        className="mb-3"
+                    ),
+                    # Advanced settings always open, below dropdown
+                    dbc.FormGroup([
+                        dbc.Label("Aspect Ratio"),
+                        dcc.Dropdown(
+                            id="aspect-ratio",
+                            options=[
+                                {"label": ar, "value": ar} for ar in 
+                                ["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"]
+                            ],
+                            value="1:1",
+                            className="mb-3"
+                        ),
+                        dbc.Label("Steps"),
+                        dcc.Slider(
+                            id="steps", 
+                            min=4, 
+                            max=50, 
+                            step=1, 
+                            value=30, 
+                            marks={4: '4', 25: '25', 50: '50'},
+                            className="mb-3"
+                        ),
+                    ]),
+                    dbc.Button("Generate Image", id="submit-btn", color="primary", className="mt-2 mb-2", style={"width": "100%"})
+                ])
+            ])
+        ], width=4, style={"minWidth": "300px", "maxWidth": "420px", "flex": "0 0 30%"}),
+        dbc.Col([
+            dbc.Card([
+                dbc.CardBody([
+                    dbc.Button("Download Image", id="download-btn", color="secondary", className="mb-3", disabled=True, style={"width": "100%"}),
+                    dcc.Loading(
+                        id="loading",
+                        type="default",
+                        children=[
+                            html.Div(id="status-message", className="mb-3"),
+                            html.Img(id="image-output", className="img-fluid mb-3"),
+                            html.Div(id="enhanced-prompt-output", className="mb-3"),
+                            dcc.Download(id="download-image")
+                        ],
+                        style={"display": "block", "margin": "auto"}
+                    ),
+                ])
+            ])
+        ], width=8, style={"flex": "0 0 70%"})
+    ], align="start")
+], fluid=True)
+
 @app.callback(
-    [Output("image-output", "src"),
-     Output("enhanced-prompt-output", "children"),
-     Output("status-message", "children"),
-     Output("download-btn", "disabled")],
+    [
+        Output("image-output", "src"),
+        Output("enhanced-prompt-output", "children"),
+        Output("status-message", "children"),
+        Output("download-btn", "disabled")
+    ],
     [Input("submit-btn", "n_clicks")],
-    [State("prompt", "value"),
-     State("style", "value"),
-     State("steps", "value"),
-     State("aspect-ratio", "value")],
+    [
+        State("prompt", "value"),
+        State("style", "value"),
+        State("steps", "value"),
+        State("aspect-ratio", "value")
+    ],
     prevent_initial_call=True
 )
 def update_output(n_clicks, prompt, style, steps, aspect_ratio):
+    ctx = callback_context
     if n_clicks is None:
         raise PreventUpdate
-    
+    session_id = flask.request.cookies.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+    lock = SESSION_LOCKS.setdefault(session_id, threading.Lock())
+    session_data = SESSION_DATA.setdefault(session_id, {'image': None, 'enhanced_prompt': None, 'status': None})
     google_api_key = os.getenv('GOOGLE_API_KEY')
     stability_api_key = os.getenv('STABILITY_API_KEY')
-    
     if not google_api_key or not stability_api_key:
         return "", "Error: API keys not found in environment variables", "API keys missing", True
-    
-    logging.debug(f"Stability API Key (first 4 chars): {stability_api_key[:4]}...")
-    
     status = {"message": "Starting process..."}
-    
     def set_status(message):
         status["message"] = message
-    
+        session_data['status'] = message
     def run_process():
-        image_bytes, enhanced_prompt = process_and_generate(google_api_key, stability_api_key, prompt, style, steps, aspect_ratio, set_status)
-        if image_bytes:
-            encoded_image = base64.b64encode(image_bytes).decode('ascii')
-            return f"data:image/jpeg;base64,{encoded_image}", f"Enhanced Prompt: {enhanced_prompt}", status["message"], False
-        else:
-            return "", f"Error: {enhanced_prompt}", status["message"], True
-    
+        with lock:
+            image_bytes, enhanced_prompt = process_and_generate(google_api_key, stability_api_key, prompt, style, steps, aspect_ratio, set_status)
+            if image_bytes:
+                encoded_image = base64.b64encode(image_bytes).decode('ascii')
+                session_data['image'] = encoded_image
+                session_data['enhanced_prompt'] = enhanced_prompt
+                return f"data:image/jpeg;base64,{encoded_image}", f"Enhanced Prompt: {enhanced_prompt}", status["message"], False
+            else:
+                session_data['image'] = None
+                session_data['enhanced_prompt'] = None
+                return "", f"Error: {enhanced_prompt}", status["message"], True
     try:
         thread = threading.Thread(target=run_process)
         thread.start()
         thread.join(timeout=90)
-        
         if thread.is_alive():
+            with lock:
+                session_data['status'] = "Process timed out"
+                session_data['image'] = None
+                session_data['enhanced_prompt'] = None
             return "", "Error: Image generation timed out", "Process timed out", True
-        
         return run_process()
     except Exception as e:
         logging.error(f"Unexpected error in update_output: {str(e)}")
+        with lock:
+            session_data['status'] = "An unexpected error occurred"
+            session_data['image'] = None
+            session_data['enhanced_prompt'] = None
         return "", f"Unexpected error: {str(e)}", "An unexpected error occurred", True
 
 @app.callback(
     Output("download-image", "data"),
-    Input("download-btn", "n_clicks"),
-    State("image-output", "src"),
+    [Input("download-btn", "n_clicks")],
+    [State("image-output", "src")],
     prevent_initial_call=True
 )
 def download_image(n_clicks, image_src):
+    ctx = callback_context
     if n_clicks is None:
         raise PreventUpdate
-    
-    image_data = image_src.split(",")[1]
-    image_bytes = base64.b64decode(image_data)
-    
-    return dcc.send_bytes(image_bytes, "generated_image.jpeg")
+    session_id = flask.request.cookies.get("session_id")
+    lock = SESSION_LOCKS.setdefault(session_id, threading.Lock())
+    session_data = SESSION_DATA.setdefault(session_id, {'image': None, 'enhanced_prompt': None, 'status': None})
+    with lock:
+        if not session_data.get('image'):
+            raise PreventUpdate
+        image_bytes = base64.b64decode(session_data['image'])
+        return dcc.send_bytes(image_bytes, "generated_image.jpeg")
 
 if __name__ == '__main__':
     print("Starting the Dash application...")
-    app.run(debug=False, host='0.0.0.0', port=7860)
+    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
     print("Dash application has finished running.")