fix error

media.py

@@ -1,3 +1,6 @@
+
+# --- START OF FILE media.py (FINAL WITH LIVE PROGRESS & FIXES) ---
+
 # --- LIBRARIES ---
 import torch
 import gradio as gr
@@ -11,9 +14,10 @@ import numpy as np
 import threading
 from queue import Queue, Empty as QueueEmpty
 from PIL import Image
+import os
 from huggingface_hub import login
 
-# --- DYNAMIC HARDWARE DETECTION & AUTH ---
+# --- DYNAMIC HARDWARE DETECTION ---
 if torch.cuda.is_available():
     device = "cuda"
     torch_dtype = torch.float16
@@ -21,13 +25,22 @@ if torch.cuda.is_available():
 else:
     device = "cpu"
     torch_dtype = torch.float32
-    print("⚠️ No GPU detected.")
+    print("⚠️ No GPU detected. Using CPU.")
 
-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!")
 
-login(token=HF_TOKEN)
+HF_TOKEN = os.environ.get('HF_TOKEN')
+
+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:
+    # This message will show when you run the app locally, which is fine.
+    print("⚠️ No HF_TOKEN secret found. This is normal for local testing.")
+    print("   The deployed app will use the secret you set on Hugging Face.")
 
 # --- CONFIGURATION & STATE ---
 available_models = {
@@ -38,34 +51,46 @@ available_models = {
 }
 model_state = { "current_pipe": None, "loaded_model_name": None }
 
-# --- THE FINAL, STABLE GENERATION FUNCTION ---
-def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_scale, width, height, seed, num_frames):
+# --- THE FINAL GENERATION FUNCTION WITH LIVE PROGRESS & FIXES ---
+def generate_media_live_progress(model_key, prompt, negative_prompt, steps, cfg_scale, width, height, seed, num_frames):
     global model_state
-
-    # --- Model Loading ---
+    
+    # --- Model Loading & Cleanup ---
     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"):
-            pipe_to_delete = model_state.pop("current_pipe", None)
-            if pipe_to_delete: del pipe_to_delete
-            gc.collect()
-            torch.cuda.empty_cache()
         
+        # --- More Aggressive & Explicit Cleanup ---
+        pipe_to_delete = model_state.pop("current_pipe", None)
+        if pipe_to_delete:
+            # FIX: Explicitly move the model to CPU before deleting to free VRAM.
+            print("Offloading previous model to CPU...")
+            pipe_to_delete.to("cpu")
+            del pipe_to_delete
+            print("Previous model deleted.")
+        
+        # Explicitly run garbage collection and empty CUDA cache.
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        # Load the new pipeline
         model_id = available_models[model_key]
         if "Video" in model_key:
-            pipe = TextToVideoSDPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+            pipe = TextToVideoSDPipeline.from_pretrained(model_id, torch_dtype=torch_dtype, variant="fp16")
         else:
             pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch_dtype, variant="fp16")
-        
+    
         pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
         pipe.to(device)
-        
+    
         if device == "cuda":
-             if "Video" not in model_key: pipe.enable_model_cpu_offload()
-             pipe.enable_vae_slicing()
+            if "Video" not in model_key:
+                pipe.enable_model_cpu_offload()
+            pipe.enable_vae_slicing()
+            
         model_state["current_pipe"] = pipe
         model_state["loaded_model_name"] = model_key
-        print(f"✅ Model loaded on {device.upper()}.")
+        print(f"✅ Model '{model_key}' loaded on {device.upper()}.")
 
     pipe = model_state["current_pipe"]
     generator = torch.Generator(device).manual_seed(seed)
@@ -73,29 +98,32 @@ def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_
     # --- Generation Logic ---
     if "Video" in model_key:
         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)
-        video_uint8 = (video_frames_4d * 255).astype(np.uint8)
-        list_of_frames = [frame for frame in video_uint8]
-        video_path = f"video_{seed}.mp4"
-        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 your brilliant text-based progress bar
+        try:
+            video_frames = pipe(prompt=prompt, num_inference_steps=int(steps), height=320, width=576, num_frames=int(num_frames), generator=generator).frames
+            
+            # FIX: More memory-efficient video saving
+            video_path = f"video_{seed}.mp4"
+            with imageio.get_writer(video_path, fps=12) as writer:
+                for frame in video_frames:
+                    writer.append_data((frame * 255).astype(np.uint8))
+
+            yield {output_image: None, output_video: video_path, status_textbox: f"Video saved! Seed: {seed}"}
+        except Exception as e:
+            print(f"An error occurred during video generation: {e}")
+            yield {status_textbox: f"Error during video generation: {e}"}
+
+    else: # Image Generation with Live Progress
         progress_queue = Queue()
         
         def run_pipe():
             start_time = time.time()
             
-            # This callback correctly accepts all arguments
-            def progress_callback(step, timestep, latents, **kwargs):
+            def progress_callback(pipe, step, timestep, callback_kwargs):
                 elapsed_time = time.time() - start_time
                 if elapsed_time > 0:
                     its_per_sec = (step + 1) / elapsed_time
-                    progress_queue.put(("progress", step + 1, its_per_sec))
-                return kwargs
+                    progress_queue.put(("progress", (step + 1, its_per_sec)))
+                return callback_kwargs
 
             try:
                 final_image = pipe(
@@ -104,7 +132,7 @@ def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_
                     generator=generator, 
                     callback_on_step_end=progress_callback
                 ).images[0]
-                progress_queue.put(("final", final_image))
+                progress_queue.put(("result", final_image))
             except Exception as e:
                 print(f"An error occurred in the generation thread: {e}")
                 progress_queue.put(("error", str(e)))
@@ -113,19 +141,17 @@ def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_
         thread.start()
 
         total_steps = int(steps)
-        final_image_result = None
         yield {status_textbox: "Generating..."}
         
         while True:
             try:
-                update_type, data = progress_queue.get(timeout=1.0)
+                update_type, payload = progress_queue.get(timeout=1.0)
                 
-                if update_type == "final":
-                    final_image_result = data
-                    yield {output_image: final_image_result, status_textbox: f"Generation complete! Seed: {seed}"}
+                if update_type == "result":
+                    yield {output_image: payload, status_textbox: f"Generation complete! Seed: {seed}"}
                     break
                 elif update_type == "progress":
-                    current_step, its_per_sec = data
+                    current_step, its_per_sec = payload
                     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
@@ -137,18 +163,19 @@ def generate_media_with_progress(model_key, prompt, negative_prompt, steps, cfg_
                     )
                     yield {status_textbox: status_text}
                 elif update_type == "error":
-                     yield {status_textbox: f"Error: {data}"}
+                     yield {status_textbox: f"Error: {payload}. Check console."}
                      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()
+        print("Generation thread joined.")
 
-# --- GRADIO UI (Unchanged) ---
+# --- GRADIO UI ---
 with gr.Blocks(theme='gradio/soft') as demo:
-    # (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)
@@ -185,15 +212,16 @@ with gr.Blocks(theme='gradio/soft') as demo:
         }
     model_selector.change(update_ui_on_model_change, model_selector, [steps_slider, cfg_slider, width_slider, height_slider, num_frames_slider, output_image, output_video])
     
-    click_event = generate_button.click(
-        fn=lambda s: (s if s != -1 else random.randint(0, 2**32 - 1)),
+    generate_button.click(
+        fn=lambda s: s if s != -1 else random.randint(0, 2**32 - 1),
         inputs=seed_input,
         outputs=seed_state,
         queue=False
     ).then(
-        fn=generate_media_with_progress,
+        fn=generate_media_live_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(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True, debug=True)