Update app.py

app.py

@@ -1,23 +1,32 @@
 import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
-from torchvision import transforms
+from torchvision import models, transforms
 from PIL import Image
-import matplotlib.pyplot as plt
 from huggingface_hub import hf_hub_download
+import torch.nn as nn
+import torch.optim as optim
 
-# Fine-tuned Stable Diffusion model from your Hugging Face repository
-model_id = "sk2003/room-styler"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-
-# VGG16 model
+# LoadING the VGG16 model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 vgg16_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_vgg", filename="vgg16_model.pth")
-vgg16 = torch.load(vgg16_model_path)
+
+vgg16 = models.vgg16(pretrained=True)
+for param in vgg16.parameters():
+    param.requires_grad = False
+
+num_classes = 8 
+vgg16.classifier[6] = nn.Linear(vgg16.classifier[6].in_features, num_classes)
+vgg16 = vgg16.to(device)
+
+# Loading the saved state dict
+checkpoint = torch.load(vgg16_model_path, map_location=device)
+vgg16.load_state_dict(checkpoint['model_state_dict'])
 vgg16.eval()
 
-# Device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-vgg16.to(device)
+# Fine-tuned Stable Diffusion model from your Hugging Face repository
+model_id = "sk2003/room-styler"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe.to(device)
 
 # Prediction function for the VGG16 model
@@ -33,14 +42,9 @@ def predict_and_show(image):
         outputs = vgg16(image_tensor)
         _, predicted = torch.max(outputs.data, 1)
 
-    class_names = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"] 
+    class_names = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"]  
     predicted_label = class_names[predicted.item()]
 
-    plt.imshow(image)
-    plt.title(f'Predicted: {predicted_label}')
-    plt.axis('off')
-    plt.show()
-
     return predicted_label
 
 # Generation function for the Stable Diffusion model