Add project files, removing large SVG

app.py

@@ -1,154 +1,210 @@
-import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+import os
 import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+import vtracer
+import tempfile
+import cairosvg
+import re
+from PIL import Image
+from datetime import datetime
+
+from flask import Flask, request, jsonify, send_from_directory
+from flask_cors import CORS
+
+from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
+
+import torchvision.transforms as transforms
+from model import Generator 
+
+def setup_directories():
+    os.makedirs(SVG_DIR, exist_ok=True)
+    os.makedirs(THUMBNAIL_DIR, exist_ok=True)
+    print(f"Directories '{SVG_DIR}' and '{THUMBNAIL_DIR}' are ready.")
+
+def sanitize_filename(prompt):
+    """Removes characters that are invalid for filenames."""
+
+    s = re.sub(r'[\\/*?:"<>|]', "", prompt)
+
+    return s[:100]
+
+SVG_DIR = os.path.join(os.getcwd(), 'generated_svgs')
+THUMBNAIL_DIR = os.path.join(os.getcwd(), 'thumbnails')
+SKETCH_MODEL_WEIGHTS = 'checkpoints/netG_A_latest.pth'
+
+class ImageToSvgPipeline:
+    """
+    A class to handle the entire pipeline from text prompt to SVG.
+    Initializes models once to be reused.
+    """
+    def __init__(self, sketch_model_path: str):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+
+        self._initialize_rinna_model()
+        self._initialize_sketch_model(sketch_model_path)
+
+    def _initialize_rinna_model(self):
+        print("Loading Rinna Stable Diffusion model...")
+        model_id = "rinna/japanese-stable-diffusion"
+
+        self.rinna_pipe = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+        )
+        self.rinna_pipe.scheduler = LMSDiscreteScheduler(
+            beta_start=0.00085, beta_end=0.012, 
+            beta_schedule="scaled_linear", num_train_timesteps=1000
+        )
+        self.rinna_pipe.tokenizer.model_max_length = 77
+        self.rinna_pipe.to(self.device)
+        print("Rinna model loaded.")
+
+    def _initialize_sketch_model(self, model_path: str):
+        print(f"Loading Sketch Generator model from {model_path}...")
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Sketch model weights not found at: {model_path}")
+
+        self.sketch_model = Generator(input_nc=3, output_nc=1, n_residual_blocks=3)
+        self.sketch_model.to(self.device)
+        self.sketch_model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.sketch_model.eval()
+
+        self.sketch_transform = transforms.Compose([
+            transforms.ToTensor(),
+        ])
+        print("Sketch model loaded.")
+
+    def _generate_image(self, prompt: str, negative_prompt: str, steps: int = 30) -> Image.Image:
+        print(f"Generating image for prompt: '{prompt}'")
+        with torch.no_grad():
+            image = self.rinna_pipe(
+                prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=7.5,
+                width=512,
+                height=512,
+            ).images[0]
+        return image
+
+    def _convert_to_sketch(self, image: Image.Image) -> Image.Image:
+        print("Converting image to sketch...")
+        with torch.no_grad():
+            input_tensor = self.sketch_transform(image.convert("RGB")).unsqueeze(0).to(self.device)
+            output_tensor = self.sketch_model(input_tensor)
+            output_tensor = output_tensor.squeeze(0).cpu()
+            sketch_image = transforms.ToPILImage()(output_tensor)
+        return sketch_image
+
+    def _extract_svg(self, image: Image.Image) -> str:
+        print("Extracting SVG from sketch...")
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            image.save(tmp_file.name)
+            tmp_path = tmp_file.name
+
+        try:
+            svg_output_path = tmp_path.replace(".png", ".svg")
+            vtracer.convert_image_to_svg_py(tmp_path, svg_output_path)
+
+            with open(svg_output_path, 'r', encoding='utf-8') as f:
+                svg_data = f.read()
+        finally:
+            if os.path.exists(tmp_path): os.remove(tmp_path)
+            if 'svg_output_path' in locals() and os.path.exists(svg_output_path): os.remove(svg_output_path)
+
+        print("SVG extraction complete.")
+        return svg_data
+
+    def process(self, prompt: str, negative_prompt: str) -> str:
+        generated_image = self._generate_image(prompt, negative_prompt)
+        sketch_image = self._convert_to_sketch(generated_image)
+        svg_content = self._extract_svg(sketch_image)
+        return svg_content
+
+app = Flask(__name__)
+
+CORS(app, resources={r"/*": {"origins": "*"}})
+
+pipeline = ImageToSvgPipeline(sketch_model_path=SKETCH_MODEL_WEIGHTS)
+
+def sanitize_filename(text):
+    text = re.sub(r'[\\/*?:"<>|]', "", text)
+    return text.strip()
+
+@app.route('/generate', methods=['POST'])
+def generate_svg():
+    data = request.json
+    prompt = data.get('prompt')
+    if not prompt: return jsonify({"error": "Prompt is required"}), 400
+
+    negative_prompt = "低品質、最悪の品質、下手な手、指が6本、指が4本、奇形、醜い、ぼやけている、ぼやけた、ウォーターマーク、署名、テキスト"
+    try:
+        svg_result = pipeline.process(prompt, negative_prompt)
+
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+        safe_prompt = sanitize_filename(prompt)[:50]
+        filename = f"{timestamp}_{safe_prompt}.svg"
+
+        svg_path = os.path.join(SVG_DIR, filename)
+        with open(svg_path, 'w', encoding='utf-8') as f:
+            f.write(svg_result)
+
+        thumbnail_path = os.path.join(THUMBNAIL_DIR, filename.replace('.svg', '.png'))
+        cairosvg.svg2png(bytestring=svg_result.encode('utf-8'), write_to=thumbnail_path, output_width=256, output_height=256)
+
+        return svg_result, 200, {'Content-Type': 'image/svg+xml'}
+    except Exception as e:
+        print(f"An error occurred during generation: {e}")
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/gallery', methods=['GET'])
+def get_gallery():
+    try:
+        page = int(request.args.get('page', 1))
+        limit = int(request.args.get('limit', 8))
+
+        svg_files = sorted([f for f in os.listdir(SVG_DIR) if f.endswith('.svg')], reverse=True)
+
+        start_index = (page - 1) * limit
+        end_index = start_index + limit
+        paginated_files = svg_files[start_index:end_index]
+
+        drawings = []
+        for filename in paginated_files:
+            prompt_match = re.match(r"\d+_(.+)\.svg", filename)
+            prompt = prompt_match.group(1).replace('_', ' ') if prompt_match else "Prompt not found"
+            drawings.append({
+                "filename": filename,
+                "thumbnail": f"/thumbnails/{filename.replace('.svg', '.png')}",
+                "prompt": prompt
+            })
+
+        has_more = end_index < len(svg_files)
+        return jsonify({"drawings": drawings, "hasMore": has_more})
+    except Exception as e:
+        print(f"Error fetching gallery: {e}")
+        return jsonify({"error": "Failed to fetch gallery"}), 500
+
+@app.route('/svgs/<path:filename>')
+def get_svg(filename):
+    return send_from_directory(SVG_DIR, filename)
+
+@app.route('/thumbnails/<path:filename>')
+def get_thumbnail(filename):
+    return send_from_directory(THUMBNAIL_DIR, filename)
+
+@app.route('/drawings/<path:filename>', methods=['DELETE'])
+def delete_drawing_file(filename):
+    try:
+        svg_path = os.path.join(SVG_DIR, filename)
+        thumb_path = os.path.join(THUMBNAIL_DIR, filename.replace('.svg', '.png'))
+        if os.path.exists(svg_path): os.remove(svg_path)
+        if os.path.exists(thumb_path): os.remove(thumb_path)
+        return jsonify({"message": f"Successfully deleted {filename}"})
+    except Exception as e:
+        print(f"Error deleting file: {e}")
+        return jsonify({"error": "Failed to delete file"}), 500
+
+if __name__ == '__main__':
+    print("Starting Flask server...")
+    app.run(host='0.0.0.0', port=5000)

checkpoints/netG_A_latest.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c686ced2a666b4850b4bb6ccf0748031c3eda9f822de73a34b8979970d90f0c6
+size 17173511

requirements.txt

@@ -1,6 +1,9 @@
-accelerate
-diffusers
-invisible_watermark
+flask
+flask-cors
 torch
+diffusers
 transformers
-xformers
+accelerate
+vtracer
+pillow
+cairosvg