Update app.py

app.py

@@ -65,7 +65,110 @@ class HistoryManager:
 # Initialize history manager
 history_manager = HistoryManager()
 
-[Previous model and generator code remains the same...]
+norm_layer = nn.InstanceNorm2d
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+        
+        conv_block = [  nn.ReflectionPad2d(1),
+                        nn.Conv2d(in_features, in_features, 3),
+                        norm_layer(in_features),
+                        nn.ReLU(inplace=True),
+                        nn.ReflectionPad2d(1),
+                        nn.Conv2d(in_features, in_features, 3),
+                        norm_layer(in_features)  ]
+        
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
+        super(Generator, self).__init__()
+        
+        # Initial convolution block
+        model0 = [   nn.ReflectionPad2d(3),
+                    nn.Conv2d(input_nc, 64, 7),
+                    norm_layer(64),
+                    nn.ReLU(inplace=True) ]
+        self.model0 = nn.Sequential(*model0)
+
+        # Downsampling
+        model1 = []
+        in_features = 64
+        out_features = in_features*2
+        for _ in range(2):
+            model1 += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                        norm_layer(out_features),
+                        nn.ReLU(inplace=True) ]
+            in_features = out_features
+            out_features = in_features*2
+        self.model1 = nn.Sequential(*model1)
+
+        # Residual blocks
+        model2 = []
+        for _ in range(n_residual_blocks):
+            model2 += [ResidualBlock(in_features)]
+        self.model2 = nn.Sequential(*model2)
+
+        # Upsampling
+        model3 = []
+        out_features = in_features//2
+        for _ in range(2):
+            model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+                        norm_layer(out_features),
+                        nn.ReLU(inplace=True) ]
+            in_features = out_features
+            out_features = in_features//2
+        self.model3 = nn.Sequential(*model3)
+
+        # Output layer
+        model4 = [  nn.ReflectionPad2d(3),
+                    nn.Conv2d(64, output_nc, 7)]
+        if sigmoid:
+            model4 += [nn.Sigmoid()]
+            
+        self.model4 = nn.Sequential(*model4)
+
+    def forward(self, x):
+        out = self.model0(x)
+        out = self.model1(out)
+        out = self.model2(out)
+        out = self.model3(out)
+        out = self.model4(out)
+        return out
+
+# Initialize models
+def load_models():
+    try:
+        print("Initializing models in CPU mode...")
+        model1 = Generator(3, 1, 3)
+        model2 = Generator(3, 1, 3)
+        
+        model1.load_state_dict(torch.load('model.pth', map_location='cpu'))
+        model2.load_state_dict(torch.load('model2.pth', map_location='cpu'))
+        
+        model1.eval()
+        model2.eval()
+        torch.set_grad_enabled(False)
+        
+        print("Models loaded successfully in CPU mode")
+        return model1, model2
+    except Exception as e:
+        error_msg = f"Error loading models: {str(e)}"
+        print(error_msg)
+        raise gr.Error("Failed to initialize models. Please check the model files and system configuration.")
+
+# Load models
+try:
+    print("Starting model initialization...")
+    model1, model2 = load_models()
+    print("Model initialization completed")
+except Exception as e:
+    print(f"Critical error during model initialization: {str(e)}")
+    raise gr.Error("Failed to start the application due to model initialization error.")
 
 def apply_preset(preset_name):
     """Apply a style preset and return the settings"""
@@ -78,29 +181,12 @@ def apply_preset(preset_name):
         )
     return (1.0, 1.0, 1.0, False)
 
-def save_image_with_metadata(image, output_path, settings):
-    """Save image with processing metadata"""
-    try:
-        # Save image
-        image.save(output_path)
-        
-        # Save metadata
-        metadata_path = output_path + ".json"
-        with open(metadata_path, 'w') as f:
-            json.dump({
-                "processing_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                "settings": settings
-            }, f)
-    except Exception as e:
-        print(f"Error saving image metadata: {e}")
-
-def get_image_download_link(image):
-    """Create a download link for the processed image"""
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    href = f'data:image/png;base64,{img_str}'
-    return href
+def enhance_lines(img, contrast=1.0, brightness=1.0):
+    """Enhance line drawing with contrast and brightness adjustments"""
+    enhanced = np.array(img)
+    enhanced = enhanced * contrast
+    enhanced = np.clip(enhanced + brightness, 0, 1)
+    return Image.fromarray((enhanced * 255).astype(np.uint8))
 
 def predict(input_img, version, preset_name, line_thickness=1.0, contrast=1.0, 
            brightness=1.0, enable_enhancement=False, output_size="Original"):
@@ -166,7 +252,7 @@ def predict(input_img, version, preset_name, line_thickness=1.0, contrast=1.0,
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Extended custom CSS
+# Custom CSS
 custom_css = """
 .gradio-container {
     font-family: 'Helvetica Neue', Arial, sans-serif;
@@ -208,7 +294,7 @@ custom_css = """
 }
 """
 
-# Create Gradio interface with enhanced UI
+# Create Gradio interface
 with gr.Blocks(css=custom_css) as iface:
     with gr.Row(elem_classes="gr-header"):
         gr.Markdown("# 🎨 Advanced Line Drawing Generator")