Update model_loading.py

model_loading.py

@@ -1,8 +1,7 @@
 from PIL import Image
 import os
-import ollama
 import torch
-from diffusers import DiffusionPipeline, AutoPipelineForImage2Image, LCMScheduler, AutoPipelineForText2Image
+from diffusers import DiffusionPipeline, AutoPipelineForImage2Image, LCMScheduler
 import time
 
 model_id = "simianluo/lcm_dreamshaper_v7"
@@ -14,42 +13,41 @@ class GenerationSession:
         self.img2img_pipeline = None
         self.current_image = None
         self.current_prompt = None
+        
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        
         self._initialize_pipelines()
 
     def _initialize_pipelines(self):
-            print("initializing pipelines...")
-
-            self.txt2img_pipeline = DiffusionPipeline.from_pretrained(
-                model_id,
-                torch_dtype = torch.float16,
-                safety_checker = None
-            )
-            
-            self.txt2img_pipeline.scheduler = LCMScheduler.from_config(self.txt2img_pipeline.scheduler.config)
-            self.txt2img_pipeline.to("cuda")
-            self.txt2img_pipeline.enable_attention_slicing()
-            self.txt2img_pipeline.enable_vae_slicing()
+        print(f"Initializing pipelines on device: {self.device}...")
 
-           # self.txt2img_pipeline.unet = torch.compile(
-           #   self.txt2img_pipeline.unet,
-           #    mode = "reduce-overhead",
-           #     fullgraph = True
-           #)
-            print("Text 2 image pipeline loaded and compiled.")
+        self.txt2img_pipeline = DiffusionPipeline.from_pretrained(
+            self.model_id,
+            torch_dtype=self.dtype,
+            safety_checker=None
+        )
+        
+        self.txt2img_pipeline.scheduler = LCMScheduler.from_config(self.txt2img_pipeline.scheduler.config)
+        self.txt2img_pipeline.to(self.device)
+        
+        self.txt2img_pipeline.enable_attention_slicing()
+        self.txt2img_pipeline.enable_vae_slicing()
 
+        print("Text 2 image pipeline loaded.")
 
-            self.img2img_pipeline = AutoPipelineForImage2Image.from_pipe(self.txt2img_pipeline)  
-            print("Image 2 image pipeline loaded (shared weights).")
+        self.img2img_pipeline = AutoPipelineForImage2Image.from_pipe(self.txt2img_pipeline)  
+        print("Image 2 image pipeline loaded (shared weights).")
 
     def GeneratingBaseImage(self, prompt: str, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
         start = time.time()
         image = self.txt2img_pipeline(
-            prompt = prompt,
-            negative_prompt= negative_prompt,
-            num_inference_steps = 4,
-            guidance_scale = 1.0,
-            height = 512,
-            width = 512
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            height=512,
+            width=512
         ).images
         print(f"Text to image generated in [{time.time() - start:.2f}s]")
         return image
@@ -57,12 +55,12 @@ class GenerationSession:
     def GeneratingVariationImage(self, prompt: str, reference_image: Image.Image, strength: float = 0.5, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
         start = time.time()
         image = self.img2img_pipeline(
-            prompt = prompt,
-            image = reference_image,
-            strength = strength,
-            num_inference_steps = 4,
-            guidance_scale = 1.0,
-            negative_prompt = negative_prompt
+            prompt=prompt,
+            image=reference_image,
+            strength=strength,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            negative_prompt=negative_prompt
         ).images
         print(f"Image to image generated in [{time.time() - start:.2f}s]")
         return image
@@ -79,5 +77,4 @@ class GenerationSession:
     def reset(self):
         self.current_image = None
         self.current_prompt = None
-        print("Session reset. Ready for new generation.")
-
+        print("Session reset. Ready for new generation.")