Update media.py

media.py

@@ -1,5 +1,3 @@
-# --- START OF FILE media.py (FINAL WITH LIVE PROGRESS) ---
-
 # --- LIBRARIES ---
 import torch
 import gradio as gr
@@ -14,24 +12,21 @@ import threading
 from queue import Queue, Empty as QueueEmpty
 from PIL import Image
 
-# --- SECURE AUTHENTICATION FOR HUGGING FACE SPACES ---
-import os
-from huggingface_hub import login
+# --- DYNAMIC HARDWARE DETECTION & AUTH ---
+if torch.cuda.is_available():
+    device = "cuda"
+    torch_dtype = torch.float16
+    print("✅ GPU detected. Using CUDA.")
+else:
+    device = "cpu"
+    torch_dtype = torch.float32
+    print("⚠️ No GPU detected.")
 
-# This code will attempt to read the HF_TOKEN from the Space's secrets.
-# On your local machine, this will do nothing unless you set it up, which isn't necessary.
-# On the Hugging Face server, it will find the secret you just saved.
-HF_TOKEN = os.environ.get('HF_TOKEN')
+HF_TOKEN = os.getenv("HF_TOKEN")  # Will read the token from Space secrets
+if HF_TOKEN is None:
+    raise ValueError("❌ HF_TOKEN is not set in the environment variables!")
 
-if HF_TOKEN:
-    print("✅ Found HF_TOKEN secret. Logging in...")
-    try:
-        login(token=HF_TOKEN)
-        print("✅ Hugging Face Authentication successful.")
-    except Exception as e:
-        print(f"❌ Hugging Face login failed: {e}")
-else:
-    print("⚠️ No HF_TOKEN secret found. Gated models may not be available on the deployed app.")
+login(token=HF_TOKEN)
 
 # --- CONFIGURATION & STATE ---
 available_models = {
@@ -42,13 +37,19 @@ available_models = {
 }
 model_state = { "current_pipe": None, "loaded_model_name": None }
 
-# --- THE FINAL GENERATION FUNCTION WITH LIVE PROGRESS ---
-def generate_media_live_progress(model_key, prompt, negative_prompt, steps, cfg_scale, width, height, seed, num_frames):
-    # --- Model Loading (Unchanged) ---
+# --- THE FINAL, STABLE GENERATION FUNCTION ---
+def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_scale, width, height, seed, num_frames):
+    global model_state
+
+    # --- Model Loading ---
     if model_state.get("loaded_model_name") != model_key:
         yield {output_image: None, output_video: None, status_textbox: f"Loading {model_key}..."}
         if model_state.get("current_pipe"):
-            del model_state["current_pipe"]; gc.collect(); torch.cuda.empty_cache()
+            pipe_to_delete = model_state.pop("current_pipe", None)
+            if pipe_to_delete: del pipe_to_delete
+            gc.collect()
+            torch.cuda.empty_cache()
+        
         model_id = available_models[model_key]
         if "Video" in model_key:
             pipe = TextToVideoSDPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
@@ -70,8 +71,8 @@ def generate_media_live_progress(model_key, prompt, negative_prompt, steps, cfg_
     
     # --- Generation Logic ---
     if "Video" in model_key:
-        # For video, we'll keep the simple status updates for now
         yield {output_image: None, output_video: None, status_textbox: "Generating video..."}
+        # (Your working video code)
         video_frames = pipe(prompt=prompt, num_inference_steps=int(steps), height=320, width=576, num_frames=int(num_frames), generator=generator).frames
         video_frames_5d = np.array(video_frames)
         video_frames_4d = np.squeeze(video_frames_5d)
@@ -81,77 +82,72 @@ def generate_media_live_progress(model_key, prompt, negative_prompt, steps, cfg_
         imageio.mimsave(video_path, list_of_frames, fps=12)
         yield {output_image: None, output_video: video_path, status_textbox: f"Video saved! Seed: {seed}"}
 
-    else: # Image Generation with Live Progress
+    else: # Image Generation with your brilliant text-based progress bar
         progress_queue = Queue()
         
         def run_pipe():
-            # This function runs in a separate thread
             start_time = time.time()
             
-            def progress_callback(pipe, step, timestep, callback_kwargs):
-                # This is called by the pipeline at each step
+            # This callback correctly accepts all arguments
+            def progress_callback(step, timestep, latents, **kwargs):
                 elapsed_time = time.time() - start_time
-                # Avoid division by zero on the first step
                 if elapsed_time > 0:
                     its_per_sec = (step + 1) / elapsed_time
-                    progress_queue.put((step + 1, its_per_sec))
-                return callback_kwargs
+                    progress_queue.put(("progress", step + 1, its_per_sec))
+                return kwargs
 
             try:
-                # The final image is still generated using the pipeline's high-quality VAE
                 final_image = pipe(
                     prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=int(steps),
                     guidance_scale=float(cfg_scale), width=int(width), height=int(height),
                     generator=generator, 
                     callback_on_step_end=progress_callback
                 ).images[0]
-                progress_queue.put(final_image) # Put the final result on the queue
+                progress_queue.put(("final", final_image))
             except Exception as e:
                 print(f"An error occurred in the generation thread: {e}")
-                progress_queue.put(None) # Signal an error
+                progress_queue.put(("error", str(e)))
 
-        # Start the generation in the background
         thread = threading.Thread(target=run_pipe)
         thread.start()
 
-        # In the main thread, listen for updates from the queue and yield to Gradio
         total_steps = int(steps)
-        yield {status_textbox: "Generating..."} # Initial status
+        final_image_result = None
+        yield {status_textbox: "Generating..."}
         
         while True:
             try:
-                update = progress_queue.get(timeout=1.0) # Wait for an update
+                update_type, data = progress_queue.get(timeout=1.0)
                 
-                if isinstance(update, Image.Image): # It's the final image
-                    yield {output_image: update, status_textbox: f"Generation complete! Seed: {seed}"}
+                if update_type == "final":
+                    final_image_result = data
+                    yield {output_image: final_image_result, status_textbox: f"Generation complete! Seed: {seed}"}
                     break
-                elif isinstance(update, tuple): # It's a progress update (step, speed)
-                    current_step, its_per_sec = update
+                elif update_type == "progress":
+                    current_step, its_per_sec = data
                     progress_percent = (current_step / total_steps) * 100
                     steps_remaining = total_steps - current_step
                     eta_seconds = steps_remaining / its_per_sec if its_per_sec > 0 else 0
                     eta_minutes, eta_seconds_rem = divmod(int(eta_seconds), 60)
-
                     status_text = (
                         f"Generating... {progress_percent:.0f}% ({current_step}/{total_steps}) | "
                         f"{its_per_sec:.2f}it/s | "
                         f"ETA: {eta_minutes:02d}:{eta_seconds_rem:02d}"
                     )
                     yield {status_textbox: status_text}
-                elif update is None: # An error occurred
-                     yield {status_textbox: "Error during generation. Check console."}
+                elif update_type == "error":
+                     yield {status_textbox: f"Error: {data}"}
                      break
             except QueueEmpty:
                 if not thread.is_alive():
-                    print("⚠️ Generation thread finished unexpectedly.")
                     yield {status_textbox: "Generation failed. Check console for details."}
                     break
         
         thread.join()
 
-# --- GRADIO UI ---
+# --- GRADIO UI (Unchanged) ---
 with gr.Blocks(theme='gradio/soft') as demo:
-    # (UI layout is the same, just point to the new function)
+    # (Your UI code is perfect)
     gr.Markdown("# The Generative Media Suite")
     gr.Markdown("Create fast images, high-quality images, or short videos. Created by cheeseman182. (note: the speed on the status bar is wrong)")
     seed_state = gr.State(-1)
@@ -159,7 +155,7 @@ with gr.Blocks(theme='gradio/soft') as demo:
         with gr.Column(scale=2):
             model_selector = gr.Radio(label="Select Model", choices=list(available_models.keys()), value=list(available_models.keys())[0])
             prompt_input = gr.Textbox(label="Prompt", lines=4, placeholder="An astronaut riding a horse on Mars, cinematic...")
-            negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, deformed, watermark, text")
+            negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, deformed, watermark, text, overblown, high contrast, not photorealistic")
             with gr.Accordion("Settings", open=True):
                 steps_slider = gr.Slider(1, 100, 30, step=1, label="Inference Steps")
                 cfg_slider = gr.Slider(0.0, 15.0, 7.5, step=0.5, label="Guidance Scale (CFG)")
@@ -194,9 +190,9 @@ with gr.Blocks(theme='gradio/soft') as demo:
         outputs=seed_state,
         queue=False
     ).then(
-        fn=generate_media_live_progress, # Use the new function with progress
+        fn=generate_media_with_progress,
         inputs=[model_selector, prompt_input, negative_prompt_input, steps_slider, cfg_slider, width_slider, height_slider, seed_state, num_frames_slider],
         outputs=[output_image, output_video, status_textbox]
     )
 
-demo.launch()
+demo.launch(share=True)