Update app.py

app.py

@@ -1,30 +1,24 @@
 import gradio as gr
 import torch
-from groq import Groq
-from cryptography.fernet import Fernet
-from huggingface_hub import login
 import os
-os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '120'
 import numpy as np
 import random
-import spaces
-from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL, UNet2DConditionModel
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast, T5Tokenizer, T5EncoderModel
-from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
-from io import BytesIO
 import base64
-from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
-import torch.nn.functional as F
-import torch.nn as nn
+from io import BytesIO
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from cryptography.fernet import Fernet
+from huggingface_hub import login
+from groq import Groq
+from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+
+# Environment setup and device configuration
 os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '120'
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-
+# --- Token & API Setup ---
 def get_hf_token(encrypted_token):
     key = "K4FlQbffvTcDxT2FIhrOPV1eue6ia45FFR3kqp2hHbM="
-    if not key:
-        raise ValueError("Missing decryption key! Set the DECRYPTION_KEY environment variable.")
     if isinstance(key, str):
         key = key.encode()
     f = Fernet(key)
@@ -34,39 +28,31 @@ def get_hf_token(encrypted_token):
 groq_client = Groq(api_key="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")
 decrypted_token = get_hf_token("gAAAAABn3GfShExoJd50nau3B5ZJNiQ9dRD1ACO3XXMwVaIQMkmi59cL-MKGr6SYnsB0E2gGITJG2j29Ar9yjaZP-EC6hHsCBmwKSj4aFtTor9_n0_NdMBv1GtlxZRmwnQwriB-Xr94e")
 login(token=decrypted_token)
-decrypted_token = get_hf_token("gAAAAABn3GfShExoJd50nau3B5ZJNiQ9dRD1ACO3XXMwVaIQMkmi59cL-MKGr6SYnsB0E2gGITJG2j29Ar9yjaZP-EC6hHsCBmwKSj4aFtTor9_n0_NdMBv1GtlxZRmwnQwriB-Xr94e")
-login(token=decrypted_token)
-groq_client = Groq(api_key="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")
-
-
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# (Repeat login and groq_client setup if needed)
 
+# --- Load Models ---
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
+pipe.to(device)
 torch.cuda.empty_cache()
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
+# --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
-
-# ----- HISTORY FUNCTIONS & GRADIO INTERFACE -----
+# --- History Functions ---
 def append_to_history(image, prompt, seed, width, height, guidance_scale, steps, history):
     if image is None:
         return history
     from PIL import Image
-    import numpy as np
     if isinstance(image, np.ndarray):
-        if image.dtype == np.uint8:
-            image = Image.fromarray(image)
-        else:
-            image = Image.fromarray((image * 255).astype(np.uint8))
+        image = Image.fromarray(image.astype("uint8")) if image.dtype != np.uint8 else Image.fromarray(image)
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
-    return history + [{
+    history.append({
          "image": img_bytes,
          "prompt": prompt,
          "seed": seed,
@@ -74,9 +60,12 @@ def append_to_history(image, prompt, seed, width, height, guidance_scale, steps,
          "height": height,
          "guidance_scale": guidance_scale,
          "steps": steps,
-    }]
+    })
+    return history
 
 def create_history_html(history):
+    if not history:
+        return "<p style='margin: 20px;'>No generations yet</p>"
     html = "<div style='display: flex; flex-direction: column; gap: 20px; margin: 20px;'>"
     for i, entry in enumerate(reversed(history)):
          img_str = base64.b64encode(entry["image"]).decode()
@@ -93,14 +82,13 @@ def create_history_html(history):
               </div>
          </div>
          """
-    return html + "</div>" if history else "<p style='margin: 20px;'>No generations yet</p>"
+    return html + "</div>"
 
-@spaces.GPU(duration=75)
+# --- Inference & Prompt Enhancement Functions ---
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
             prompt=prompt,
             guidance_scale=guidance_scale,
@@ -120,16 +108,10 @@ def enhance_prompt(user_prompt):
                  {
                      "role": "system",
                      "content": (
-                         """Enhance user input into prompts that paint a clear picture for image generation. Be precise, detailed and direct, describe not only the content of the image but also such details as tone, style, color palette, and point of view, for photorealistic images, include the name of the device used (e.g., “shot on iPhone 16”), aperture, lens, and shot type. Use precise, visual descriptions (rather than metaphorical concepts). 
-Try to keep prompts to contain only keywords, yet precise, and awe-inspiring. 
-Medium:
-Consider what form of art this image should be simulating.
-
--Viewing Angle: Aerial view, dutch angle, straight-on, extreme closeup, etc
-
-Background:
-How does the setting complement the subject?
-
+                         """Enhance user input into prompts that paint a clear picture for image generation. Be precise, detailed and direct, describing content, tone, style, color palette, and point of view. Use precise, visual descriptions with keywords for photorealistic images.
+                         
+Viewing Angle: Aerial view, dutch angle, straight-on, extreme closeup, etc.
+Background: How does the setting complement the subject?
 Environment: Indoor, outdoor, abstract, etc.
 Colors: How do they contrast or harmonize with the subject?
 Lighting: Time of day, intensity, direction (e.g., backlighting).
@@ -142,56 +124,55 @@ Lighting: Time of day, intensity, direction (e.g., backlighting).
              temperature=0.5,
              max_completion_tokens=1024,
              top_p=1,
-             stop=None,
-             stream=False,
          )
          enhanced = chat_completion.choices[0].message.content
     except Exception as e:
          enhanced = f"Error enhancing prompt: {str(e)}"
     return enhanced
 
-css = """
+# --- Gradio Interface with Enhanced UI ---
+custom_css = """
 #col-container {
     margin: 0 auto;
-    max-width: 520px;
+    max-width: 600px;
+    padding: 20px;
 }
+.user-msg { background: #e3f2fd; border-radius: 15px; padding: 10px; margin: 5px; }
+.bot-msg { background: #f5f5f5; border-radius: 15px; padding: 10px; margin: 5px; }
 """
 
-with gr.Blocks(css=css) as demo:
-    history_state = gr.State([])
-    with gr.Column(elem_id="col-container"):
-         gr.Markdown("# FLUX.1 [dev] with History Tracking")
-         gr.Markdown("### Step 1: Enhance Your Prompt")
-         original_prompt = gr.Textbox(label="Original Prompt", lines=2)
-         enhance_button = gr.Button("Enhance Prompt")
-         enhanced_prompt = gr.Textbox(label="Enhanced Prompt (Editable)", lines=2)
-         enhance_button.click(enhance_prompt, original_prompt, enhanced_prompt)
-         gr.Markdown("### Step 2: Generate Image")
-         with gr.Row():
-              run_button = gr.Button("Generate Image", variant="primary")
-         result = gr.Image(label="Result", show_label=False)
-         with gr.Accordion("Advanced Settings"):
-              seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
-              randomize_seed = gr.Checkbox(True, label="Randomize seed")
-              with gr.Row():
-                   width = gr.Slider(256, MAX_IMAGE_SIZE, 1024, step=32, label="Width")
-                   height = gr.Slider(256, MAX_IMAGE_SIZE, 1024, step=32, label="Height")
-              with gr.Row():
-                   guidance_scale = gr.Slider(1, 15, 3.5, step=0.1, label="Guidance Scale")
-                   num_inference_steps = gr.Slider(1, 50, 28, step=1, label="Inference Steps")
-         with gr.Accordion("Generation History", open=False):
-              history_display = gr.HTML("<p style='margin: 20px;'>No generations yet</p>")
-         gr.Examples(
-              examples=[
-                   "a tiny astronaut hatching from an egg on the moon",
-                   "a cat holding a sign that says hello world",
-                   "an anime illustration of a wiener schnitzel",
-              ],
-              inputs=enhanced_prompt,
-              outputs=[result, seed],
-              fn=infer,
-              cache_examples="lazy"
-         )
+with gr.Blocks(css=custom_css, title="FLUX.1 [dev] Enhanced UI") as demo:
+    gr.Markdown("# FLUX.1 [dev] with Enhanced UI")
+    # Using Tabs to separate prompt enhancement and image generation
+    with gr.Tabs():
+        with gr.Tab("Prompt Enhancement"):
+            gr.Markdown("### Step 1: Enhance Your Prompt")
+            original_prompt = gr.Textbox(label="Original Prompt", placeholder="Enter your creative idea here...", lines=3)
+            enhance_button = gr.Button("Enhance Prompt", variant="secondary")
+            enhanced_prompt = gr.Textbox(label="Enhanced Prompt (Editable)", placeholder="Enhanced prompt appears here...", lines=3)
+            enhance_button.click(enhance_prompt, original_prompt, enhanced_prompt)
+
+        with gr.Tab("Generate Image"):
+            gr.Markdown("### Step 2: Generate Image")
+            with gr.Row():
+                run_button = gr.Button("Generate Image", variant="primary")
+                clear_history_button = gr.Button("Clear History", variant="secondary")
+            result = gr.Image(label="Generated Image", show_label=False)
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(0, MAX_SEED, value=0, label="Seed", info="Seed for reproducibility")
+                randomize_seed = gr.Checkbox(True, label="Randomize Seed")
+                with gr.Row():
+                    width = gr.Slider(256, MAX_IMAGE_SIZE, 1024, step=32, label="Width")
+                    height = gr.Slider(256, MAX_IMAGE_SIZE, 1024, step=32, label="Height")
+                with gr.Row():
+                    guidance_scale = gr.Slider(1, 15, 3.5, step=0.1, label="Guidance Scale")
+                    num_inference_steps = gr.Slider(1, 50, 28, step=1, label="Inference Steps")
+            with gr.Accordion("Generation History", open=False):
+                history_display = gr.HTML("<p style='margin: 20px;'>No generations yet</p>")
+            # State to track generation history
+            history_state = gr.State([])
+
+    # --- Define interactions ---
     generation_event = run_button.click(
          fn=infer,
          inputs=[enhanced_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
@@ -206,17 +187,11 @@ with gr.Blocks(css=css) as demo:
          inputs=history_state,
          outputs=history_display
     )
-    enhanced_prompt.submit(
-         fn=infer,
-         inputs=[enhanced_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-         outputs=[result, seed]
-    ).then(
-         fn=append_to_history,
-         inputs=[result, enhanced_prompt, seed, width, height, guidance_scale, num_inference_steps, history_state],
-         outputs=history_state
-    ).then(
-         fn=create_history_html,
+    # Clear history action
+    clear_history_button.click(fn=lambda: [], inputs=[], outputs=history_state).then(
+         fn=lambda hist: "<p style='margin: 20px;'>No generations yet</p>",
          inputs=history_state,
          outputs=history_display
     )
+
 demo.launch(share=True)