Initial commitcq

app.py

@@ -15,6 +15,7 @@ from PIL import Image
 import io
 import requests
 from transformers import DPTImageProcessor, DPTForDepthEstimation
+import gc
 
 # Available models
 AVAILABLE_MODELS = {
@@ -28,6 +29,21 @@ AVAILABLE_LORAS = {
     "Papercut": "TheLastBen/Papercut_SDXL",
 }
 
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+def get_dtype():
+    if torch.cuda.is_available():
+        return torch.float16
+    return torch.float32
+
+def cleanup_memory():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
 def get_depth_map(image):
     # Initialize depth estimator
     depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")
@@ -61,15 +77,20 @@ def load_image_from_url(url):
 
 def generate_image(prompt, seed, num_steps, guidance_scale, eta):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         tcd_lora_id = "h1t/TCD-SDXL-LoRA"
         
-        # Use CPU for inference
         pipe = StableDiffusionXLPipeline.from_pretrained(
             base_model_id,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
+        
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
         
         # Load and fuse LoRA weights with prefix=None
@@ -77,7 +98,7 @@ def generate_image(prompt, seed, num_steps, guidance_scale, eta):
         pipe.fuse_lora()
         
         # Generate the image
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         image = pipe(
             prompt=prompt,
             num_inference_steps=num_steps,
@@ -86,21 +107,31 @@ def generate_image(prompt, seed, num_steps, guidance_scale, eta):
             generator=generator,
         ).images[0]
         
+        # Cleanup
+        del pipe
+        cleanup_memory()
+        
         return image, "Image generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating image: {str(e)}"
 
 def generate_community_image(prompt, model_name, seed, num_steps, guidance_scale, eta):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = AVAILABLE_MODELS[model_name]
         tcd_lora_id = "h1t/TCD-SDXL-LoRA"
         
-        # Use CPU for inference
         pipe = StableDiffusionXLPipeline.from_pretrained(
             base_model_id,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
+        
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
         
         # Load and fuse LoRA weights with prefix=None
@@ -108,7 +139,7 @@ def generate_community_image(prompt, model_name, seed, num_steps, guidance_scale
         pipe.fuse_lora()
         
         # Generate the image
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         image = pipe(
             prompt=prompt,
             num_inference_steps=num_steps,
@@ -117,22 +148,32 @@ def generate_community_image(prompt, model_name, seed, num_steps, guidance_scale
             generator=generator,
         ).images[0]
         
+        # Cleanup
+        del pipe
+        cleanup_memory()
+        
         return image, "Image generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating image: {str(e)}"
 
 def generate_style_mix(prompt, seed, num_steps, guidance_scale, eta, style_weight):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         tcd_lora_id = "h1t/TCD-SDXL-LoRA"
         styled_lora_id = "TheLastBen/Papercut_SDXL"
         
-        # Use CPU for inference
         pipe = StableDiffusionXLPipeline.from_pretrained(
             base_model_id,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
+        
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
         
         # Load multiple LoRA weights with prefix=None
@@ -143,7 +184,7 @@ def generate_style_mix(prompt, seed, num_steps, guidance_scale, eta, style_weigh
         pipe.set_adapters(["tcd", "style"], adapter_weights=[1.0, style_weight])
         
         # Generate the image
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         image = pipe(
             prompt=prompt,
             num_inference_steps=num_steps,
@@ -152,12 +193,20 @@ def generate_style_mix(prompt, seed, num_steps, guidance_scale, eta, style_weigh
             generator=generator,
         ).images[0]
         
+        # Cleanup
+        del pipe
+        cleanup_memory()
+        
         return image, "Image generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating image: {str(e)}"
 
 def generate_controlnet(prompt, init_image, seed, num_steps, guidance_scale, eta, controlnet_scale):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         controlnet_id = "diffusers/controlnet-depth-sdxl-1.0"
@@ -166,15 +215,20 @@ def generate_controlnet(prompt, init_image, seed, num_steps, guidance_scale, eta
         # Initialize ControlNet
         controlnet = ControlNetModel.from_pretrained(
             controlnet_id,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
         
         # Initialize pipeline
         pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
             base_model_id,
             controlnet=controlnet,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
+        
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
         
         # Load and fuse LoRA weights with prefix=None
@@ -185,7 +239,7 @@ def generate_controlnet(prompt, init_image, seed, num_steps, guidance_scale, eta
         depth_image = get_depth_map(init_image)
         
         # Generate the image
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         image = pipe(
             prompt=prompt,
             image=depth_image,
@@ -198,21 +252,32 @@ def generate_controlnet(prompt, init_image, seed, num_steps, guidance_scale, eta
         
         # Create a grid of the depth map and result
         grid = make_image_grid([depth_image, image], rows=1, cols=2)
+        
+        # Cleanup
+        del pipe, controlnet
+        cleanup_memory()
+        
         return grid, "Image generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating image: {str(e)}"
 
 def inpaint_image(prompt, init_image, mask_image, seed, num_steps, guidance_scale, eta, strength):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
         tcd_lora_id = "h1t/TCD-SDXL-LoRA"
         
-        # Use CPU for inference
         pipe = AutoPipelineForInpainting.from_pretrained(
             base_model_id,
-            torch_dtype=torch.float32  # Use float32 for CPU
-        )
+            torch_dtype=dtype,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
+        
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
         
         # Load and fuse LoRA weights with prefix=None
@@ -220,7 +285,7 @@ def inpaint_image(prompt, init_image, mask_image, seed, num_steps, guidance_scal
         pipe.fuse_lora()
         
         # Generate the image
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         image = pipe(
             prompt=prompt,
             image=init_image,
@@ -234,12 +299,21 @@ def inpaint_image(prompt, init_image, mask_image, seed, num_steps, guidance_scal
         
         # Create a grid of the original image, mask, and result
         grid = make_image_grid([init_image, mask_image, image], rows=1, cols=3)
+        
+        # Cleanup
+        del pipe
+        cleanup_memory()
+        
         return grid, "Image generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating image: {str(e)}"
 
 def generate_animation(prompt, seed, num_steps, guidance_scale, eta, num_frames, motion_scale):
     try:
+        device = get_device()
+        dtype = get_dtype()
+        
         # Initialize the pipeline
         base_model_id = "frankjoshua/toonyou_beta6"
         motion_adapter_id = "guoyww/animatediff-motion-adapter-v1-5"
@@ -247,16 +321,17 @@ def generate_animation(prompt, seed, num_steps, guidance_scale, eta, num_frames,
         motion_lora_id = "guoyww/animatediff-motion-lora-zoom-in"
         
         # Load motion adapter
-        adapter = MotionAdapter.from_pretrained(motion_adapter_id)
+        adapter = MotionAdapter.from_pretrained(motion_adapter_id).to(device)
         
-        # Initialize pipeline with CPU optimization
+        # Initialize pipeline with optimization
         pipe = AnimateDiffPipeline.from_pretrained(
             base_model_id,
             motion_adapter=adapter,
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            low_cpu_mem_usage=True,     # Enable low CPU memory usage
-            use_safetensors=False       # Use standard PyTorch weights
-        )
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+            use_safetensors=True,
+            variant="fp16" if device == "cuda" else None
+        ).to(device)
         
         # Set TCD scheduler
         pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
@@ -273,7 +348,7 @@ def generate_animation(prompt, seed, num_steps, guidance_scale, eta, num_frames,
         pipe.set_adapters(["tcd", "motion-lora"], adapter_weights=[1.0, motion_scale])
         
         # Generate animation
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         frames = pipe(
             prompt=prompt,
             num_inference_steps=num_steps,
@@ -287,8 +362,14 @@ def generate_animation(prompt, seed, num_steps, guidance_scale, eta, num_frames,
         # Export to GIF
         gif_path = "animation.gif"
         export_to_gif(frames, gif_path)
+        
+        # Cleanup
+        del pipe, adapter
+        cleanup_memory()
+        
         return gif_path, "Animation generated successfully!"
     except Exception as e:
+        cleanup_memory()
         return None, f"Error generating animation: {str(e)}"
 
 # Create the Gradio interface