Update app.py

app.py

@@ -31,81 +31,39 @@ plastic, cartoonish, artificial, fake, unnatural
 app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
 
 app.layout = dbc.Container([
-    html.H1("Insta-Image", className="my-4"),
-    dbc.Card([
-        dbc.CardBody([
-            dbc.Input(id="google-api-key", type="password", placeholder="Enter Google AI API Key", className="mb-3"),
-            dbc.Input(id="stability-api-key", type="password", placeholder="Enter Stability AI API Key", className="mb-3"),
-            dbc.Textarea(id="prompt", placeholder="Enter your prompt", className="mb-3"),
-            dcc.Dropdown(id="style", options=[{"label": s, "value": s} for s in STYLES], placeholder="Select style", className="mb-3"),
-            dbc.Accordion([
-                dbc.AccordionItem(
-                    [
-                        dbc.Row([
-                            dbc.Col([
-                                html.Label("CFG Scale:", className="mr-2"),
-                                dcc.Slider(id="cfg-scale", min=1, max=30, step=0.5, value=7, 
-                                           marks={1: '1', 15: '15', 30: '30'}, 
-                                           tooltip={"placement": "bottom", "always_visible": True}),
-                                dbc.Tooltip(
-                                    "Controls the influence of the prompt. Higher values adhere more closely to the prompt.",
-                                    target="cfg-scale",
-                                ),
-                            ], width=12, className="mb-3"),
-                        ]),
-                        dbc.Row([
-                            dbc.Col([
-                                html.Label("Steps:", className="mr-2"),
-                                dcc.Slider(id="steps", min=4, max=50, step=1, value=20, 
-                                           marks={4: '4', 25: '25', 50: '50'}, 
-                                           tooltip={"placement": "bottom", "always_visible": True}),
-                                dbc.Tooltip(
-                                    "Number of denoising steps. More steps can lead to higher quality but longer generation time.",
-                                    target="steps",
-                                ),
-                            ], width=12, className="mb-3"),
-                        ]),
-                        dbc.Row([
-                            dbc.Col([
-                                html.Label("Sampler:", className="mr-2"),
-                                dcc.Dropdown(
-                                    id="sampler",
-                                    options=[
-                                        {"label": "DDIM", "value": "DDIM"},
-                                        {"label": "PLMS", "value": "PLMS"},
-                                        {"label": "K_EULER", "value": "K_EULER"},
-                                        {"label": "K_EULER_ANCESTRAL", "value": "K_EULER_ANCESTRAL"},
-                                        {"label": "DPM_2", "value": "DPM_2"},
-                                        {"label": "DPM_2_ANCESTRAL", "value": "DPM_2_ANCESTRAL"},
-                                    ],
-                                    value="K_EULER_ANCESTRAL",
-                                ),
-                                dbc.Tooltip(
-                                    "The algorithm used for image generation. Different samplers can produce varying results.",
-                                    target="sampler",
-                                ),
-                            ], width=12, className="mb-3"),
-                        ]),
-                    ],
-                    title="Advanced Settings",
-                )
-            ], start_collapsed=True, className="mb-3"),
-            dbc.Button("Generate Image", id="submit-btn", color="primary", className="mb-3"),
-        ])
-    ], className="mb-4"),
-    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"),
-                    html.Div(id="enhanced-prompt-output", className="mt-3"),
-                ]
-            ),
-        ])
-    ])
+    html.H1("Stability AI SD3.5 Large Turbo Image Generator with Google Gemini Prompt Enhancement", className="my-4"),
+    dbc.Row([
+        # Left column: Form entry
+        dbc.Col([
+            dbc.Card([
+                dbc.CardBody([
+                    dbc.Input(id="google-api-key", type="password", placeholder="Enter Google AI API Key", className="mb-3"),
+                    dbc.Input(id="stability-api-key", type="password", placeholder="Enter Stability AI API Key", className="mb-3"),
+                    dbc.Textarea(id="prompt", placeholder="Enter your prompt", className="mb-3"),
+                    dcc.Dropdown(id="style", options=[{"label": s, "value": s} for s in STYLES], placeholder="Select style", className="mb-3"),
+                    dbc.Button("Generate Image", id="submit-btn", color="primary", 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)
 
 def enhance_prompt(google_api_key, prompt, style):
@@ -117,11 +75,11 @@ def enhance_prompt(google_api_key, prompt, style):
     Original prompt: '{prompt}'
     
     Instructions:
-    1. Expand the prompt to be more detailed, vivid, and realism and always include the right camera used for the shot
+    1. Expand the prompt to be more detailed, vivid, and photorealistic.
     2. Incorporate elements of the specified style, focusing on realism and natural appearances.
     3. Add details that enhance the realism of the scene, especially for elements like trees, textures, and lighting.
     4. Emphasize natural lighting and enhance the realism of textures and colors.
-    5. Avoid terms that might result in artificial or cartoonish appearances unless specifically requested
+    5. Avoid terms that might result in artificial or cartoonish appearances.
     6. Maintain the original intent of the prompt while significantly improving its descriptive quality.
     7. Provide ONLY the enhanced prompt, without any explanations or options.
     8. Keep the enhanced prompt concise, ideally under 100 words.
@@ -145,7 +103,7 @@ def enhance_prompt(google_api_key, prompt, style):
         logging.error(f"Error in enhance_prompt: {str(e)}")
         raise
 
-def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, cfg_scale, steps, sampler):
+def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
     url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
     
     headers = {
@@ -161,9 +119,8 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, c
         "width": 1024,
         "height": 1024,
         "num_images": 1,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "sampler": sampler,
+        "steps": 20,
+        "cfg_scale": 7,
     }
     
     try:
@@ -184,13 +141,13 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt, c
         logging.error(f"Request failed: {str(e)}")
         raise Exception(f"Request failed: {str(e)}")
 
-def process_and_generate(google_api_key, stability_api_key, prompt, style, cfg_scale, steps, sampler, set_status):
+def process_and_generate(google_api_key, stability_api_key, prompt, style, set_status):
     try:
         set_status("Enhancing prompt...")
         enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
         
         set_status("Generating image...")
-        image_bytes = generate_image(stability_api_key, enhanced_prompt, style, DEFAULT_NEGATIVE_PROMPT, cfg_scale, steps, sampler)
+        image_bytes = generate_image(stability_api_key, enhanced_prompt, style, DEFAULT_NEGATIVE_PROMPT)
         
         set_status("Image generated successfully!")
         return image_bytes, enhanced_prompt
@@ -202,18 +159,16 @@ def process_and_generate(google_api_key, stability_api_key, prompt, style, cfg_s
 @app.callback(
     [Output("image-output", "src"),
      Output("enhanced-prompt-output", "children"),
-     Output("status-message", "children")],
+     Output("status-message", "children"),
+     Output("download-btn", "disabled")],
     [Input("submit-btn", "n_clicks")],
     [State("google-api-key", "value"),
      State("stability-api-key", "value"),
      State("prompt", "value"),
-     State("style", "value"),
-     State("cfg-scale", "value"),
-     State("steps", "value"),
-     State("sampler", "value")],
+     State("style", "value")],
     prevent_initial_call=True
 )
-def update_output(n_clicks, google_api_key, stability_api_key, prompt, style, cfg_scale, steps, sampler):
+def update_output(n_clicks, google_api_key, stability_api_key, prompt, style):
     if n_clicks is None:
         raise PreventUpdate
     
@@ -225,12 +180,12 @@ def update_output(n_clicks, google_api_key, stability_api_key, prompt, style, cf
         status["message"] = message
     
     def run_process():
-        image_bytes, enhanced_prompt = process_and_generate(google_api_key, stability_api_key, prompt, style, cfg_scale, steps, sampler, set_status)
+        image_bytes, enhanced_prompt = process_and_generate(google_api_key, stability_api_key, prompt, style, 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"]
+            return f"data:image/jpeg;base64,{encoded_image}", f"Enhanced Prompt: {enhanced_prompt}", status["message"], False
         else:
-            return "", f"Error: {enhanced_prompt}", status["message"]
+            return "", f"Error: {enhanced_prompt}", status["message"], True
     
     # Run the process in a separate thread
     thread = threading.Thread(target=run_process)
@@ -239,6 +194,24 @@ def update_output(n_clicks, google_api_key, stability_api_key, prompt, style, cf
     
     return run_process()
 
+@app.callback(
+    Output("download-image", "data"),
+    Input("download-btn", "n_clicks"),
+    State("image-output", "src"),
+    prevent_initial_call=True
+)
+def download_image(n_clicks, image_src):
+    if n_clicks is None:
+        raise PreventUpdate
+    
+    # Extract the base64 encoded image data
+    image_data = image_src.split(",")[1]
+    
+    # Decode the base64 data
+    image_bytes = base64.b64decode(image_data)
+    
+    return dcc.send_bytes(image_bytes, "generated_image.jpeg")
+
 if __name__ == '__main__':
     print("Starting the Dash application...")
     app.run(debug=True, host='0.0.0.0', port=7860)