Update app.py

app.py

@@ -1,23 +1,67 @@
 import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
+from torchvision import transforms
+from PIL import Image
+import matplotlib.pyplot as plt
+from huggingface_hub import hf_hub_download
 
-# Fine-tuned Stable Diffusion model
+# 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
+vgg16_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_vgg", filename="vgg16_model.pth")
+vgg16 = torch.load(vgg16_model_path)
+vgg16.eval()
+
+# Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+vgg16.to(device)
 pipe.to(device)
 
+# Prediction function for the VGG16 model
+def predict_and_show(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    image_tensor = transform(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = vgg16(image_tensor)
+        _, predicted = torch.max(outputs.data, 1)
+
+    class_names = ["Class1", "Class2", "Class3"]  # Replace with your actual class names
+    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
 def generate_image(prompt):
     image = pipe(prompt).images[0]
     return image
 
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(lines=2, placeholder="Enter prompt here..."),
-    outputs="image",
-    title="Room Styler",
-    description="Generate room style images with Stable Diffusion."
-)
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Room Style Recognition and Generation")
+
+    with gr.Tab("Recognize Room Style"):
+        image_input = gr.Image(type="pil")
+        label_output = gr.Textbox()
+        btn_predict = gr.Button("Predict Style")
+        btn_predict.click(predict_and_show, inputs=image_input, outputs=label_output)
+
+    with gr.Tab("Generate Room Style"):
+        text_input = gr.Textbox(placeholder="Enter a prompt for room style...")
+        image_output = gr.Image()
+        btn_generate = gr.Button("Generate Image")
+        btn_generate.click(generate_image, inputs=text_input, outputs=image_output)
 
-iface.launch()
+demo.launch()