Update app.py

app.py

@@ -12,10 +12,6 @@ from PIL import Image
 from huggingface_hub import snapshot_download
 import gc
 
-# Force CPU usage
-device = "cpu"
-dtype = torch.float32
-
 # Clear memory
 gc.collect()
 if torch.cuda.is_available():
@@ -28,14 +24,17 @@ css = """
 }
 """
 
+# Device configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float32
+
 huggingface_token = os.getenv("HF_TOKEN")
 
-# Minimal model configuration
+# Modified model configuration
 model_config = {
     "low_cpu_mem_usage": True,
     "torch_dtype": dtype,
-    "use_safetensors": True,
-    "device_map": "cpu"
+    "use_safetensors": False,  # Disabled safetensors
 }
 
 model_path = snapshot_download(
@@ -46,38 +45,42 @@ model_path = snapshot_download(
     token=huggingface_token,
 )
 
-# Load models on CPU
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler",
-    **model_config
-)
+# Load models with modified configuration
+try:
+    controlnet = FluxControlNetModel.from_pretrained(
+        "jasperai/Flux.1-dev-Controlnet-Upscaler",
+        **model_config
+    )
+    controlnet.to(device)
 
-pipe = FluxControlNetPipeline.from_pretrained(
-    model_path,
-    controlnet=controlnet,
-    **model_config
-)
+    pipe = FluxControlNetPipeline.from_pretrained(
+        model_path,
+        controlnet=controlnet,
+        **model_config
+    )
+    pipe.to(device)
+
+except Exception as e:
+    print(f"Error loading models: {str(e)}")
+    raise
 
 # Enable optimizations
 pipe.enable_attention_slicing(1)
 pipe.enable_vae_slicing()
 
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 64 * 64  # Extremely reduced
+MAX_PIXEL_BUDGET = 64 * 64
 
 def process_input(input_image, upscale_factor):
     input_image = input_image.convert('RGB')
     
-    # Aggressive size reduction
     w, h = input_image.size
     max_size = int(np.sqrt(MAX_PIXEL_BUDGET))
     
-    # Resize to very small size
     new_w = min(w, max_size)
     new_h = min(h, max_size)
     input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
     
-    # Ensure dimensions are multiples of 8
     w = new_w - new_w % 8
     h = new_h - new_h % 8
     
@@ -94,6 +97,8 @@ def infer(
 ):
     try:
         gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
@@ -101,7 +106,7 @@ def infer(
         input_image, w, h = process_input(input_image, upscale_factor)
         
         with torch.inference_mode():
-            generator = torch.Generator().manual_seed(seed)
+            generator = torch.Generator(device=device).manual_seed(seed)
             image = pipe(
                 prompt="",
                 control_image=input_image,
@@ -114,6 +119,8 @@ def infer(
             ).images[0]
             
         gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
         return [input_image, image, seed]
 
@@ -197,7 +204,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         show_api=False,
     )
 
-# Minimal launch configuration
+# Launch configuration
 demo.queue(max_size=1).launch(
     share=False,
     debug=True,