Update SegCloth.py

SegCloth.py

@@ -1,14 +1,19 @@
+import torch
 from transformers import pipeline
+from diffusers import StableDiffusionPipeline
 from PIL import Image
 import numpy as np
 from io import BytesIO
 import base64
-import torch
-from torchvision.transforms.functional import to_pil_image
 
-# Initialize segmentation and super-resolution pipelines
+# Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
-super_res = pipeline("image-super-resolution", model="CompVis/stable-diffusion-v1-4")
+
+# Initialize Stable Diffusion pipeline
+model_id = "CompVis/stable-diffusion-v1-4"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+pipe = pipe.to(device)
 
 def encode_image_to_base64(image):
     buffered = BytesIO()
@@ -36,9 +41,12 @@ def segment_and_enhance_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt",
             final_mask = Image.fromarray(current_mask)
             resized_img.putalpha(final_mask)
 
-            # Enhance image using super-resolution
-            enhanced_img = super_res(resized_img)
-            
+            # Generate prompt for Stable Diffusion
+            prompt = f"a clear photo of {s['label']}"
+
+            # Enhance image using Stable Diffusion
+            enhanced_img = pipe(prompt).images[0]
+
             # Convert the final image to base64
             imageBase64 = encode_image_to_base64(enhanced_img)
             result_images.append((s['label'], imageBase64))