Update app.py

app.py

@@ -30,14 +30,6 @@ class SketchToRealisticFace:
         #     requires_safety_checker=False
         # )
 
-        # Load Realistic LoRA model
-        # self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
-        #     "SG161222/Realistic_Vision_V6.0_B1_noVAE",
-        #     controlnet=self.controlnet,
-        #     torch_dtype=torch.float16,
-        #     safety_checker=None,
-        #     requires_safety_checker=False
-        # )
         self.pipe = StableDiffusionControlNetPipeline.from_single_file(
             "https://huggingface.co/NikhilJoson/Realistic_Vision_V6_B1/blob/main/realisticVisionV60B1_v51HyperVAE.safetensors",
             controlnet=self.controlnet,
@@ -106,6 +98,12 @@ print("Loading model... This may take a few minutes.")
 generator = SketchToRealisticFace()
 print("Model loaded successfully!")
 
+def pil_to_base64(image):
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/png;base64,{img_str}"
+
 @spaces.GPU
 def generate_face(sketch_image, custom_prompt, seed):
     """Wrapper function for Gradio interface"""
@@ -116,7 +114,12 @@ def generate_face(sketch_image, custom_prompt, seed):
         # Generate the realistic face
         realistic_face, control_image = generator.generate(sketch_image=sketch_image,custom_prompt=custom_prompt,seed=seed_int)
         
-        return realistic_face, control_image
+        #return realistic_face, control_image
+
+        # Convert to base64 (uncomment if you want base64 strings)
+        realistic_b64 = pil_to_base64(realistic_face)
+        control_b64 = pil_to_base64(control_image)
+        return realistic_b64, control_b64
         
     except Exception as e:
         print(f"Error: {str(e)}")