Upload Lora app

.gitignore

@@ -0,0 +1,3 @@
+**/__pycache__/
+/.venv/
+/scripts/

app.py

@@ -0,0 +1,211 @@
+import gradio as gr
+import json
+import sys
+import io
+import subprocess
+import tempfile
+from pathlib import Path
+from safetensors_worker import PrintMetadata
+
+class Context:
+    def __init__(self):
+        self.obj = {'quiet': True, 'parse_more': True}
+
+ctx = Context()
+
+def debug_log(message: str):
+    print(f"[DEBUG] {message}")
+
+def load_metadata(file_path: str) -> tuple:
+    try:
+        debug_log(f"Loading file: {file_path}")
+        
+        if not file_path:
+            return {"status": "Awaiting input"}, {}, "", "", ""
+
+        old_stdout = sys.stdout
+        sys.stdout = buffer = io.StringIO()
+        exit_code = PrintMetadata(ctx.obj, file_path.name)
+        sys.stdout = old_stdout
+        
+        metadata_str = buffer.getvalue().strip()
+        
+        if exit_code != 0:
+            error_msg = f"Error code {exit_code}"
+            return {"error": error_msg}, {}, "", error_msg, ""
+
+        try:
+            full_metadata = json.loads(metadata_str)
+        except json.JSONDecodeError:
+            error_msg = "Invalid metadata structure"
+            return {"error": error_msg}, {}, "", error_msg, ""
+
+        training_params = full_metadata.get("__metadata__", {})
+        key_metrics = {
+            key: training_params.get(key, "N/A")
+            for key in [
+                "ss_optimizer", "ss_num_epochs", "ss_unet_lr",
+                "ss_text_encoder_lr", "ss_steps"
+            ]
+        }
+        
+        return full_metadata, key_metrics, json.dumps(full_metadata, indent=2), "", file_path.name
+    
+    except Exception as e:
+        return {"error": str(e)}, {}, "", str(e), ""
+
+def validate_json(edited_json: str) -> tuple:
+    try:
+        return True, json.loads(edited_json), ""
+    except Exception as e:
+        return False, None, str(e)
+
+def update_metadata(edited_json: str) -> tuple:
+    try:
+        modified_data = json.loads(edited_json)
+        metadata = modified_data.get("__metadata__", {})
+        
+        key_fields = {
+            param: metadata.get(param, "N/A")
+            for param in [
+                "ss_optimizer", "ss_num_epochs", "ss_unet_lr",
+                "ss_text_encoder_lr", "ss_steps"
+            ]
+        }
+        return key_fields, modified_data, ""
+    except:
+        return gr.update(), gr.update(), ""
+
+def save_metadata(edited_json: str, source_file: str, output_name: str) -> tuple:
+    debug_log("Initiating save process")
+    try:
+        if not source_file:
+            return None, "No source file provided"
+
+        is_valid, parsed_data, error = validate_json(edited_json)
+        if not is_valid:
+            return None, f"Validation error: {error}"
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp:
+            json.dump(parsed_data, tmp, indent=2)
+            temp_path = tmp.name
+
+        source_path = Path(source_file)
+        
+        if output_name.strip():
+            base_name = output_name.strip()
+            if not base_name.endswith(".safetensors"):
+                base_name += ".safetensors"
+        else:
+            base_name = f"{source_path.stem}_modified.safetensors"
+        
+        output_path = Path(base_name)
+        version = 1
+        while output_path.exists():
+            output_path = Path(f"{source_path.stem}_modified_{version}.safetensors")
+            version += 1
+
+        cmd = [
+            sys.executable,
+            "safetensors_util.py",
+            "writemd",
+            source_file,
+            temp_path,
+            str(output_path),
+            "-f"
+        ]
+
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False
+        )
+
+        Path(temp_path).unlink(missing_ok=True)
+
+        if result.returncode != 0:
+            error_msg = f"Save failure: {result.stderr}"
+            return None, error_msg
+
+        return str(output_path), ""
+    
+    except Exception as e:
+        return None, f"Critical error: {str(e)}"
+
+def create_interface():
+    with gr.Blocks(title="LoRA Metadata Editor") as app:
+        gr.Markdown("# LoRA Metadata Editor")
+        
+        with gr.Tabs():
+            with gr.Tab("Metdata Viewer"):
+                gr.Markdown("### LoRa Upload")
+                file_input = gr.File(
+                    file_types=[".safetensors"],
+                    show_label=False
+                )
+                
+                with gr.Row():
+                    with gr.Column():
+                        gr.Markdown("### Full Metadata")
+                        full_viewer = gr.JSON(show_label=False)
+                    
+                    with gr.Column():
+                        gr.Markdown("### Key Metrics")
+                        key_viewer = gr.JSON(show_label=False)
+
+            with gr.Tab("Edit Metadata"):
+                with gr.Row():
+                    with gr.Column():
+                        gr.Markdown("### JSON Workspace")
+                        metadata_editor = gr.Textbox(
+                            lines=25,
+                            show_label=False,
+                            placeholder="Edit metadata JSON here"
+                        )
+                        gr.Markdown("### Output Name")
+                        filename_input = gr.Textbox(
+                            placeholder="Leave empty for auto-naming",
+                            show_label=False
+                        )
+                    
+                    with gr.Column():
+                        gr.Markdown("### Live Preview")
+                        modified_viewer = gr.JSON(show_label=False)
+                        save_btn = gr.Button("💾 Save Metadata", variant="primary")
+                        gr.Markdown("### Download Modified LoRa") 
+                        output_file = gr.File(
+                            visible=False,
+                            show_label=False
+                        )
+
+        status_display = gr.HTML(visible=False)
+        source_tracker = gr.State()
+
+        file_input.upload(
+            load_metadata,
+            inputs=file_input,
+            outputs=[full_viewer, key_viewer, metadata_editor, status_display, source_tracker]
+        )
+
+        metadata_editor.change(
+            update_metadata,
+            inputs=metadata_editor,
+            outputs=[key_viewer, modified_viewer, status_display]
+        )
+
+        save_btn.click(
+            save_metadata,
+            inputs=[metadata_editor, source_tracker, filename_input],
+            outputs=[output_file, status_display],
+        ).then(
+            lambda x: gr.File(value=x, visible=True),
+            inputs=output_file,
+            outputs=output_file
+        )
+
+    return app
+
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch()

safetensors_file.py

@@ -0,0 +1,125 @@
+import os, sys, json
+
+class SafeTensorsException(Exception):
+    def __init__(self, msg:str):
+        self.msg=msg
+        super().__init__(msg)
+
+    @staticmethod
+    def invalid_file(filename:str,whatiswrong:str):
+        s=f"{filename} is not a valid .safetensors file: {whatiswrong}"
+        return SafeTensorsException(msg=s)
+
+    def __str__(self):
+        return self.msg
+
+class SafeTensorsChunk:
+    def __init__(self,name:str,dtype:str,shape:list[int],offset0:int,offset1:int):
+        self.name=name
+        self.dtype=dtype
+        self.shape=shape
+        self.offset0=offset0
+        self.offset1=offset1
+
+class SafeTensorsFile:
+    def __init__(self):
+        self.f=None         #file handle
+        self.hdrbuf=None    #header byte buffer
+        self.header=None    #parsed header as a dict
+        self.error=0
+
+    def __del__(self):
+        self.close_file()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close_file()
+
+    def close_file(self):
+        if self.f is not None:
+            self.f.close()
+            self.f=None
+            self.filename=""
+
+    #test file: duplicate_keys_in_header.safetensors
+    def _CheckDuplicateHeaderKeys(self):
+        def parse_object_pairs(pairs):
+            return [k for k,_ in pairs]
+
+        keys=json.loads(self.hdrbuf,object_pairs_hook=parse_object_pairs)
+        #print(keys)
+        d={}
+        for k in keys:
+            if k in d: d[k]=d[k]+1
+            else: d[k]=1
+        hasError=False
+        for k,v in d.items():
+            if v>1:
+                print(f"key {k} used {v} times in header",file=sys.stderr)
+                hasError=True
+        if hasError:
+            raise SafeTensorsException.invalid_file(self.filename,"duplicate keys in header")
+
+    @staticmethod
+    def open_file(filename:str,quiet=False,parseHeader=True):
+        s=SafeTensorsFile()
+        s.open(filename,quiet,parseHeader)
+        return s
+
+    def open(self,fn:str,quiet=False,parseHeader=True)->int:
+        st=os.stat(fn)
+        if st.st_size<8: #test file: zero_len_file.safetensors
+            raise SafeTensorsException.invalid_file(fn,"length less than 8 bytes")
+
+        f=open(fn,"rb")
+        b8=f.read(8) #read header size
+        if len(b8)!=8:
+            raise SafeTensorsException.invalid_file(fn,f"read only {len(b8)} bytes at start of file")
+        headerlen=int.from_bytes(b8,'little',signed=False)
+
+        if (8+headerlen>st.st_size): #test file: header_size_too_big.safetensors
+            raise SafeTensorsException.invalid_file(fn,"header extends past end of file")
+
+        if quiet==False:
+            print(f"{fn}: length={st.st_size}, header length={headerlen}")
+        hdrbuf=f.read(headerlen)
+        if len(hdrbuf)!=headerlen:
+            raise SafeTensorsException.invalid_file(fn,f"header size is {headerlen}, but read {len(hdrbuf)} bytes")
+        self.filename=fn
+        self.f=f
+        self.st=st
+        self.hdrbuf=hdrbuf
+        self.error=0
+        self.headerlen=headerlen
+        if parseHeader==True:
+            self._CheckDuplicateHeaderKeys()
+            self.header=json.loads(self.hdrbuf)
+        return 0
+
+    def get_header(self):
+        return self.header
+
+    def load_one_tensor(self,tensor_name:str):
+        self.get_header()
+        if tensor_name not in self.header: return None
+
+        t=self.header[tensor_name]
+        self.f.seek(8+self.headerlen+t['data_offsets'][0])
+        bytesToRead=t['data_offsets'][1]-t['data_offsets'][0]
+        bytes=self.f.read(bytesToRead)
+        if len(bytes)!=bytesToRead:
+            print(f"{tensor_name}: length={bytesToRead}, only read {len(bytes)} bytes",file=sys.stderr)
+        return bytes
+
+    def copy_data_to_file(self,file_handle) -> int:
+
+        self.f.seek(8+self.headerlen)
+        bytesLeft:int=self.st.st_size - 8 - self.headerlen
+        while bytesLeft>0:
+            chunklen:int=min(bytesLeft,int(16*1024*1024)) #copy in blocks of 16 MB
+            file_handle.write(self.f.read(chunklen))
+            bytesLeft-=chunklen
+
+        return 0

safetensors_util.py

@@ -0,0 +1,98 @@
+import sys, click
+
+import safetensors_worker
+# This file deals with command line only. If the command line is parsed successfully,
+# we will call one of the functions in safetensors_worker.py.
+
+readonly_input_file=click.argument("input_file", metavar='input_file',
+                                   type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True))
+output_file=click.argument("output_file", metavar='output_file',
+                            type=click.Path(file_okay=True, dir_okay=False, writable=True))
+
+force_overwrite_flag=click.option("-f","--force-overwrite",default=False,is_flag=True, show_default=True,
+                                  help="overwrite existing files")
+fix_ued_flag=click.option("-pm","--parse-more",default=False,is_flag=True, show_default=True,
+                          help="when printing metadata, unescaped doublequotes to make text more readable" )
+quiet_flag=click.option("-q","--quiet",default=False,is_flag=True, show_default=True,
+                        help="Quiet mode, don't print informational stuff" )
+
+@click.group()
+@click.version_option(version=7)
+@quiet_flag
+
+@click.pass_context
+def cli(ctx,quiet:bool):
+    # ensure that ctx.obj exists and is a dict (in case `cli()` is called
+    # by means other than the `if` block below)
+    ctx.ensure_object(dict)
+    ctx.obj['quiet'] = quiet
+
+
+@cli.command(name="header",short_help="print file header")
+@readonly_input_file
+@click.pass_context
+def cmd_header(ctx,input_file:str) -> int:
+    sys.exit( safetensors_worker.PrintHeader(ctx.obj,input_file) )
+
+
+@cli.command(name="metadata",short_help="print only __metadata__ in file header")
+@readonly_input_file
+@fix_ued_flag
+@click.pass_context
+def cmd_meta(ctx,input_file:str,parse_more:bool)->int:
+    ctx.obj['parse_more'] = parse_more
+    sys.exit( safetensors_worker.PrintMetadata(ctx.obj,input_file) )
+
+
+@cli.command(name="listkeys",short_help="print header key names (except __metadata__) as a Python list")
+@readonly_input_file
+@click.pass_context
+def cmd_keyspy(ctx,input_file:str) -> int:
+    sys.exit( safetensors_worker.HeaderKeysToLists(ctx.obj,input_file) )
+
+
+@cli.command(name="writemd",short_help="read __metadata__ from json and write to safetensors file")
+@click.argument("in_st_file", metavar='input_st_file',
+                 type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True))
+@click.argument("in_json_file", metavar='input_json_file',
+                 type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True))
+@output_file
+@force_overwrite_flag
+@click.pass_context
+def cmd_writemd(ctx,in_st_file:str,in_json_file:str,output_file:str,force_overwrite:bool) -> int:
+    """Read "__metadata__" from json file and write to safetensors header"""
+    ctx.obj['force_overwrite'] = force_overwrite
+    sys.exit( safetensors_worker.WriteMetadataToHeader(ctx.obj,in_st_file,in_json_file,output_file) )
+
+
+@cli.command(name="extracthdr",short_help="extract file header and save to output file")
+@readonly_input_file
+@output_file
+@force_overwrite_flag
+@click.pass_context
+def cmd_extractheader(ctx,input_file:str,output_file:str,force_overwrite:bool) -> int:
+    ctx.obj['force_overwrite'] = force_overwrite
+    sys.exit( safetensors_worker.ExtractHeader(ctx.obj,input_file,output_file) )
+
+
+@cli.command(name="extractdata",short_help="extract one tensor and save to file")
+@readonly_input_file
+@click.argument("key_name", metavar='key_name',type=click.STRING)
+@output_file
+@force_overwrite_flag
+@click.pass_context
+def cmd_extractheader(ctx,input_file:str,key_name:str,output_file:str,force_overwrite:bool) -> int:
+    ctx.obj['force_overwrite'] = force_overwrite
+    sys.exit( safetensors_worker.ExtractData(ctx.obj,input_file,key_name,output_file) )
+
+
+@cli.command(name="checklora",short_help="see if input file is a SD 1.x LoRA file")
+@readonly_input_file
+@click.pass_context
+def cmd_checklora(ctx,input_file:str)->int:
+    sys.exit( safetensors_worker.CheckLoRA(ctx.obj,input_file) )
+
+
+if __name__ == '__main__':
+    sys.stdout.reconfigure(encoding='utf-8')
+    cli(obj={},max_content_width=96)

safetensors_worker.py

@@ -0,0 +1,243 @@
+import os, sys, json
+from safetensors_file import SafeTensorsFile
+
+def _need_force_overwrite(output_file:str,cmdLine:dict) -> bool:
+    if cmdLine["force_overwrite"]==False:
+        if os.path.exists(output_file):
+            print(f'output file "{output_file}" already exists, use -f flag to force overwrite',file=sys.stderr)
+            return True
+    return False
+
+def WriteMetadataToHeader(cmdLine:dict,in_st_file:str,in_json_file:str,output_file:str) -> int:
+    if _need_force_overwrite(output_file,cmdLine): return -1
+
+    with open(in_json_file,"rt") as f:
+        inmeta=json.load(f)
+    if not "__metadata__" in inmeta:
+        print(f"file {in_json_file} does not contain a top-level __metadata__ item",file=sys.stderr)
+        #json.dump(inmeta,fp=sys.stdout,indent=2)
+        return -2
+    inmeta=inmeta["__metadata__"] #keep only metadata
+    #json.dump(inmeta,fp=sys.stdout,indent=2)
+
+    s=SafeTensorsFile.open_file(in_st_file)
+    js=s.get_header()
+
+    if inmeta==[]:
+        js.pop("__metadata__",0)
+        print("loaded __metadata__ is an empty list, output file will not contain __metadata__ in header")
+    else:
+        print("adding __metadata__ to header:")
+        json.dump(inmeta,fp=sys.stdout,indent=2)
+        if isinstance(inmeta,dict):
+            for k in inmeta:
+                inmeta[k]=str(inmeta[k])
+        else:
+            inmeta=str(inmeta)
+        #js["__metadata__"]=json.dumps(inmeta,ensure_ascii=False)
+        js["__metadata__"]=inmeta
+        print()
+
+    newhdrbuf=json.dumps(js,separators=(',',':'),ensure_ascii=False).encode('utf-8')
+    newhdrlen:int=int(len(newhdrbuf))
+    pad:int=((newhdrlen+7)&(~7))-newhdrlen #pad to multiple of 8
+
+    with open(output_file,"wb") as f:
+        f.write(int(newhdrlen+pad).to_bytes(8,'little'))
+        f.write(newhdrbuf)
+        if pad>0: f.write(bytearray([32]*pad))
+        i:int=s.copy_data_to_file(f)
+    if i==0:
+        print(f"file {output_file} saved successfully")
+    else:
+        print(f"error {i} occurred when writing to file {output_file}")
+    return i
+
+def PrintHeader(cmdLine:dict,input_file:str) -> int:
+    s=SafeTensorsFile.open_file(input_file,cmdLine['quiet'])
+    js=s.get_header()
+
+    # All the .safetensors files I've seen have long key names, and as a result,
+    # neither json nor pprint package prints text in very readable format,
+    # so we print it ourselves, putting key name & value on one long line.
+    # Note the print out is in Python format, not valid JSON format.
+    firstKey=True
+    print("{")
+    for key in js:
+        if firstKey:
+            firstKey=False
+        else:
+            print(",")
+        json.dump(key,fp=sys.stdout,ensure_ascii=False,separators=(',',':'))
+        print(": ",end='')
+        json.dump(js[key],fp=sys.stdout,ensure_ascii=False,separators=(',',':'))
+    print("\n}")
+    return 0
+
+def _ParseMore(d:dict):
+    '''Basically try to turn this:
+
+        "ss_dataset_dirs":"{\"abc\": {\"n_repeats\": 2, \"img_count\": 60}}",
+
+    into this:
+
+        "ss_dataset_dirs":{
+         "abc":{
+          "n_repeats":2,
+          "img_count":60
+         }
+        },
+
+    '''
+    for key in d:
+        value=d[key]
+        #print("+++",key,value,type(value),"+++",sep='|')
+        if isinstance(value,str):
+            try:
+                v2=json.loads(value)
+                d[key]=v2
+                value=v2
+            except json.JSONDecodeError as e:
+                pass
+        if isinstance(value,dict):
+            _ParseMore(value)
+
+def PrintMetadata(cmdLine:dict,input_file:str) -> int:
+    with SafeTensorsFile.open_file(input_file,cmdLine['quiet']) as s:
+        js=s.get_header()
+
+        if not "__metadata__" in js:
+            print("file header does not contain a __metadata__ item",file=sys.stderr)
+            return -2
+
+        md=js["__metadata__"]
+        if cmdLine['parse_more']:
+            _ParseMore(md)
+        json.dump({"__metadata__":md},fp=sys.stdout,ensure_ascii=False,separators=(',',':'),indent=1)
+    return 0
+
+def HeaderKeysToLists(cmdLine:dict,input_file:str) -> int:
+    s=SafeTensorsFile.open_file(input_file,cmdLine['quiet'])
+    js=s.get_header()
+
+    _lora_keys:list[tuple(str,bool)]=[] # use list to sort by name
+    for key in js:
+        if key=='__metadata__': continue
+        v=js[key]
+        isScalar=False
+        if isinstance(v,dict):
+            if 'shape' in v:
+                if 0==len(v['shape']):
+                    isScalar=True
+        _lora_keys.append((key,isScalar))
+    _lora_keys.sort(key=lambda x:x[0])
+
+    def printkeylist(kl):
+        firstKey=True
+        for key in kl:
+            if firstKey: firstKey=False
+            else: print(",")
+            print(key,end='')
+        print()
+
+    print("# use list to keep insertion order")
+    print("_lora_keys:list[tuple[str,bool]]=[")
+    printkeylist(_lora_keys)
+    print("]")
+
+    return 0
+
+
+def ExtractHeader(cmdLine:dict,input_file:str,output_file:str)->int:
+    if _need_force_overwrite(output_file,cmdLine): return -1
+
+    s=SafeTensorsFile.open_file(input_file,parseHeader=False)
+    if s.error!=0: return s.error
+
+    hdrbuf=s.hdrbuf
+    s.close_file() #close it in case user wants to write back to input_file itself
+    with open(output_file,"wb") as fo:
+        wn=fo.write(hdrbuf)
+        if wn!=len(hdrbuf):
+            print(f"write output file failed, tried to write {len(hdrbuf)} bytes, only wrote {wn} bytes",file=sys.stderr)
+            return -1
+    print(f"raw header saved to file {output_file}")
+    return 0
+
+
+def _CheckLoRA_internal(s:SafeTensorsFile)->int:
+    import lora_keys_sd15 as lora_keys
+    js=s.get_header()
+    set_scalar=set()
+    set_nonscalar=set()
+    for x in lora_keys._lora_keys:
+        if x[1]==True: set_scalar.add(x[0])
+        else: set_nonscalar.add(x[0])
+
+    bad_unknowns:list[str]=[] # unrecognized keys
+    bad_scalars:list[str]=[] #bad scalar
+    bad_nonscalars:list[str]=[] #bad nonscalar
+    for key in js:
+        if key in set_nonscalar:
+            if js[key]['shape']==[]: bad_nonscalars.append(key)
+            set_nonscalar.remove(key)
+        elif key in set_scalar:
+            if js[key]['shape']!=[]: bad_scalars.append(key)
+            set_scalar.remove(key)
+        else:
+            if "__metadata__"!=key:
+                bad_unknowns.append(key)
+
+    hasError=False
+
+    if len(bad_unknowns)!=0:
+        print("INFO: unrecognized items:")
+        for x in bad_unknowns: print(" ",x)
+        #hasError=True
+
+    if len(set_scalar)>0:
+        print("missing scalar keys:")
+        for x in set_scalar: print(" ",x)
+        hasError=True
+    if len(set_nonscalar)>0:
+        print("missing nonscalar keys:")
+        for x in set_nonscalar: print(" ",x)
+        hasError=True
+
+    if len(bad_scalars)!=0:
+        print("keys expected to be scalar but are nonscalar:")
+        for x in bad_scalars: print(" ",x)
+        hasError=True
+
+    if len(bad_nonscalars)!=0:
+        print("keys expected to be nonscalar but are scalar:")
+        for x in bad_nonscalars: print(" ",x)
+        hasError=True
+
+    return (1 if hasError else 0)
+
+def CheckLoRA(cmdLine:dict,input_file:str)->int:
+    s=SafeTensorsFile.open_file(input_file)
+    i:int=_CheckLoRA_internal(s)
+    if i==0: print("looks like an OK SD 1.x LoRA file")
+    return 0
+
+def ExtractData(cmdLine:dict,input_file:str,key_name:str,output_file:str)->int:
+    if _need_force_overwrite(output_file,cmdLine): return -1
+
+    s=SafeTensorsFile.open_file(input_file,cmdLine['quiet'])
+    if s.error!=0: return s.error
+
+    bindata=s.load_one_tensor(key_name)
+    s.close_file() #close it just in case user wants to write back to input_file itself
+    if bindata is None:
+        print(f'key "{key_name}" not found in header (key names are case-sensitive)',file=sys.stderr)
+        return -1
+
+    with open(output_file,"wb") as fo:
+        wn=fo.write(bindata)
+        if wn!=len(bindata):
+            print(f"write output file failed, tried to write {len(bindata)} bytes, only wrote {wn} bytes",file=sys.stderr)
+            return -1
+    if cmdLine['quiet']==False: print(f"{key_name} saved to {output_file}, len={wn}")
+    return 0