multi gpu

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import numpy as np
 import random
 import torch
+import torch.multiprocessing as mp
+from torch.cuda.amp import autocast
 from diffusers import (
     DiffusionPipeline, StableDiffusion3Pipeline, FluxPipeline, PixArtSigmaPipeline,
     AuraFlowPipeline, Kandinsky3Pipeline, HunyuanDiTPipeline,
@@ -18,15 +20,20 @@ import time
 import glob
 from datetime import datetime
 from PIL import Image
+from queue import Queue
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+TORCH_DTYPE = torch.bfloat16
 OUTPUT_DIR = "generated_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# Get available GPU devices
+AVAILABLE_GPUS = list(range(torch.cuda.device_count()))
+print(f"Available GPUs: {AVAILABLE_GPUS}")
+
 # Model configurations
 MODEL_CONFIGS = {
     "FLUX": {
@@ -37,94 +44,44 @@ MODEL_CONFIGS = {
         "repo_id": "stabilityai/stable-diffusion-3.5-large",
         "pipeline_class": StableDiffusion3Pipeline
     }
-    
 }
 
-# Dictionary to store model pipelines
-pipes = {}
-model_locks = {model_name: threading.Lock() for model_name in MODEL_CONFIGS.keys()}
-
-def get_process_memory():
-    """Get memory usage of current process in GB"""
-    process = psutil.Process(os.getpid())
-    return process.memory_info().rss / 1024 / 1024 / 1024
-
-def clear_torch_cache():
-    """Clear PyTorch's CUDA cache"""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-
-def remove_cache_dir(model_name):
-    """Remove the model's cache directory"""
-    cache_dir = Path.home() / '.cache' / 'huggingface' / 'diffusers' / MODEL_CONFIGS[model_name]['repo_id'].replace('/', '--')
-    if cache_dir.exists():
-        shutil.rmtree(cache_dir, ignore_errors=True)
-
-def deep_cleanup(model_name, pipe):
-    """Perform deep cleanup of model resources"""
-    try:
-        # 1. Move model to CPU first (helps prevent CUDA memory fragmentation)
-        if hasattr(pipe, 'to'):
-            pipe.to('cpu')
-
-        # 2. Delete all model components explicitly
-        for attr_name in list(pipe.__dict__.keys()):
-            if hasattr(pipe, attr_name):
-                delattr(pipe, attr_name)
-
-        # 3. Remove from pipes dictionary
-        if model_name in pipes:
-            del pipes[model_name]
-
-        # 4. Clear CUDA cache
-        clear_torch_cache()
-
-        # 5. Run garbage collection multiple times
-        for _ in range(3):
-            gc.collect()
+# GPU allocation queue and model cache
+gpu_queue = Queue()
+for gpu_id in AVAILABLE_GPUS:
+    gpu_queue.put(gpu_id)
 
-        # 6. Remove cached files
-        remove_cache_dir(model_name)
-
-        # 7. Additional CUDA cleanup if available
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
+model_cache = {}
+model_locks = {model_name: threading.Lock() for model_name in MODEL_CONFIGS.keys()}
 
-        # 8. Wait a small amount of time to ensure cleanup
-        time.sleep(1)
+def get_next_available_gpu():
+    """Get the next available GPU from the queue"""
+    gpu_id = gpu_queue.get()
+    return gpu_id
 
-    except Exception as e:
-        print(f"Error during cleanup of {model_name}: {str(e)}")
-    
-    finally:
-        # Final garbage collection
-        gc.collect()
-        clear_torch_cache()
+def release_gpu(gpu_id):
+    """Release GPU back to the queue"""
+    gpu_queue.put(gpu_id)
 
-def load_pipeline(model_name):
-    """Load model pipeline with memory tracking"""
-    initial_memory = get_process_memory()
+def load_pipeline_on_gpu(model_name, gpu_id):
+    """Load model pipeline on specific GPU with memory tracking"""
     config = MODEL_CONFIGS[model_name]
     
-    pipe = config["pipeline_class"].from_pretrained(
-        config["repo_id"],
-        torch_dtype=TORCH_DTYPE
-    )
-    pipe = pipe.to(DEVICE)
-    
-    if hasattr(pipe, 'enable_model_cpu_offload'):
-        pipe.enable_model_cpu_offload()
-    
-    final_memory = get_process_memory()
-    print(f"Memory used by {model_name}: {final_memory - initial_memory:.2f} GB")
+    with torch.cuda.device(gpu_id):
+        pipe = config["pipeline_class"].from_pretrained(
+            config["repo_id"],
+            torch_dtype=TORCH_DTYPE
+        )
+        pipe = pipe.to(f"cuda:{gpu_id}")
+        
+        if hasattr(pipe, 'enable_model_cpu_offload'):
+            pipe.enable_model_cpu_offload()
     
     return pipe
 
 def save_generated_image(image, model_name, prompt):
     """Save generated image with timestamp and model name"""
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    # Create sanitized filename from prompt (first 30 chars)
     prompt_part = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).strip()
     filename = f"{timestamp}_{model_name}_{prompt_part}.png"
     filepath = os.path.join(OUTPUT_DIR, filename)
@@ -134,7 +91,7 @@ def save_generated_image(image, model_name, prompt):
 def get_generated_images():
     """Get list of generated images with their details"""
     files = glob.glob(os.path.join(OUTPUT_DIR, "*.png"))
-    files.sort(key=os.path.getctime, reverse=True)  # Sort by creation time
+    files.sort(key=os.path.getctime, reverse=True)
     return [
         {
             "path": f,
@@ -145,34 +102,25 @@ def get_generated_images():
         for f in files
     ]
 
-def generate_image(
-    model_name,
-    prompt,
-    negative_prompt="",
-    seed=42,
-    randomize_seed=False,
-    width=1024,
-    height=1024,
-    guidance_scale=4.5,
-    num_inference_steps=40,
-    progress=gr.Progress(track_tqdm=True)
-):
-    with model_locks[model_name]:
-        try:
-            #progress(0, desc=f"Loading {model_name} model...")
-            
-            if model_name not in pipes:
-                pipes[model_name] = load_pipeline(model_name)
-            
-            pipe = pipes[model_name]
-            
-            if randomize_seed:
-                seed = random.randint(0, MAX_SEED)
-            
-            generator = torch.Generator(DEVICE).manual_seed(seed)
-            print(f"Generating image with {model_name}...")
-            #progress(0.3, desc=f"Generating image with {model_name}...")
-            
+def generate_image_on_gpu(args):
+    """Generate image on specific GPU"""
+    model_name, prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps = args
+    
+    try:
+        gpu_id = get_next_available_gpu()
+        print(f"Generating {model_name} on GPU {gpu_id}")
+        
+        # Load or get cached pipeline
+        cache_key = f"{model_name}_{gpu_id}"
+        if cache_key not in model_cache:
+            with model_locks[model_name]:
+                model_cache[cache_key] = load_pipeline_on_gpu(model_name, gpu_id)
+        
+        pipe = model_cache[cache_key]
+        
+        # Generate image
+        with torch.cuda.device(gpu_id), autocast():
+            generator = torch.Generator(f"cuda:{gpu_id}").manual_seed(seed)
             image = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -182,21 +130,52 @@ def generate_image(
                 height=height,
                 generator=generator,
             ).images[0]
-            
-            filepath = save_generated_image(image, model_name, prompt)
-            print(f"Saved image to: {filepath}")
-            
-            #progress(0.9, desc=f"Cleaning up {model_name} resources...")
-            #deep_cleanup(model_name, pipe)
-            
-            #progress(1.0, desc=f"Generation complete with {model_name}")
-            return image, seed
-            
-        except Exception as e:
-            print(f"Error with {model_name}: {str(e)}")
-            if model_name in pipes:
-                deep_cleanup(model_name, pipes[model_name])
-            raise e
+        
+        filepath = save_generated_image(image, model_name, prompt)
+        print(f"Saved image from {model_name} to: {filepath}")
+        
+        release_gpu(gpu_id)
+        return image, seed
+        
+    except Exception as e:
+        print(f"Error with {model_name} on GPU {gpu_id}: {str(e)}")
+        release_gpu(gpu_id)
+        raise e
+
+def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress()):
+    outputs = [None] * (len(MODEL_CONFIGS) * 2)
+    
+    # Prepare generation tasks
+    tasks = []
+    for model_name in MODEL_CONFIGS.keys():
+        current_seed = random.randint(0, MAX_SEED) if randomize_seed else seed
+        tasks.append((
+            model_name, prompt, negative_prompt, current_seed,
+            width, height, guidance_scale, num_inference_steps
+        ))
+    
+    # Run generation in parallel using thread pool
+    with ThreadPoolExecutor(max_workers=len(AVAILABLE_GPUS)) as executor:
+        future_to_model = {
+            executor.submit(generate_image_on_gpu, task): idx 
+            for idx, task in enumerate(tasks)
+        }
+        
+        for future in as_completed(future_to_model):
+            idx = future_to_model[future]
+            try:
+                image, used_seed = future.result()
+                outputs[idx * 2] = image
+                outputs[idx * 2 + 1] = used_seed
+                yield outputs + [None]
+            except Exception as e:
+                print(f"Generation failed for model {idx}: {str(e)}")
+                outputs[idx * 2] = None
+                outputs[idx * 2 + 1] = None
+    
+    # Update gallery after all generations complete
+    gallery_images = update_gallery()
+    return outputs
 
 # Gradio Interface
 css = """
@@ -208,7 +187,7 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# Multi-Model Image Generation")
+        gr.Markdown(f"# Multi-GPU Image Generation ({len(AVAILABLE_GPUS)} GPUs Available)")
         
         with gr.Row():
             prompt = gr.Text(
@@ -269,8 +248,6 @@ with gr.Blocks(css=css) as demo:
                     value=40,
                 )
         
-        memory_indicator = gr.Markdown("Current memory usage: 0 GB")
-        
         with gr.Row():
             with gr.Column(scale=2):
                 with gr.Tabs() as tabs:
@@ -291,7 +268,6 @@ with gr.Blocks(css=css) as demo:
                     height=400
                 )
                 refresh_button = gr.Button("Refresh Gallery")
-        
 
     def update_gallery():
         """Update the file gallery"""
@@ -301,41 +277,6 @@ with gr.Blocks(css=css) as demo:
             for f in files
         ]
 
-    @spaces.GPU(duration=600)
-    def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress()):
-        outputs = [None] * (len(MODEL_CONFIGS) * 2)
-        for idx, model_name in enumerate(MODEL_CONFIGS.keys()):
-            try:
-                # Display progress for the specific model
-                #progress(0, desc=f"Starting generation for {model_name}...")
-                print(f"IMAGE GENERATING {model_name} ")
-                image, used_seed = generate_image(
-                    model_name, prompt, negative_prompt, seed,
-                    randomize_seed, width, height, guidance_scale,
-                    num_inference_steps, progress
-                )
-                print(f"IMAGE GENERATIED {model_name} ")
-                # Update the respective model's tab with the generated image
-                #results[model_name].update(image)
-                #seeds[model_name].update(used_seed)
-                outputs[idx * 2] = image  # Image slot
-                outputs[idx * 2 + 1] = seed  # Seed slot
-                #outputs.extend([image, used_seed])
-                # Add intermediate results to progress * (len(all_outputs) - len(all_outputs))
-                print("YELID")
-                yield outputs + [None]
-                
-            
-            except Exception as e:
-                print(f"Error generating with {model_name}: {str(e)}")
-                outputs[idx * 2] = None
-                outputs[idx * 2 + 1] = None
-        
-        # Update the gallery after generation
-        gallery_images = update_gallery()
-        #file_gallery.update(value=gallery_images)
-        return outputs
-
     output_components = []
     for model_name in MODEL_CONFIGS.keys():
         output_components.extend([results[model_name], seeds[model_name]])
@@ -368,4 +309,6 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
+    # Initialize multiprocessing for PyTorch
+    mp.set_start_method('spawn', force=True)
     demo.launch()