new changes

app.py

@@ -2,12 +2,10 @@ 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,
-    LuminaText2ImgPipeline
+    LuminaText2ImgPipeline, SanaPipeline
 )
 import spaces
 import gc
@@ -20,108 +18,169 @@ import time
 import glob
 from datetime import datetime
 from PIL import Image
-from queue import Queue
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-from dataclasses import dataclass
-from typing import Optional, List
-
-@dataclass
-class MultiGPUConfig:
-    count: int = 2  # Number of GPUs to request
-    memory: int = 16  # Memory per GPU in GB
-    duration: int = 3600  # Duration in seconds
-
-class SpacesMultiGPU:
-    def __init__(self, config: Optional[MultiGPUConfig] = None):
-        self.config = config or MultiGPUConfig()
-    
-    def __call__(self, func):
-        # Apply multiple GPU decorators
-        decorated_func = func
-        for gpu_idx in range(self.config.count):
-            decorated_func = spaces.GPU(
-                device=gpu_idx,  # Specify which GPU to request
-                memory=self.config.memory,
-                duration=self.config.duration
-            )(decorated_func)
-        return decorated_func
-
-# Example usage in your generation code
-gpu_config = MultiGPUConfig(
-    count=2,  # Request 2 GPUs
-    duration=400  # 1 hour duration
-)
+
+#import os
+#cache_dir = '/workspace/hf_cache'
+
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
-TORCH_DTYPE = torch.bfloat16
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.bfloat16  if torch.cuda.is_available() else torch.float32
 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": {
         "repo_id": "black-forest-labs/FLUX.1-dev",
-        "pipeline_class": FluxPipeline
+        "pipeline_class": FluxPipeline,
+        #"cache_dir" : cache_dir
     },
     "Stable Diffusion 3.5": {
         "repo_id": "stabilityai/stable-diffusion-3.5-large",
-        "pipeline_class": StableDiffusion3Pipeline
+        "pipeline_class": StableDiffusion3Pipeline,
+        #"cache_dir" : cache_dir
+    },
+     "PixArt": {
+        "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
+        "pipeline_class": PixArtSigmaPipeline,
+        #"cache_dir" : cache_dir
+    },
+         "SANA": {
+        "repo_id": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
+        "pipeline_class": SanaPipeline,
+        #"cache_dir" : cache_dir
+    },
+        "AuraFlow": {
+        "repo_id": "fal/AuraFlow",
+        "pipeline_class": AuraFlowPipeline,
+        #"cache_dir" : cache_dir
+    },
+    "Kandinsky": {
+        "repo_id": "kandinsky-community/kandinsky-3",
+        "pipeline_class": Kandinsky3Pipeline,
+        #"cache_dir" : cache_dir
+    },
+    "Hunyuan": {
+        "repo_id": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
+        "pipeline_class": HunyuanDiTPipeline,
+        #"cache_dir" : cache_dir
+    },
+    "Lumina": {
+        "repo_id": "Alpha-VLLM/Lumina-Next-SFT-diffusers",
+        "pipeline_class": LuminaText2ImgPipeline,
+        #"cache_dir" : cache_dir
     }
-}
 
-# GPU allocation queue and model cache
-gpu_queue = Queue()
-for gpu_id in AVAILABLE_GPUS:
-    gpu_queue.put(gpu_id)
 
-model_cache = {}
+}
+
+# Dictionary to store model pipelines
+pipes = {}
 model_locks = {model_name: threading.Lock() for model_name in MODEL_CONFIGS.keys()}
 
-def get_next_available_gpu():
-    """Get the next available GPU from the queue"""
-    gpu_id = gpu_queue.get()
-    return gpu_id
 
-def release_gpu(gpu_id):
-    """Release GPU back to the queue"""
-    gpu_queue.put(gpu_id)
+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)
 
-def load_pipeline_on_gpu(model_name, gpu_id):
-    """Load model pipeline on specific GPU with memory tracking"""
+        # 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()
+
+        # 6. Remove cached files
+        remove_cache_dir(model_name)
+
+        # 7. Additional CUDA cleanup if available
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        # 8. Wait a small amount of time to ensure cleanup
+        time.sleep(1)
+
+    except Exception as e:
+        print(f"Error during cleanup of {model_name}: {str(e)}")
+
+    finally:
+        # Final garbage collection
+        gc.collect()
+        clear_torch_cache()
+
+
+def load_pipeline(model_name):
+    """Load model pipeline with memory tracking"""
+    initial_memory = get_process_memory()
     config = MODEL_CONFIGS[model_name]
-    
-    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()
-    
+
+    pipe = config["pipeline_class"].from_pretrained(
+        config["repo_id"],
+        torch_dtype=TORCH_DTYPE,
+        cache_dir=cache_dir
+    )
+    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")
+
     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)
     image.save(filepath)
     return filepath
 
+
 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)
+    files.sort(key=os.path.getctime, reverse=True)  # Sort by creation time
     return [
         {
             "path": f,
@@ -132,25 +191,35 @@ def get_generated_images():
         for f in files
     ]
 
-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)
+
+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}...")
+
             image = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -160,53 +229,22 @@ def generate_image_on_gpu(args):
                 height=height,
                 generator=generator,
             ).images[0]
-        
-        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
-        
-@SpacesMultiGPU(gpu_config)
-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
+
+            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
+
 
 # Gradio Interface
 css = """
@@ -218,8 +256,8 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"# Multi-GPU Image Generation ({len(AVAILABLE_GPUS)} GPUs Available)")
-        
+        gr.Markdown("# Multi-Model Image Generation")
+
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -229,14 +267,14 @@ with gr.Blocks(css=css) as demo:
                 container=False,
             )
             run_button = gr.Button("Generate", scale=0, variant="primary")
-        
+
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
             )
-            
+
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -244,9 +282,9 @@ with gr.Blocks(css=css) as demo:
                 step=1,
                 value=0,
             )
-            
+
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
+
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
@@ -262,7 +300,7 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024,
                 )
-            
+
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
@@ -278,7 +316,9 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=40,
                 )
-        
+
+        memory_indicator = gr.Markdown("Current memory usage: 0 GB")
+
         with gr.Row():
             with gr.Column(scale=2):
                 with gr.Tabs() as tabs:
@@ -287,18 +327,21 @@ with gr.Blocks(css=css) as demo:
                     for model_name in MODEL_CONFIGS.keys():
                         with gr.Tab(model_name):
                             results[model_name] = gr.Image(label=f"{model_name} Result")
-                            seeds[model_name] = gr.Number(label="Seed used", visible=False)
-            
-            with gr.Column(scale=1):
-                gr.Markdown("### Generated Images")
-                file_gallery = gr.Gallery(
-                    label="Generated Images",
-                    show_label=False,
-                    elem_id="file_gallery",
-                    columns=2,
-                    height=400
-                )
-                refresh_button = gr.Button("Refresh Gallery")
+                            seeds[model_name] = gr.Number(label="Seed used", visible=True)
+                with gr.Column(scale=1):
+                    gr.Markdown("### Generated Images")
+                    file_gallery = gr.Gallery(
+                        label="Generated Images",
+                        show_label=False,
+                        elem_id="file_gallery",
+                        columns=3,
+                        height=800,
+                        visible=True
+                    )
+                    refresh_button = gr.Button("Refresh Gallery")
+
+
+
 
     def update_gallery():
         """Update the file gallery"""
@@ -308,6 +351,44 @@ 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]])
@@ -340,6 +421,4 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    # Initialize multiprocessing for PyTorch
-    mp.set_start_method('spawn', force=True)
-    demo.launch()
+    demo.launch(server_name='0.0.0.0')