One_Model

options/Banner_Model/Image2Image.py

@@ -8,10 +8,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device for I2I: {device}")
 
 # Load the inpainting pipeline
-pipe = AutoPipelineForInpainting.from_pretrained(
-    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-    torch_dtype=torch.float16, variant="fp16").to(device)
-
 
 def resize_image(image, height, width):
     """Resize image tensor to the desired height and width."""
@@ -30,6 +26,11 @@ def dummy(img):
 
 
 def I2I(prompt, image, width=1024, height=1024, guidance_scale=8.0, num_inference_steps=20, strength=0.99):
+    
+    pipe = AutoPipelineForInpainting.from_pretrained(
+        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        torch_dtype=torch.float16, variant="fp16").to(device)
+
     img_url, mask = dummy(image)
 
     # Resize image and mask to the target dimensions (height x width)

options/Banner_Model/Image2Image_2.py

@@ -5,17 +5,18 @@ from PIL import Image
 
 device= "cuda" if torch.cuda.is_available() else "cpu"
 print("Using device for I2I_2:", device)
-processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
 
-checkpoint = "ControlNet-1-1-preview/control_v11p_sd15_lineart"
-controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(device)
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "radames/stable-diffusion-v1-5-img2img", controlnet=controlnet, torch_dtype=torch.float16
-).to(device)
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-# pipe.enable_model_cpu_offload()
 
 def I2I_2(image, prompt,size,num_inference_steps):
+    processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
+
+    checkpoint = "ControlNet-1-1-preview/control_v11p_sd15_lineart"
+    controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(device)
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        "radames/stable-diffusion-v1-5-img2img", controlnet=controlnet, torch_dtype=torch.float16
+    ).to(device)
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_model_cpu_offload()
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
     image.resize((size,size))

options/Video_model/Model.py

@@ -7,17 +7,19 @@ login(token=os.getenv("TOKEN"))
 # Check if CUDA (GPU) is available, otherwise use CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the pipeline and move it to the appropriate device (GPU or CPU)
-pipeline = StableVideoDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float32
-).to(device)
-
-# Enable model offloading if using the CPU
-if device == "cpu":
-    pipeline.enable_model_cpu_offload()
 
 # Function to generate the video
 def Video(image):
+
+
+    pipeline = StableVideoDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float32
+    ).to(device)
+
+    # Enable model offloading if using the CPU
+    if device == "cpu":
+        pipeline.enable_model_cpu_offload()
+    
     image = Image.fromarray(image)
     image = image.resize((1024, 576))