Update app.py

app.py

@@ -1,24 +1,30 @@
 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 base64
+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
-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
+import base64
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
+# For voice transcription
+import speech_recognition as sr
 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)
@@ -28,31 +34,39 @@ 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)
-# (Repeat login and groq_client setup if needed)
+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"
 
-# --- 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
 
-# --- History Functions ---
+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 -----
 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):
-        image = Image.fromarray(image.astype("uint8")) if image.dtype != np.uint8 else Image.fromarray(image)
+        if image.dtype == np.uint8:
+            image = Image.fromarray(image)
+        else:
+            image = Image.fromarray((image * 255).astype(np.uint8))
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
-    history.append({
+    return history + [{
          "image": img_bytes,
          "prompt": prompt,
          "seed": seed,
@@ -60,12 +74,9 @@ 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()
@@ -82,13 +93,14 @@ def create_history_html(history):
               </div>
          </div>
          """
-    return html + "</div>"
+    return html + "</div>" if history else "<p style='margin: 20px;'>No generations yet</p>"
 
-# --- Inference & Prompt Enhancement Functions ---
+@spaces.GPU(duration=75)
 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,
@@ -108,10 +120,16 @@ 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, 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?
+                         """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 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?
+
 Environment: Indoor, outdoor, abstract, etc.
 Colors: How do they contrast or harmonize with the subject?
 Lighting: Time of day, intensity, direction (e.g., backlighting).
@@ -124,13 +142,25 @@ 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
+# --- Voice Transcription Function ---
+def transcribe_audio(audio_file):
+    r = sr.Recognizer()
+    try:
+        with sr.AudioFile(audio_file) as source:
+            audio_data = r.record(source)
+        text = r.recognize_google(audio_data)
+    except Exception as e:
+        text = f"Transcription error: {str(e)}"
+    return text
 
-# --- Gradio Interface with Enhanced UI ---
+# --- Gradio Interface with Enhanced UI and Voice Recognition ---
 custom_css = """
 #col-container {
     margin: 0 auto;
@@ -141,9 +171,9 @@ custom_css = """
 .bot-msg { background: #f5f5f5; border-radius: 15px; padding: 10px; margin: 5px; }
 """
 
-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.Blocks(css=custom_css, title="FLUX.1 [dev] Enhanced UI with Voice Recognition") as demo:
+    gr.Markdown("# FLUX.1 [dev] with Enhanced UI and Voice Recognition")
+    # Using Tabs to separate functionalities
     with gr.Tabs():
         with gr.Tab("Prompt Enhancement"):
             gr.Markdown("### Step 1: Enhance Your Prompt")
@@ -152,6 +182,13 @@ with gr.Blocks(css=custom_css, title="FLUX.1 [dev] Enhanced UI") as demo:
             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("Voice Recognition"):
+            gr.Markdown("### Step 1A: Record Your Prompt")
+            audio_input = gr.Audio(source="microphone", type="filepath", label="Record your prompt")
+            transcribe_button = gr.Button("Transcribe Audio", variant="secondary")
+            voice_text = gr.Textbox(label="Transcribed Prompt", placeholder="Your spoken prompt will appear here...", lines=3)
+            transcribe_button.click(transcribe_audio, audio_input, voice_text)
+
         with gr.Tab("Generate Image"):
             gr.Markdown("### Step 2: Generate Image")
             with gr.Row():
@@ -171,8 +208,6 @@ with gr.Blocks(css=custom_css, title="FLUX.1 [dev] Enhanced UI") as demo:
                 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],
@@ -187,11 +222,17 @@ with gr.Blocks(css=custom_css, title="FLUX.1 [dev] Enhanced UI") as demo:
          inputs=history_state,
          outputs=history_display
     )
-    # 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>",
+    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,
          inputs=history_state,
          outputs=history_display
     )
-
-demo.launch(share=True)
+demo.launch(share=True)