Update app.py

app.py

@@ -71,53 +71,66 @@ img_transforms = transforms.Compose([
     transforms.Normalize(means,stds)
 ])
 
-# Load models globally (outside GPU-decorated functions)
-modelv4 = torch.jit.load(modelarcanev4).eval()
-modelv3 = torch.jit.load(modelarcanev3).eval()
-modelv2 = torch.jit.load(modelarcanev2).eval()
+# Load models globally on CPU
+modelv4_cpu = torch.jit.load(modelarcanev4, map_location='cpu').eval()
+modelv3_cpu = torch.jit.load(modelarcanev3, map_location='cpu').eval()
+modelv2_cpu = torch.jit.load(modelarcanev2, map_location='cpu').eval()
 
 @spaces.GPU
-def proc_pil_img(input_image, model):
+def proc_pil_img(input_image, model_path):
     """GPU-accelerated image processing"""
-    # Move tensors to GPU inside the decorated function
-    t_stds = torch.tensor(stds).cuda().half()[:,None,None]
-    t_means = torch.tensor(means).cuda().half()[:,None,None]
+    # Load model fresh on GPU to avoid device mismatch
+    model = torch.jit.load(model_path, map_location='cuda').eval()
     
-    # Move model to GPU
-    model = model.cuda().half()
+    # Create tensors on GPU
+    t_stds = torch.tensor(stds).cuda().view(3, 1, 1)
+    t_means = torch.tensor(means).cuda().view(3, 1, 1)
     
-    transformed_image = img_transforms(input_image)[None,...].cuda().half()
+    # Transform and move to GPU
+    transformed_image = img_transforms(input_image).unsqueeze(0).cuda()
     
     with torch.no_grad():
         result_image = model(transformed_image)[0]
-        output_image = result_image.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)
-        output_image = output_image.detach().cpu().numpy().astype('uint8')
+        output_image = result_image.mul(t_stds).add(t_means).mul(255.).clamp(0, 255).permute(1, 2, 0)
+        output_image = output_image.cpu().numpy().astype('uint8')
         output_image = PIL.Image.fromarray(output_image)
     
+    # Clean up
+    del model
+    torch.cuda.empty_cache()
+    
     return output_image
 
 @spaces.GPU
 def process(im, version):
     """Main processing function with GPU acceleration"""
     if im is None:
-        return None
-    
-    # Ensure image is PIL Image
-    if not isinstance(im, Image.Image):
-        im = Image.fromarray(im)
-    
-    # Scale image (CPU operation)
-    im = scale_by_face_size(im, target_face=256, max_res=1_500_000, max_upscale=1)
+        raise gr.Error("Please upload an image first!")
     
-    # Select model based on version
-    if version == 'v0.4 (Recommended)':
-        res = proc_pil_img(im, modelv4)
-    elif version == 'v0.3':
-        res = proc_pil_img(im, modelv3)
-    else:
-        res = proc_pil_img(im, modelv2)
+    try:
+        # Ensure image is PIL Image
+        if not isinstance(im, Image.Image):
+            im = Image.fromarray(im)
+        
+        # Convert to RGB if needed
+        if im.mode != 'RGB':
+            im = im.convert('RGB')
+        
+        # Scale image (CPU operation)
+        im = scale_by_face_size(im, target_face=256, max_res=1_500_000, max_upscale=1)
+        
+        # Select model based on version
+        if version == 'v0.4 (Recommended)':
+            res = proc_pil_img(im, modelarcanev4)
+        elif version == 'v0.3':
+            res = proc_pil_img(im, modelarcanev3)
+        else:
+            res = proc_pil_img(im, modelarcanev2)
+        
+        return res
     
-    return res
+    except Exception as e:
+        raise gr.Error(f"Error processing image: {str(e)}")
 
 # Custom theme
 custom_theme = gr.themes.Soft(
@@ -287,21 +300,16 @@ with gr.Blocks() as demo:
                 buttons=["download", "share"]
             )
     
-    # Examples section
-    gr.Markdown("### 🖼️ Try These Examples", elem_id="examples-header")
+    # Tips section
     with gr.Row():
-        gr.Examples(
-            examples=[
-                ['bill.png', 'v0.3'],
-                ['keanu.png', 'v0.4 (Recommended)'],
-                ['will.jpeg', 'v0.4 (Recommended)']
-            ],
-            inputs=[input_image, version_selector],
-            outputs=output_image,
-            fn=process,
-            cache_examples=False,
-            label="Example Images",
-            examples_per_page=3
+        gr.Markdown(
+            """
+            ### 💡 Tips for Best Results
+            - Use clear, well-lit portrait photos
+            - Face should be clearly visible and not too small
+            - Works best with frontal or slightly angled faces
+            - Try different model versions for varied artistic styles
+            """
         )
     
     # Footer
@@ -315,6 +323,11 @@ with gr.Blocks() as demo:
             
             **⚡ Zero-GPU Optimization**: This Space uses Hugging Face's Zero-GPU infrastructure for efficient GPU allocation.
             
+            **Model Versions:**
+            - **v0.4**: Latest and recommended - best quality and style accuracy
+            - **v0.3**: Alternative style interpretation
+            - **v0.2**: Original version with unique characteristics
+            
             <div style='margin-top: 1rem;'>
                 <img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_arcanegan' alt='visitor badge'>
             </div>