fix bugs

app.py

@@ -4,14 +4,14 @@ import numpy as np
 from PIL import Image
 from model import CycleGAN, get_val_transform, de_normalize
 
-# Configure page
+
 st.set_page_config(
     page_title="CycleGAN Image Converter",
     page_icon="🎨",
     layout="wide"
 )
 
-# Get the best available device
+
 @st.cache_resource
 def get_device():
     if torch.cuda.is_available():
@@ -25,7 +25,7 @@ def get_device():
         st.sidebar.info("Using CPU 💻")
     return device
 
-# Add custom CSS
+
 st.markdown("""
     <style>
     .stApp {
@@ -38,7 +38,7 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-# Title and description
+
 st.title("CycleGAN Image Converter 🎨")
 st.markdown("""
     Transform images between different domains using CycleGAN.
@@ -47,7 +47,7 @@ st.markdown("""
     *Note: Images will be resized to 256x256 pixels during conversion.*
 """)
 
-# Available models and their configurations
+
 MODELS = [
     {
         "name": "Cezanne ↔ Photo",
@@ -57,25 +57,25 @@ MODELS = [
     }
 ]
 
-# Sidebar controls
+
 with st.sidebar:
     st.header("Settings")
     
-    # Model selection
+    
     selected_model = st.selectbox(
         "Conversion Type",
         options=range(len(MODELS)),
         format_func=lambda x: MODELS[x]["name"]
     )
     
-    # Direction selection
+    
     direction = st.radio(
         "Conversion Direction",
         options=["A → B", "B → A"],
         help="A → B: Convert from domain A to B\nB → A: Convert from domain B to A"
     )
 
-# Load model
+
 @st.cache_resource
 def load_model(model_path):
     device = get_device()
@@ -84,29 +84,21 @@ def load_model(model_path):
     model.eval()
     return model
 
-# Process image
+
 def process_image(image, model, direction):
-    # Prepare transform
     transform = get_val_transform(model, direction)
-    
-    # Convert PIL image to tensor
     tensor = transform(np.array(image)).unsqueeze(0)
-    
-    # Move to appropriate device
     tensor = tensor.to(next(model.parameters()).device)
-    
-    # Process
     with torch.no_grad():
         if direction == "A → B":
             output = model.generator_ab(tensor)
         else:
             output = model.generator_ba(tensor)
-    
-    # Convert back to image
     result = de_normalize(output[0], model, direction)
-    return result
+    image = Image.fromarray(result.cpu().detach().numpy(), 'RGB')
+    return image
+
 
-# Main interface
 col1, col2 = st.columns(2)
 
 with col1:
@@ -115,19 +107,23 @@ with col1:
     
     if uploaded_file is not None:
         input_image = Image.open(uploaded_file)
-        st.image(input_image, use_column_width=True)
+        
+        if input_image.mode == 'RGBA':
+            input_image = input_image.convert('RGB')
+        st.image(input_image)
 
 with col2:
     st.subheader("Converted Image")
     if uploaded_file is not None:
         try:
-            # Load and process
+            
             model = load_model(MODELS[selected_model]["model_path"])
             result = process_image(input_image, model, direction)
             
-            # Display
-            st.image(result, use_column_width=True)
+            
+            st.image(result)
         except Exception as e:
-            st.error(f"Error during conversion: {str(e)}")
+            st.error(f"Error during conversion: {str(e)} {e.__traceback__}")
+            raise
     else:
-        st.info("Upload an image to see the conversion result") 
+        st.info("Upload an image to see the conversion result")

model.py

@@ -77,7 +77,6 @@ class CycleGAN(nn.Module, PyTorchModelHubMixin, pipeline_tag="image-to-image"):
 def get_val_transform(model, direction="a_to_b", size=256):
     mean = model.channel_mean_a if direction == "a_to_b" else model.channel_mean_b
     std = model.channel_std_a if direction == "a_to_b" else model.channel_std_b
-    
     return tr.Compose([
         tr.ToPILImage(),
         tr.Resize(size),
@@ -87,9 +86,8 @@ def get_val_transform(model, direction="a_to_b", size=256):
     ])
 
 def de_normalize(tensor, model, direction="a_to_b"):
-    img_tensor = tensor.cpu().detach().clone()
+    img_tensor = tensor
     mean = model.channel_mean_a if direction == "a_to_b" else model.channel_mean_b
     std = model.channel_std_a if direction == "a_to_b" else model.channel_std_b
-    
     img_tensor = img_tensor * std[:, None, None] + mean[:, None, None]
-    return torch.clamp(img_tensor.permute(1, 2, 0), 0.0, 1.0) 
+    return torch.clamp(img_tensor.permute(1, 2, 0) * 255.0, 0.0, 255.0).to(torch.uint8)