Update: Add original model implementation

Dockerfile

@@ -17,20 +17,32 @@ RUN pip install torch==2.0.0 torchvision==0.15.1 --extra-index-url https://downl
 # Install CLIP
 RUN pip install git+https://github.com/openai/CLIP.git
 
+# Install diffusers and other dependencies
+RUN pip install diffusers transformers accelerate xformers omegaconf einops kornia
+
 # Install cairosvg and other dependencies
 RUN pip install cairosvg cairocffi cssselect2 defusedxml tinycss2
 
 # Install FastAPI and other dependencies
 RUN pip install fastapi uvicorn pydantic pillow numpy requests
 
+# Install SVG dependencies
+RUN pip install svgwrite svgpathtools cssutils numba
+
 # Copy the model files
 COPY . /code/
 
+# Download model weights if they don't exist
+RUN if [ ! -f /code/ViT-B-32.pt ]; then \
+    pip install gdown && \
+    python -c "import clip; clip.load('ViT-B-32')" ; \
+    fi
+
 # Make sure the handler and model are available
-RUN if [ -f /code/simplified_diffsketcher.py ]; then \
-    echo "Simplified DiffSketcher found"; \
+RUN if [ -f /code/diffsketcher_endpoint.py ]; then \
+    echo "DiffSketcher endpoint found"; \
     else \
-    echo "Simplified DiffSketcher not found, using placeholder"; \
+    echo "DiffSketcher endpoint not found, using placeholder"; \
     fi
 
 # Set environment variables

handler.py

@@ -15,21 +15,12 @@ except ImportError:
     subprocess.check_call(["pip", "install", "cairosvg", "cairocffi", "cssselect2", "defusedxml", "tinycss2"])
     import cairosvg
 
-# Safely import clip with fallback
+# Import the DiffSketcher endpoint
 try:
-    import clip
+    from diffsketcher_endpoint import DiffSketcherEndpoint
 except ImportError:
-    print("Warning: clip not found. Installing...")
-    import subprocess
-    subprocess.check_call(["pip", "install", "git+https://github.com/openai/CLIP.git"])
-    import clip
-
-# Import the versatile SVG generator
-try:
-    from versatile_svg_generator import VersatileSVGGenerator
-except ImportError:
-    print("Warning: versatile_svg_generator not found. Using placeholder.")
-    VersatileSVGGenerator = None
+    print("Warning: diffsketcher_endpoint not found. Using placeholder.")
+    DiffSketcherEndpoint = None
 
 class EndpointHandler:
     def __init__(self, model_dir):
@@ -38,14 +29,14 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Initializing model on device: {self.device}")
         
-        # Initialize the versatile SVG generator if available
-        if VersatileSVGGenerator is not None:
+        # Initialize the DiffSketcher endpoint if available
+        if DiffSketcherEndpoint is not None:
             try:
-                self.model = VersatileSVGGenerator(model_dir)
+                self.model = DiffSketcherEndpoint(model_dir)
                 self.use_model = True
-                print("Versatile SVG generator initialized successfully")
+                print("DiffSketcher endpoint initialized successfully")
             except Exception as e:
-                print(f"Error initializing versatile SVG generator: {e}")
+                print(f"Error initializing DiffSketcher endpoint: {e}")
                 self.use_model = False
         else:
             self.use_model = False
@@ -76,11 +67,11 @@ class EndpointHandler:
             # Generate SVG using the model or placeholder
             if self.use_model:
                 try:
-                    # Use the versatile SVG generator
+                    # Use the DiffSketcher endpoint
                     result = self.model(prompt)
                     image = result["image"]
                 except Exception as e:
-                    print(f"Error using versatile SVG generator: {e}")
+                    print(f"Error using DiffSketcher endpoint: {e}")
                     # Fall back to placeholder
                     svg_content = self.generate_placeholder_svg(prompt)
                     png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))